use std::path::Path;
use super::{Categorization, FileCategorizer};
use limnifs_core::codec::{CODEC_BCJ_ARM64_LZ4, CODEC_BCJ_X86_LZ4};
const MIN_EXEC_SIZE: usize = 1024;
const ELF_MAGIC: [u8; 4] = [0x7F, b'E', b'L', b'F'];
const DOS_MAGIC: [u8; 2] = [b'M', b'Z'];
const MACHO_MAGICS: &[[u8; 4]] = &[
[0xFE, 0xED, 0xFA, 0xCE], [0xFE, 0xED, 0xFA, 0xCF], [0xCE, 0xFA, 0xED, 0xFE], [0xCF, 0xFA, 0xED, 0xFE], ];
const EM_X86_64: u16 = 0x3E;
const EM_AARCH64: u16 = 0xB7;
const PE_MACHINE_AMD64: u16 = 0x8664;
const PE_MACHINE_ARM64: u16 = 0xAA64;
const CPU_TYPE_X86_64: u32 = 0x0100_0007;
const CPU_TYPE_ARM64: u32 = 0x0100_000C;
fn pick_codec(data: &[u8]) -> Option<u8> {
if data.len() < MIN_EXEC_SIZE {
return None;
}
if data.starts_with(&ELF_MAGIC) {
return parse_elf(data);
}
if data.starts_with(&DOS_MAGIC) {
return parse_pe(data);
}
if MACHO_MAGICS.iter().any(|m| data.starts_with(m)) {
return parse_macho(data);
}
None
}
fn parse_elf(data: &[u8]) -> Option<u8> {
if data.len() < 20 {
return None;
}
let machine = u16::from_le_bytes([data[18], data[19]]);
route_by_arch(machine == EM_X86_64, machine == EM_AARCH64)
}
fn parse_pe(data: &[u8]) -> Option<u8> {
if data.len() < 0x40 {
return None;
}
let lfanew = u32::from_le_bytes([data[0x3C], data[0x3D], data[0x3E], data[0x3F]]) as usize;
let pe_off = lfanew.checked_add(4)?;
if data.len() < pe_off + 2 {
return None;
}
if &data[lfanew..lfanew + 4] != b"PE\0\0" {
return None;
}
let machine = u16::from_le_bytes([data[pe_off], data[pe_off + 1]]);
route_by_arch(machine == PE_MACHINE_AMD64, machine == PE_MACHINE_ARM64)
}
fn parse_macho(data: &[u8]) -> Option<u8> {
if data.len() < 12 {
return None;
}
let magic = [data[0], data[1], data[2], data[3]];
let is_le = matches!(magic, [0xFE, 0xED, 0xFA, 0xCE] | [0xFE, 0xED, 0xFA, 0xCF]);
let cputype = if is_le {
u32::from_le_bytes([data[4], data[5], data[6], data[7]])
} else {
u32::from_be_bytes([data[4], data[5], data[6], data[7]])
};
route_by_arch(cputype == CPU_TYPE_X86_64, cputype == CPU_TYPE_ARM64)
}
fn route_by_arch(x86_64: bool, arm64: bool) -> Option<u8> {
if x86_64 {
Some(CODEC_BCJ_X86_LZ4)
} else if arm64 {
Some(CODEC_BCJ_ARM64_LZ4)
} else {
None
}
}
pub struct ExecutableCategorizer;
impl FileCategorizer for ExecutableCategorizer {
fn name(&self) -> &'static str {
"executable"
}
fn categories(&self) -> &'static [&'static str] {
&["binary/executable"]
}
fn first_byte_hint(&self) -> Option<&'static [u8]> {
Some(&[0x7F, b'M', 0xFE, 0xCE, 0xCF])
}
fn categorize(&self, _path: &Path, data: &[u8]) -> Option<Categorization> {
let codec_id = pick_codec(data)?;
Some(Categorization {
codec_id,
codec_params: Vec::new(),
category: "binary/executable",
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn elf_x86_64() -> Vec<u8> {
let mut v = vec![0u8; MIN_EXEC_SIZE];
v[0..4].copy_from_slice(&ELF_MAGIC);
v[4] = 2; v[5] = 1; v[16..18].copy_from_slice(&2u16.to_le_bytes()); v[18..20].copy_from_slice(&EM_X86_64.to_le_bytes());
v
}
fn elf_aarch64() -> Vec<u8> {
let mut v = elf_x86_64();
v[18..20].copy_from_slice(&EM_AARCH64.to_le_bytes());
v
}
fn elf_unknown_arch() -> Vec<u8> {
let mut v = elf_x86_64();
v[18..20].copy_from_slice(&0x1234u16.to_le_bytes());
v
}
#[test]
fn detects_elf_x86_64() {
let data = elf_x86_64();
assert_eq!(
pick_codec(&data),
Some(CODEC_BCJ_X86_LZ4),
"ELF x86_64 should route to BCJ-x86+LZ4"
);
}
#[test]
fn detects_elf_aarch64() {
let data = elf_aarch64();
assert_eq!(
pick_codec(&data),
Some(CODEC_BCJ_ARM64_LZ4),
"ELF aarch64 should route to BCJ-ARM64+LZ4"
);
}
#[test]
fn unknown_arch_returns_none() {
let data = elf_unknown_arch();
assert_eq!(pick_codec(&data), None);
}
#[test]
fn small_input_returns_none() {
let mut v = elf_x86_64();
v.truncate(100); assert_eq!(pick_codec(&v), None);
}
#[test]
fn pe_x86_64_routes_correctly() {
let mut v = vec![0u8; MIN_EXEC_SIZE];
v[0..2].copy_from_slice(&DOS_MAGIC);
let pe_offset: u32 = 0x40;
v[0x3C..0x40].copy_from_slice(&pe_offset.to_le_bytes());
v[0x40..0x44].copy_from_slice(b"PE\0\0");
v[0x44..0x46].copy_from_slice(&PE_MACHINE_AMD64.to_le_bytes());
assert_eq!(pick_codec(&v), Some(CODEC_BCJ_X86_LZ4));
}
#[test]
fn macho_x86_64_routes_correctly() {
let mut v = vec![0u8; MIN_EXEC_SIZE];
v[0..4].copy_from_slice(&[0xFE, 0xED, 0xFA, 0xCF]); v[4..8].copy_from_slice(&CPU_TYPE_X86_64.to_le_bytes());
assert_eq!(pick_codec(&v), Some(CODEC_BCJ_X86_LZ4));
}
#[test]
fn non_executable_returns_none() {
assert_eq!(pick_codec(b"hello world text not exec"), None);
assert_eq!(pick_codec(&vec![0u8; 4096]), None);
}
}