use alloc::vec::Vec;
pub(crate) const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;
const DOS_MAGIC_MZ: u16 = 0x5A4D;
const NT_SIGNATURE_PE: u32 = 0x0000_4550;
const DOS_E_LFANEW_OFFSET: usize = 0x3C;
const FILE_HEADER_SIZE: usize = 20;
const SECTION_HEADER_SIZE: usize = 40;
const OPTIONAL_HEADER_OFFSET: usize = 4 + FILE_HEADER_SIZE;
const OPTIONAL_HEADER_SIZE_OF_IMAGE_OFFSET: usize = 56;
#[derive(Debug, Clone, Copy)]
struct PeHeaders {
module_base: usize,
nt: usize,
section_table: usize,
num_sections: usize,
}
#[inline]
fn parse_pe_headers(module_base: usize) -> Option<PeHeaders> {
if module_base == 0 {
return None;
}
unsafe {
if *(module_base as *const u16) != DOS_MAGIC_MZ {
return None;
}
let nt_offset = *((module_base + DOS_E_LFANEW_OFFSET) as *const u32) as usize;
let nt = module_base + nt_offset;
if *(nt as *const u32) != NT_SIGNATURE_PE {
return None;
}
let file_hdr = nt + 4;
let num_sections = *((file_hdr + 2) as *const u16) as usize;
let opt_hdr_size = *((file_hdr + 16) as *const u16) as usize;
let section_table = file_hdr + FILE_HEADER_SIZE + opt_hdr_size;
Some(PeHeaders {
module_base,
nt,
section_table,
num_sections,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct SectionInfo {
pub(crate) name: [u8; 8],
pub(crate) virtual_address: usize,
pub(crate) virtual_size: usize,
pub(crate) characteristics: u32,
}
impl SectionInfo {
pub(crate) fn is_executable(&self) -> bool {
self.characteristics & IMAGE_SCN_MEM_EXECUTE != 0
}
}
#[inline]
unsafe fn read_section_at(sec: usize, module_base: usize) -> SectionInfo {
let mut name = [0u8; 8];
core::ptr::copy_nonoverlapping(sec as *const u8, name.as_mut_ptr(), 8);
let virtual_size = *((sec + 8) as *const u32) as usize;
let virtual_address = *((sec + 12) as *const u32) as usize;
let characteristics = *((sec + 36) as *const u32);
SectionInfo {
name,
virtual_address: module_base + virtual_address,
virtual_size,
characteristics,
}
}
#[must_use]
pub(crate) fn iter_sections(module_base: usize) -> Option<impl Iterator<Item = SectionInfo>> {
let hdr = parse_pe_headers(module_base)?;
Some((0..hdr.num_sections).map(move |i| {
let sec = hdr.section_table + i * SECTION_HEADER_SIZE;
unsafe { read_section_at(sec, hdr.module_base) }
}))
}
#[must_use]
pub(crate) fn find_section(module_base: usize, prefix: &[u8]) -> Option<SectionInfo> {
iter_sections(module_base)?.find(|s| s.name.starts_with(prefix))
}
pub(crate) fn exec_sections(module_base: usize) -> Option<Vec<(usize, usize)>> {
Some(
iter_sections(module_base)?
.filter(SectionInfo::is_executable)
.map(|s| (s.virtual_address, s.virtual_size))
.collect(),
)
}
pub(crate) fn text_section_bounds(module_base: usize) -> Option<(usize, usize)> {
let s = find_section(module_base, b".text")?;
Some((s.virtual_address, s.virtual_size))
}
#[must_use]
pub fn module_size(module_base: usize) -> Option<usize> {
let hdr = parse_pe_headers(module_base)?;
unsafe {
let size_of_image_addr =
hdr.nt + OPTIONAL_HEADER_OFFSET + OPTIONAL_HEADER_SIZE_OF_IMAGE_OFFSET;
Some(*(size_of_image_addr as *const u32) as usize)
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
use alloc::vec::Vec;
pub(super) fn synthetic_pe(sections: &[([u8; 8], u32, &[u8], u32)]) -> Vec<u8> {
let needed = sections
.iter()
.map(|(_, vaddr, bytes, _)| *vaddr as usize + bytes.len())
.max()
.unwrap_or(0)
.max(0x400);
let mut buf = vec![0u8; needed];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
let nt = nt_offset as usize;
buf[nt..nt + 4].copy_from_slice(b"PE\0\0");
let num_sections: u16 = sections.len() as u16;
buf[nt + 4 + 2..nt + 4 + 4].copy_from_slice(&num_sections.to_le_bytes());
let opt_size: u16 = 0xF0;
buf[nt + 4 + 16..nt + 4 + 18].copy_from_slice(&opt_size.to_le_bytes());
let size_of_image: u32 = needed as u32;
let soi_offset = nt + OPTIONAL_HEADER_OFFSET + 56;
buf[soi_offset..soi_offset + 4].copy_from_slice(&size_of_image.to_le_bytes());
let section_table = nt + 4 + 20 + opt_size as usize;
for (i, (name, vaddr, bytes, characteristics)) in sections.iter().enumerate() {
let sec = section_table + i * 40;
buf[sec..sec + 8].copy_from_slice(name);
let vsize: u32 = bytes.len() as u32;
buf[sec + 8..sec + 12].copy_from_slice(&vsize.to_le_bytes());
buf[sec + 12..sec + 16].copy_from_slice(&vaddr.to_le_bytes());
buf[sec + 36..sec + 40].copy_from_slice(&characteristics.to_le_bytes());
let v = *vaddr as usize;
buf[v..v + bytes.len()].copy_from_slice(bytes);
}
buf
}
#[test]
fn parse_pe_headers_returns_none_for_zero_base() {
assert!(parse_pe_headers(0).is_none());
}
#[test]
fn parse_pe_headers_returns_none_for_missing_mz() {
let buf = vec![0u8; 0x400];
assert!(parse_pe_headers(buf.as_ptr() as usize).is_none());
}
#[test]
fn parse_pe_headers_returns_none_for_missing_pe_sig() {
let mut buf = vec![0u8; 0x400];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
assert!(parse_pe_headers(buf.as_ptr() as usize).is_none());
}
#[test]
fn parse_pe_headers_walks_to_section_table() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let hdr = parse_pe_headers(buf.as_ptr() as usize).unwrap();
assert_eq!(hdr.num_sections, 1);
assert_eq!(hdr.module_base, buf.as_ptr() as usize);
assert_eq!(hdr.section_table, buf.as_ptr() as usize + 0x80 + 24 + 0xF0);
}
#[test]
fn iter_sections_yields_every_section_in_order() {
let body_a = [0x90u8, 0xC3];
let body_b = [0xAAu8, 0xBB];
let body_c = [0xCCu8];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
(*b".rdata\0\0", 0x310, &body_b, IMAGE_SCN_MEM_READ),
(
*b".data\0\0\0",
0x320,
&body_c,
IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
),
]);
let base = buf.as_ptr() as usize;
let secs: Vec<SectionInfo> = iter_sections(base).unwrap().collect();
assert_eq!(secs.len(), 3);
assert_eq!(&secs[0].name, b".text\0\0\0");
assert_eq!(secs[0].virtual_address, base + 0x300);
assert_eq!(secs[0].virtual_size, body_a.len());
assert!(secs[0].is_executable());
assert_eq!(&secs[1].name, b".rdata\0\0");
assert!(!secs[1].is_executable());
assert_eq!(&secs[2].name, b".data\0\0\0");
assert!(!secs[2].is_executable());
}
#[test]
fn iter_sections_returns_none_for_zero_base() {
assert!(iter_sections(0).is_none());
}
#[test]
fn iter_sections_returns_none_for_malformed_module() {
let buf = vec![0u8; 0x400];
assert!(iter_sections(buf.as_ptr() as usize).is_none());
}
#[test]
fn find_section_prefix_matches_companion_sections() {
let body = [0x90u8];
let buf = synthetic_pe(&[
(*b".text$mn", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
(*b".data\0\0\0", 0x310, &body, 0),
]);
let base = buf.as_ptr() as usize;
let s = find_section(base, b".text").unwrap();
assert_eq!(&s.name, b".text$mn");
assert_eq!(s.virtual_address, base + 0x300);
}
#[test]
fn find_section_returns_none_when_absent() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
let base = buf.as_ptr() as usize;
assert!(find_section(base, b".text").is_none());
}
#[test]
fn find_section_exact_8_byte_name_disambiguates() {
let body = [0x90u8];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let s = find_section(base, b".text$mn").unwrap();
assert_eq!(s.virtual_address, base + 0x310);
}
#[test]
fn text_section_bounds_finds_text() {
let body = [0x90u8, 0xC3];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let (start, size) = text_section_bounds(base).unwrap();
assert_eq!(start, base + 0x300);
assert_eq!(size, body.len());
}
#[test]
fn text_section_bounds_rejects_missing_mz() {
let buf = vec![0u8; 0x400];
assert!(text_section_bounds(buf.as_ptr() as usize).is_none());
}
#[test]
fn text_section_bounds_rejects_missing_pe_sig() {
let mut buf = vec![0u8; 0x400];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
assert!(text_section_bounds(buf.as_ptr() as usize).is_none());
}
#[test]
fn text_section_bounds_skips_non_text_sections() {
let data_body = [0xAAu8, 0xBB];
let text_body = [0x90u8, 0xC3];
let buf = synthetic_pe(&[
(*b".data\0\0\0", 0x300, &data_body, 0),
(*b".text\0\0\0", 0x310, &text_body, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let (start, size) = text_section_bounds(base).unwrap();
assert_eq!(start, base + 0x310);
assert_eq!(size, text_body.len());
}
#[test]
fn text_section_bounds_returns_none_when_no_text() {
let body = [0xAAu8, 0xBB];
let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
assert!(text_section_bounds(buf.as_ptr() as usize).is_none());
}
#[test]
fn exec_sections_includes_only_executable() {
let exec_body = [0x90u8, 0xC3];
let data_body = [0xAAu8, 0xBB];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &exec_body, IMAGE_SCN_MEM_EXECUTE),
(*b".data\0\0\0", 0x310, &data_body, 0),
]);
let base = buf.as_ptr() as usize;
let secs = exec_sections(base).unwrap();
assert_eq!(secs.len(), 1);
assert_eq!(secs[0], (base + 0x300, exec_body.len()));
}
#[test]
fn exec_sections_returns_multiple_when_multiple_exec() {
let body_a = [0x90u8];
let body_b = [0xC3u8];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let secs = exec_sections(base).unwrap();
assert_eq!(secs.len(), 2);
assert_eq!(secs[0], (base + 0x300, body_a.len()));
assert_eq!(secs[1], (base + 0x310, body_b.len()));
}
#[test]
fn exec_sections_rejects_missing_mz() {
let buf = vec![0u8; 0x400];
assert!(exec_sections(buf.as_ptr() as usize).is_none());
}
#[test]
fn exec_sections_rejects_missing_pe_sig() {
let mut buf = vec![0u8; 0x400];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
assert!(exec_sections(buf.as_ptr() as usize).is_none());
}
#[test]
fn exec_sections_empty_when_no_exec_sections() {
let body = [0xAAu8];
let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
let secs = exec_sections(buf.as_ptr() as usize).unwrap();
assert!(secs.is_empty());
}
#[cfg(feature = "section-info")]
mod module_size_tests {
use super::synthetic_pe;
use crate::pe::{module_size, IMAGE_SCN_MEM_EXECUTE};
use alloc::vec;
#[test]
fn returns_size_of_image() {
let body = [0x90u8, 0xC3];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
assert_eq!(module_size(base), Some(buf.len()));
}
#[test]
fn rejects_zero_base() {
assert_eq!(module_size(0), None);
}
#[test]
fn rejects_missing_mz() {
let buf = vec![0u8; 0x400];
assert!(module_size(buf.as_ptr() as usize).is_none());
}
#[test]
fn rejects_missing_pe_sig() {
let mut buf = vec![0u8; 0x400];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
assert!(module_size(buf.as_ptr() as usize).is_none());
}
#[test]
fn reads_actual_size_of_image_field() {
let body = [0x90u8];
let mut buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let soi_offset = 0x80 + 4 + 20 + 56;
let sentinel: u32 = 0xDEAD_BEEF;
buf[soi_offset..soi_offset + 4].copy_from_slice(&sentinel.to_le_bytes());
let base = buf.as_ptr() as usize;
assert_eq!(module_size(base), Some(0xDEAD_BEEF));
}
#[test]
fn callable_via_public_re_export() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
assert_eq!(crate::module_size(base), Some(buf.len()));
}
}
}