fix(windows): restore managed companion DLLs

docs/design.md

@@ -32,7 +32,7 @@ senbei-wasm/src/
 
 `senbei-engine/src/windows/` contains PE detection, layout discovery, EXE and DLL restoration, deterministic block parallelism, and structural integrity checks. Candidate layouts are trial-decrypted and validated before an output is accepted.
 
-External companion inputs are reconstructed as `stub[..4096]` followed by the matching `._` payload. The stub's export and TLS data is overlaid after unpacking because those regions are not present in the encrypted companion.
+External companion inputs are reconstructed as `stub[..4096]` followed by the matching `._` payload. The stub's export, TLS, and declared CLR regions are overlaid after unpacking because those regions are not present in the encrypted companion. Managed restoration follows the COR20 directory and referenced metadata, resources, and vtable fixups through each file's RVA mapping, preserving the decrypted method bodies.
 
 ## Android Engine
 

docs/usage.md

@@ -25,6 +25,8 @@ If a restored library contains embedded metadata, the unwrapped blob is written
 
 Folder mode walks recursively, skips directories named `unpack`, and mirrors recognized outputs below `<root>/unpack/` or `--out DIR`. Windows candidates are `.exe`, `.dll`, and `global-metadata.dat`; Android candidates are `.so` and `global-metadata.dat`. A matching `.exe._` or `.dll._` payload is consumed by its stub and is excluded from the skipped count.
 
+Managed DLL companions retain CLR metadata and related runtime tables in the original DLL. Both the DLL and its matching `._` file must be available; Senbei restores the declared CLR regions from the DLL while retaining method bodies decrypted from the companion. Invalid or missing referenced regions are reported as errors.
+
 The summary has the form `12 unpacked · 3 skipped · 0 errors · 1 suspect · 2 metadata`; the package count is appended when packages were opened. Each file is isolated so one failed target does not stop the folder run.
 
 ## Integrity Check

senbei-engine/src/windows/error.rs

@@ -149,6 +149,13 @@ pub enum UnpackError {
         buffer_len: usize,
     },
 
+    #[error("managed stub {region} restoration failed: {source}")]
+    ManagedStubRestoreFailed {
+        region: &'static str,
+        #[source]
+        source: senbei_pe::Error,
+    },
+
     #[error(
         "EXE checksum descriptor at 0x{descriptor:08X} points outside input (offset {offset}, size {size}, input length {image_len})"
     )]

senbei-engine/src/windows/exe/pipeline/pe32.rs

@@ -1,6 +1,27 @@
 use super::super::super::layout;
 use super::*;
 
+fn stage_key_rounds(data: &[u8], table: u32, slots: usize) -> Result<u32, UnpackError> {
+    let offset = table as usize;
+    let size = slots.saturating_mul(16);
+    let descriptors = offset
+        .checked_add(size)
+        .and_then(|end| data.get(offset..end))
+        .ok_or(UnpackError::BufferRangeOutOfBounds {
+            operation: BufferOperation::Read,
+            offset,
+            size,
+            buffer_len: data.len(),
+        })?;
+    // The loader stops at the first empty helper, even if later slots are nonempty.
+    Ok(descriptors
+        .as_chunks::<16>()
+        .0
+        .iter()
+        .take_while(|descriptor| get_u32(descriptor.as_slice(), 4) > 4)
+        .count() as u32)
+}
+
 impl<'a> Unpacker<'a> {
     /// PE32 (32-bit) unpack pipeline. The shared Stage 1/2 setup (info decrypt,
     /// payload decrypt, raw copy, header restore) has already run in `run()`
@@ -224,11 +245,15 @@ impl<'a> Unpacker<'a> {
 
         // ---- ForthStage ----
         let second_stage_cs = self.calculate_checksum(second_stage_cs_addr);
+        let dp_base = ss.wrapping_add(dp_base_off);
+        let forth_key_rounds = stage_key_rounds(&self.decompressed, dp_base, 4)?;
         let forth_stage_key = advance_key(
             get_u32(&self.decompressed, ss.wrapping_add(forth_key_off)),
-            4,
+            forth_key_rounds,
         );
-        let dp_base = ss.wrapping_add(dp_base_off);
+        if verbose {
+            println!("  fourth-stage key rounds = {forth_key_rounds}");
+        }
         let forth_addr = dp_base.wrapping_add(0x40);
         let fk = header_checksum ^ second_stage_cs ^ forth_stage_key;
         if let Err(reason) = self.decrypt_and_decompress_data(forth_addr, fk, None) {
@@ -313,6 +338,11 @@ impl<'a> Unpacker<'a> {
         )?;
 
         let seven_cs = self.calculate_checksum(seven_stage_cs_addr);
+        let eighth_key_rounds =
+            stage_key_rounds(&self.decompressed, dp_base.wrapping_add(0x80), 4)?;
+        if verbose {
+            println!("  eighth-stage key rounds = {eighth_key_rounds}");
+        }
         let eighth_addr = dp_base.wrapping_add(0xC0);
         let eighth_dsz = get_u32(&self.decompressed, eighth_addr.wrapping_add(12));
         let eighth_src = get_u32(&self.decompressed, eighth_addr);
@@ -381,7 +411,7 @@ impl<'a> Unpacker<'a> {
             self.decompressed[eighth_addr as usize..(eighth_addr + 16) as usize]
                 .copy_from_slice(&eighth_pair_bak);
             let raw = get_u32(&self.decompressed, seven_start_actual.wrapping_add(ek_off));
-            let test_key = advance_key(raw, 3);
+            let test_key = advance_key(raw, eighth_key_rounds);
             let fk8 = header_checksum ^ fifth_cs ^ seven_cs ^ test_key;
             let result = primitives::decrypt_and_decompress_data(
                 &mut self.decompressed,
@@ -445,7 +475,7 @@ impl<'a> Unpacker<'a> {
                 let mut best: Option<(u32 /*dist*/, u32 /*off*/)> = None;
                 let mut o = 0u32;
                 let dlen = self.decompressed.len() as u32;
-                while o + 8 <= eighth_dsz.saturating_sub(0x4B4u32.saturating_sub(0x30)) {
+                while o + 8 <= eighth_dsz {
                     let fc = get_u32(&self.decompressed, eighth_start.wrapping_add(o));
                     let sz = get_u32(&self.decompressed, eighth_start.wrapping_add(o + 4));
                     if fc > info3
@@ -454,8 +484,9 @@ impl<'a> Unpacker<'a> {
                         && (0x10..=0x200).contains(&sz)
                         && (sz & 0xF) == 0
                     {
-                        // Cluster base must leave room for the +0x4B4 LFSR slot
+                        // Compact stages place the decryptor closer to this
-                        // (even if the exact LFSR is later adjusted by scan).
+                        // cluster. Only the config fields must fit here; the
+                        // actual LFSR location is trial-validated below.
                         if o >= 0x30 {
                             let base = o - 0x30;
                             if base.wrapping_add(0x4C) <= eighth_dsz {
@@ -1061,3 +1092,36 @@ impl<'a> Unpacker<'a> {
         Ok(compact)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn stage_key_rounds_follow_active_descriptor_prefix() {
+        for (sizes, expected) in [
+            ([0, 0x45, 0x45, 0x45], 0),
+            ([0x45, 4, 0x45, 0x45], 1),
+            ([0x45, 0x45, 0x45, 3], 3),
+            ([0x45, 0x45, 0x45, 0x45], 4),
+        ] {
+            let mut data = [0u8; 80];
+            for (index, size) in sizes.into_iter().enumerate() {
+                write_u32(&mut data, 16 + index as u32 * 16 + 4, size);
+            }
+            assert_eq!(stage_key_rounds(&data, 16, 4).unwrap(), expected);
+        }
+    }
+
+    #[test]
+    fn stage_key_rounds_reject_truncated_tables() {
+        assert!(matches!(
+            stage_key_rounds(&[0u8; 63], 0, 4),
+            Err(UnpackError::BufferRangeOutOfBounds { .. })
+        ));
+        assert!(matches!(
+            stage_key_rounds(&[0u8; 64], u32::MAX, 4),
+            Err(UnpackError::BufferRangeOutOfBounds { .. })
+        ));
+    }
+}

senbei-io/src/windows/mod.rs

@@ -142,6 +142,97 @@ pub(crate) fn overlay_exports_from_stub(out: &mut [u8], stub: &[u8]) {
     }
 }
 
+/// Restore the CLR regions retained by an external-companion loader stub.
+/// Method bodies come from the unpacked payload and must not be overlaid.
+fn restore_managed_from_stub(out: &mut [u8], stub: &[u8]) -> Result<(), unpacker::UnpackError> {
+    let failure =
+        |region, source| unpacker::UnpackError::ManagedStubRestoreFailed { region, source };
+    let source_headers = senbei_pe::parse(stub).map_err(|e| failure("PE headers", e))?;
+    let (clr_rva, clr_size) = senbei_pe::data_directory(stub, source_headers, 14)
+        .map_err(|e| failure("CLR directory", e))?;
+    if clr_rva == 0 && clr_size == 0 {
+        return Ok(());
+    }
+    if clr_rva == 0 || clr_size < 0x48 {
+        return Err(failure("CLR directory", senbei_pe::Error::Invalid));
+    }
+    let destination_headers = senbei_pe::parse(out).map_err(|e| failure("output PE headers", e))?;
+    senbei_pe::data_directory(out, destination_headers, 14)
+        .map_err(|e| failure("output CLR directory", e))?;
+    let cor = senbei_pe::rva_range(stub, source_headers, clr_rva, 0x48)
+        .map_err(|e| failure("COR20 header", e))?;
+    if read_u32(stub, cor.start) != Some(0x48) {
+        return Err(failure("COR20 header", senbei_pe::Error::Invalid));
+    }
+    let range_pair = |rva, size, region| {
+        let source = senbei_pe::rva_range(stub, source_headers, rva, size)
+            .map_err(|e| failure(region, e))?;
+        let destination = senbei_pe::rva_range(out, destination_headers, rva, size)
+            .map_err(|e| failure(region, e))?;
+        Ok::<_, unpacker::UnpackError>((source, destination))
+    };
+    let mut copies = vec![range_pair(clr_rva, 0x48, "COR20 header")?];
+    for (field, region) in [
+        (0x08, "metadata"),
+        (0x18, "resources"),
+        (0x20, "strong-name signature"),
+        (0x28, "code-manager table"),
+        (0x30, "vtable fixups"),
+        (0x38, "export address jumps"),
+        (0x40, "managed native header"),
+    ] {
+        let rva = read_u32(stub, cor.start + field)
+            .ok_or_else(|| failure(region, senbei_pe::Error::OutOfBounds))?;
+        let size = read_u32(stub, cor.start + field + 4)
+            .ok_or_else(|| failure(region, senbei_pe::Error::OutOfBounds))?;
+        if field != 0x08 && rva == 0 && size == 0 {
+            continue;
+        }
+        if rva == 0 || size == 0 {
+            return Err(failure(region, senbei_pe::Error::Invalid));
+        }
+        let (source, destination) = range_pair(rva, size, region)?;
+        if field == 0x08 && !stub[source.clone()].starts_with(b"BSJB") {
+            return Err(failure(region, senbei_pe::Error::Invalid));
+        }
+        if field == 0x30 {
+            if !size.is_multiple_of(8) {
+                return Err(failure(region, senbei_pe::Error::Invalid));
+            }
+            for fixup in stub[source.clone()].as_chunks::<8>().0 {
+                let slots_rva =
+                    u32::from_le_bytes(fixup[..4].try_into().expect("eight-byte fixup"));
+                let count = u16::from_le_bytes([fixup[4], fixup[5]]) as u32;
+                let flags = u16::from_le_bytes([fixup[6], fixup[7]]);
+                let width = match flags & 3 {
+                    1 => 4,
+                    2 => 8,
+                    _ => return Err(failure(region, senbei_pe::Error::Invalid)),
+                };
+                if count != 0 {
+                    copies.push(range_pair(slots_rva, count * width, "vtable slots")?);
+                }
+            }
+        }
+        copies.push((source, destination));
+    }
+    // Validate all referenced ranges before changing the output.
+    for (source, destination) in copies {
+        out[destination].copy_from_slice(&stub[source]);
+    }
+    let directory = destination_headers.pe_offset
+        + 24
+        + if destination_headers.is_pe32_plus {
+            112
+        } else {
+            96
+        }
+        + 14 * 8;
+    out[directory..directory + 4].copy_from_slice(&clr_rva.to_le_bytes());
+    out[directory + 4..directory + 8].copy_from_slice(&clr_size.to_le_bytes());
+    Ok(())
+}
+
 /// Restore the TLS directory from the loader `stub` onto the unpacked image
 /// `out`, for the external-companion layout.
 ///
@@ -413,6 +504,7 @@ fn unpack_bytes_impl(
     if spliced.is_some() {
         overlay_exports_from_stub(&mut out, input);
         restore_tls_from_stub(&mut out, input);
+        restore_managed_from_stub(&mut out, input)?;
     }
     let integrity = unpacker::check_integrity(&out);
     Ok(UnpackedImage {
@@ -440,6 +532,7 @@ pub fn unpack_one_v(
         // re-installs at runtime; the ordinary loader needs it or thread_local
         // access crashes (see [`restore_tls_from_stub`]).
         restore_tls_from_stub(&mut out, &stub);
+        restore_managed_from_stub(&mut out, &stub)?;
     }
     let report = unpacker::check_integrity(&out);
     if let Some(parent) = dest.parent() {
@@ -496,4 +589,89 @@ mod tests {
         let short_companion = vec![1_u8; 16];
         assert!(splice_companion(&stub, &short_companion).is_none());
     }
+
+    fn managed_fixture(is_pe32_plus: bool, raw: usize) -> Vec<u8> {
+        let mut data = vec![0; raw + 0x600];
+        data[..2].copy_from_slice(b"MZ");
+        data[0x80..0x84].copy_from_slice(b"PE\0\0");
+        let optional_size = if is_pe32_plus { 0xf0u16 } else { 0xe0 };
+        let section = 0x98 + optional_size as usize;
+        let dirs = 0x98 + if is_pe32_plus { 112 } else { 96 };
+        for (offset, value) in [
+            (0x86, 1u16),
+            (0x94, optional_size),
+            (0x98, if is_pe32_plus { 0x20b } else { 0x10b }),
+            (raw + 0x204, 2),
+            (raw + 0x206, if is_pe32_plus { 2 } else { 1 }),
+        ] {
+            data[offset..offset + 2].copy_from_slice(&value.to_le_bytes());
+        }
+        for (offset, value) in [
+            (0x3c, 0x80u32),
+            (0xd0, 0x3000),
+            (0xd4, 0x400),
+            (dirs + 14 * 8, 0x2010),
+            (dirs + 14 * 8 + 4, 0x48),
+            (section + 8, 0x600),
+            (section + 12, 0x2000),
+            (section + 16, 0x600),
+            (section + 20, raw as u32),
+            (raw + 0x10, 0x48),
+            (raw + 0x18, 0x2100),
+            (raw + 0x1c, 0x20),
+            (raw + 0x28, 0x2180),
+            (raw + 0x2c, 8),
+            (raw + 0x40, 0x2200),
+            (raw + 0x44, 8),
+            (raw + 0x200, 0x2280),
+            (raw + 0x280, 0x0600_0001),
+        ] {
+            data[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
+        }
+        data[raw + 0x100..raw + 0x104].copy_from_slice(b"BSJB");
+        data[raw + 0x180..raw + 0x188].copy_from_slice(b"resource");
+        data
+    }
+
+    #[test]
+    fn managed_companion_restores_rva_mapped_regions_without_overwriting_il() {
+        for is_pe32_plus in [false, true] {
+            let stub = managed_fixture(is_pe32_plus, 0x600);
+            let mut out = managed_fixture(is_pe32_plus, 0x400);
+            out[0x400..].fill(0xcc);
+            restore_managed_from_stub(&mut out, &stub).unwrap();
+            for (offset, size) in [
+                (0x10, 0x48),
+                (0x100, 0x20),
+                (0x180, 8),
+                (0x200, 8),
+                (0x280, if is_pe32_plus { 16 } else { 8 }),
+            ] {
+                assert_eq!(
+                    &out[0x400 + offset..0x400 + offset + size],
+                    &stub[0x600 + offset..0x600 + offset + size]
+                );
+            }
+            assert!(out[0x700..0x740].iter().all(|&b| b == 0xcc));
+        }
+    }
+
+    #[test]
+    fn managed_companion_rejects_invalid_metadata_and_unbacked_vtable_slots() {
+        for broken_metadata in [true, false] {
+            let mut stub = managed_fixture(false, 0x600);
+            if broken_metadata {
+                stub[0x700..0x704].fill(0);
+            } else {
+                stub[0x800..0x804].copy_from_slice(&0x2600u32.to_le_bytes());
+            }
+            let mut out = managed_fixture(false, 0x400);
+            let before = out.clone();
+            assert!(matches!(
+                restore_managed_from_stub(&mut out, &stub),
+                Err(unpacker::UnpackError::ManagedStubRestoreFailed { .. })
+            ));
+            assert_eq!(out, before);
+        }
+    }
 }

senbei-pe/src/lib.rs

@@ -139,6 +139,36 @@ pub fn rva_to_offset(data: &[u8], headers: Headers, rva: u32) -> Result<usize> {
     Err(Error::OutOfBounds)
 }
 
+/// Map a complete RVA range backed by file bytes in the headers or one section.
+/// Unlike a virtual mapping, this rejects a section's zero-filled tail.
+pub fn rva_range(
+    data: &[u8],
+    headers: Headers,
+    rva: u32,
+    size: u32,
+) -> Result<std::ops::Range<usize>> {
+    let header_size = read_u32(data, headers.pe_offset + 24 + 60)?;
+    let offset = if rva < header_size && size <= header_size - rva {
+        rva
+    } else {
+        sections(data, headers)?
+            .into_iter()
+            .find_map(|section| {
+                let delta = rva.checked_sub(section.virtual_address)?;
+                if delta >= section.raw_size || size > section.raw_size - delta {
+                    return None;
+                }
+                section.raw_offset.checked_add(delta)
+            })
+            .ok_or(Error::OutOfBounds)?
+    } as usize;
+    let end = offset
+        .checked_add(size as usize)
+        .ok_or(Error::OutOfBounds)?;
+    data.get(offset..end).ok_or(Error::OutOfBounds)?;
+    Ok(offset..end)
+}
+
 fn read_u16(data: &[u8], offset: usize) -> Result<u16> {
     let bytes: [u8; 2] = data
         .get(offset..offset + 2)
@@ -165,3 +195,39 @@ fn read_u64(data: &[u8], offset: usize) -> Result<u64> {
         .map_err(|_| Error::OutOfBounds)?;
     Ok(u64::from_le_bytes(bytes))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn rva_ranges_require_file_backing_for_every_byte() {
+        let mut data = [0u8; 0x400];
+        let headers = Headers {
+            pe_offset: 0x40,
+            is_pe32_plus: false,
+            image_base: 0,
+            size_of_image: 0x2000,
+            entry_rva: 0x1000,
+            sections_offset: 0x100,
+            sections: 1,
+        };
+        for (offset, value) in [
+            (0x94, 0x200u32),
+            (0x108, 0x100),
+            (0x10c, 0x1000),
+            (0x110, 0x80),
+            (0x114, 0x200),
+        ] {
+            data[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
+        }
+        assert_eq!(rva_range(&data, headers, 0x1000, 0x80), Ok(0x200..0x280));
+        assert_eq!(rva_range(&data, headers, 0x100, 0x100), Ok(0x100..0x200));
+        for (rva, size) in [(0x1070, 0x20), (0x1080, 1), (0x1f0, 0x20), (u32::MAX, 4)] {
+            assert_eq!(
+                rva_range(&data, headers, rva, size),
+                Err(Error::OutOfBounds)
+            );
+        }
+    }
+}