refactor: consolidate platform engines into senbei-engine

senbei-engine/src/windows/mod.rs

@@ -0,0 +1,391 @@
+//! PE detection, unpacking, and structural validation.
+
+pub mod dll;
+mod error;
+pub mod exe;
+pub mod integrity;
+mod layout;
+pub(crate) mod parallel;
+
+use senbei_crypto::primitives;
+use std::cell::RefCell;
+use std::sync::{Arc, Mutex};
+
+pub use dll::{unpack_dll, unpack_dll_v};
+pub use error::*;
+pub use exe::{unpack as unpack_exe, unpack_v as unpack_exe_v};
+pub use integrity::{IntegrityReport, check as check_integrity};
+pub use parallel::thread_cap;
+
+/// Maximum plausible PE `SizeOfImage` we are willing to allocate a zero buffer
+/// for. Guards against a corrupt/crafted header requesting a multi-gigabyte
+/// (or, as a sign-extended negative `i32`, multi-exabyte) allocation, which
+/// would abort the process — an abort that `catch_unpack` below cannot trap.
+/// Real protected binaries are far below this.
+pub(crate) const MAX_IMAGE_SIZE: u64 = senbei_crypto::MAX_IMAGE_SIZE;
+
+#[derive(Clone)]
+pub(crate) struct PanicCapture(Arc<Mutex<Option<PanicDetails>>>);
+
+#[derive(Clone)]
+struct PanicDetails {
+    message: String,
+    file: String,
+    line: u32,
+    column: u32,
+}
+
+thread_local! {
+    static ACTIVE_PANIC_CAPTURE: RefCell<Option<PanicCapture>> = const { RefCell::new(None) };
+}
+
+struct PanicCaptureGuard(Option<PanicCapture>);
+
+impl Drop for PanicCaptureGuard {
+    fn drop(&mut self) {
+        ACTIVE_PANIC_CAPTURE.with(|slot| {
+            slot.replace(self.0.take());
+        });
+    }
+}
+
+impl PanicCapture {
+    fn new() -> Self {
+        Self(Arc::new(Mutex::new(None)))
+    }
+
+    fn record(&self, info: &std::panic::PanicHookInfo<'_>) {
+        let location = info.location();
+        let details = PanicDetails {
+            message: panic_message(info.payload()),
+            file: location
+                .map(|value| value.file().to_owned())
+                .unwrap_or_else(|| "<unknown>".to_owned()),
+            line: location.map_or(0, std::panic::Location::line),
+            column: location.map_or(0, std::panic::Location::column),
+        };
+        let mut captured = self
+            .0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if captured.is_none() {
+            *captured = Some(details);
+        }
+    }
+
+    fn into_error(self, payload: &(dyn std::any::Any + Send)) -> UnpackError {
+        let details = self
+            .0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+            .unwrap_or_else(|| PanicDetails {
+                message: panic_message(payload),
+                file: "<unknown>".to_owned(),
+                line: 0,
+                column: 0,
+            });
+        UnpackError::InternalPanic {
+            message: details.message,
+            file: details.file,
+            line: details.line,
+            column: details.column,
+        }
+    }
+
+    fn merge_from(&self, other: &Self) {
+        let details = other
+            .0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone();
+        let Some(details) = details else { return };
+        let mut captured = self
+            .0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if captured.is_none() {
+            *captured = Some(details);
+        }
+    }
+}
+
+fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
+    if let Some(message) = payload.downcast_ref::<&str>() {
+        (*message).to_owned()
+    } else if let Some(message) = payload.downcast_ref::<String>() {
+        message.clone()
+    } else {
+        "non-string panic payload".to_owned()
+    }
+}
+
+fn install_panic_capture_hook() {
+    static INSTALL: std::sync::Once = std::sync::Once::new();
+    INSTALL.call_once(|| {
+        let previous = std::panic::take_hook();
+        std::panic::set_hook(Box::new(move |info| {
+            let capture = ACTIVE_PANIC_CAPTURE
+                .try_with(|slot| slot.borrow().clone())
+                .ok()
+                .flatten();
+            if let Some(capture) = capture {
+                capture.record(info);
+            } else {
+                previous(info);
+            }
+        }));
+    });
+}
+
+pub(crate) fn current_panic_capture() -> Option<PanicCapture> {
+    ACTIVE_PANIC_CAPTURE.with(|slot| slot.borrow().clone())
+}
+
+pub(crate) fn with_panic_capture<R>(capture: Option<PanicCapture>, f: impl FnOnce() -> R) -> R {
+    let previous = ACTIVE_PANIC_CAPTURE.with(|slot| slot.replace(capture));
+    let _guard = PanicCaptureGuard(previous);
+    f()
+}
+
+/// Run an unpack pipeline, converting any internal panic into a clean
+/// [`UnpackError::InternalPanic`] so the public API stays panic-free on any input
+/// (truncated/garbled files chase offsets out of bounds). The panic location and
+/// payload are captured for diagnostics without printing a backtrace to stderr.
+///
+/// Note: allocation *failures* abort the process and are NOT caught here; size
+/// requests are bounds-checked against [`MAX_IMAGE_SIZE`] before allocating.
+pub(crate) fn catch_unpack<F>(f: F) -> Result<Vec<u8>, UnpackError>
+where
+    F: FnOnce() -> Result<Vec<u8>, UnpackError>,
+{
+    install_panic_capture_hook();
+    let capture = PanicCapture::new();
+    let r = with_panic_capture(Some(capture.clone()), || {
+        std::panic::catch_unwind(std::panic::AssertUnwindSafe(f))
+    });
+    match r {
+        Ok(result) => result,
+        Err(payload) => Err(capture.into_error(payload.as_ref())),
+    }
+}
+
+/// Crackproof header magic stored in `keys[1]`/`info[1]`.
+pub(crate) const MAGIC_KONN: u32 = 0x4E4E4F4B; // b"KONN" little-endian (= 1313754955)
+
+/// True if `magic` is the Crackproof magic this unpacker supports.
+pub(crate) fn is_supported_magic(magic: u32) -> bool {
+    magic == MAGIC_KONN
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Kind {
+    NativeExe,
+    ManagedExe,
+    NativeDll,
+    ManagedDll,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Detected {
+    pub kind: Kind,
+    pub magic: u32,
+}
+
+// ---------------------------------------------------------------------------
+// Content-based detection
+// ---------------------------------------------------------------------------
+
+/// Derive the 8-element Crackproof key table from the header at offset 4096.
+/// Returns `None` if the input is too short or doesn't have a valid PE signature.
+fn key_table(input: &[u8]) -> Option<[u32; 8]> {
+    // Need at least 4128 bytes: the key-table loop below reads dwords up to
+    // offset 4124 (bytes 4124..4127). Guarding only `< 4096` would let a
+    // 4096..4127-byte PE (e.g. a 4 KiB stub) panic in `get_u32`.
+    if input.len() < 4128 {
+        return None;
+    }
+    // Validate the PE signature with checked arithmetic so a crafted offset
+    // cannot wrap the bounds check on a narrower target.
+    let e_lfanew = primitives::get_u32(input, 0x3C);
+    let pe_start = e_lfanew as usize;
+    if pe_start.checked_add(4).is_none_or(|end| end > input.len()) {
+        return None;
+    }
+    if &input[pe_start..pe_start + 4] != b"PE\0\0" {
+        return None;
+    }
+    // Derive 8 keys per the Crackproof header-key formula.
+    let mut keys = [0u32; 8];
+    keys[0] = primitives::get_u32(input, 4096);
+    let mut k = keys[0];
+    for i in 0u32..7 {
+        let cell = primitives::get_u32(input, 4100u32.wrapping_add(i.wrapping_mul(4)));
+        keys[(i + 1) as usize] = k ^ cell;
+        k = i.wrapping_mul(i) ^ (k.wrapping_add(cell).wrapping_sub(i));
+    }
+    Some(keys)
+}
+
+/// Detect whether `input` is a Crackproof-protected binary and classify it.
+/// Returns `None` if the magic doesn't match.
+///
+/// Routing: `keys[1]` must be the Crackproof magic (`KONN`).
+/// The PE IMAGE_FILE_DLL characteristic distinguishes EXE vs DLL;
+/// the CLR data-directory RVA distinguishes managed from native for both.
+pub fn detect(input: &[u8]) -> Option<Detected> {
+    let keys = key_table(input)?;
+    let magic = keys[1];
+    // Anything whose magic doesn't match is left untouched rather than
+    // detected-then-errored, honoring the "anything that doesn't match is
+    // left untouched" contract.
+    if !is_supported_magic(magic) {
+        return None;
+    }
+    // Use the PE DLL characteristic to distinguish EXE from DLL.
+    // IMAGE_FILE_HEADER.Characteristics is at peOff+4+18; bit 0x2000 = IMAGE_FILE_DLL.
+    let pe_off = primitives::get_u32(input, 0x3C);
+    let chars_offset = pe_off.wrapping_add(4).wrapping_add(18);
+    if (chars_offset as usize)
+        .checked_add(2)
+        .is_none_or(|end| end > input.len())
+    {
+        return None;
+    }
+    let chars =
+        (input[chars_offset as usize] as u16) | ((input[chars_offset as usize + 1] as u16) << 8);
+    let is_dll = (chars & 0x2000) != 0;
+    // Managed vs native via the CLR data-directory RVA.
+    // peOff + 24 = start of optional header. The data directories start at a
+    // magic-dependent offset within it: PE32 (0x10B) at +96, PE32+ (0x20B) at
+    // +112. Using the PE32+ offset on a PE32 image reads the wrong dword and
+    // can mis-flag a native image as managed.
+    //
+    // `get_u16`/`get_u32` index unchecked, so every read past the already-
+    // checked Characteristics word must be bounds-checked first: a truncated
+    // file (e.g. `e_lfanew` pointing at len-24) would otherwise panic here,
+    // and this detector runs on the folder scan threads where a panic aborts
+    // the whole run.
+    let opt_magic_off = pe_off.wrapping_add(24) as usize;
+    let b = input.get(opt_magic_off..opt_magic_off.checked_add(2)?)?;
+    let opt_magic = u16::from_le_bytes([b[0], b[1]]);
+    let dd_off: u32 = if opt_magic == 0x20B { 112 } else { 96 };
+    // + 14*8 = IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR
+    let clr_rva_offset = pe_off
+        .wrapping_add(24)
+        .wrapping_add(dd_off)
+        .wrapping_add(14u32.wrapping_mul(8));
+    if (clr_rva_offset as usize)
+        .checked_add(4)
+        .is_none_or(|end| end > input.len())
+    {
+        return None;
+    }
+    let clr_rva = primitives::get_u32(input, clr_rva_offset);
+    let kind = match (is_dll, clr_rva != 0) {
+        (false, false) => Kind::NativeExe,
+        (false, true) => Kind::ManagedExe,
+        (true, false) => Kind::NativeDll,
+        (true, true) => Kind::ManagedDll,
+    };
+    Some(Detected { kind, magic })
+}
+
+/// Detect the file type and dispatch to the matching pipeline.
+/// Returns the detected `Kind` together with the unpacked image bytes.
+pub fn unpack_auto(input: &[u8]) -> Result<(Kind, Vec<u8>), UnpackError> {
+    unpack_auto_v(input, false)
+}
+
+/// Like [`unpack_auto`], but prints detailed `[N/9]` unpack-step progress to
+/// stdout when `verbose` is true. Output bytes are identical regardless.
+pub fn unpack_auto_v(input: &[u8], verbose: bool) -> Result<(Kind, Vec<u8>), UnpackError> {
+    let detected = detect(input).ok_or(UnpackError::NotCrackproof)?;
+    let out = match detected.kind {
+        Kind::NativeExe | Kind::ManagedExe => unpack_exe_v(input, verbose)?,
+        Kind::NativeDll | Kind::ManagedDll => {
+            // Two Crackproof DLL layouts exist. The older one (the byte-identical
+            // DLL goldens) follows the pipeline in `dll.rs`. Newer builds protect
+            // DLLs with the EXE-style shell layout instead — `dll::unpack_dll`
+            // cannot parse them and errors. Try the DLL pipeline first; on
+            // failure, fall back to the EXE pipeline, which handles the new
+            // layout (including managed-DLL CLR metadata restore). The DLL-first
+            // order keeps the old-layout goldens byte-identical (the EXE
+            // pipeline "succeeds" on them but with different bytes).
+            match dll::unpack_dll_v(input, verbose) {
+                Ok(out) => out,
+                Err(dll_err) => match exe::unpack_v(input, verbose) {
+                    Ok(out) => out,
+                    Err(exe_err) => {
+                        return Err(UnpackError::PipelineFallbackFailed {
+                            dll: Box::new(dll_err),
+                            exe: Box::new(exe_err),
+                        });
+                    }
+                },
+            }
+        }
+    };
+    Ok((detected.kind, out))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn caught_panic_reports_location_and_message() {
+        let error = catch_unpack(|| -> Result<Vec<u8>, UnpackError> {
+            panic!("test panic");
+        })
+        .expect_err("panic must become an error");
+        let UnpackError::InternalPanic {
+            message,
+            file,
+            line,
+            column,
+        } = error
+        else {
+            panic!("unexpected error: {error}");
+        };
+        assert_eq!(message, "test panic");
+        assert!(
+            file.ends_with("senbei-engine/src/windows/mod.rs")
+                || file.ends_with("senbei-engine\\src\\windows\\mod.rs")
+        );
+        assert!(line > 0);
+        assert!(column > 0);
+    }
+
+    #[test]
+    fn worker_panic_keeps_the_worker_source_location() {
+        let error = catch_unpack(|| -> Result<Vec<u8>, UnpackError> {
+            let capture = current_panic_capture();
+            let result = std::thread::spawn(move || {
+                with_panic_capture(capture, || panic!("worker panic"));
+            })
+            .join();
+            if let Err(payload) = result {
+                std::panic::resume_unwind(payload);
+            }
+            Ok(Vec::new())
+        })
+        .expect_err("worker panic must become an error");
+        let UnpackError::InternalPanic {
+            message,
+            file,
+            line,
+            column,
+        } = error
+        else {
+            panic!("unexpected error: {error}");
+        };
+        assert_eq!(message, "worker panic");
+        assert!(
+            file.ends_with("senbei-engine/src/windows/mod.rs")
+                || file.ends_with("senbei-engine\\src\\windows\\mod.rs")
+        );
+        assert!(line > 0);
+        assert!(column > 0);
+    }
+}