use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use memmap2::Mmap;
use object::pe::{ImageNtHeaders32, ImageNtHeaders64};
use object::read::pe::{ImageNtHeaders, ImageOptionalHeader, PeFile};
use object::FileKind;
use wholesym::{CodeId, PeCodeId};
use super::avma_range::AvmaRange;
#[derive(Debug, Clone)]
pub struct SuspectedPeMapping {
pub path: PathBuf,
pub code_id: CodeId,
pub avma_range: AvmaRange,
}
pub struct PeMappings {
suspected_pe_mappings: BTreeMap<u64, SuspectedPeMapping>,
}
impl PeMappings {
pub fn new() -> Self {
Self {
suspected_pe_mappings: BTreeMap::new(),
}
}
pub fn find_mapping(&self, avma_range: &AvmaRange) -> Option<&SuspectedPeMapping> {
let (_, mapping) = self
.suspected_pe_mappings
.range(..=avma_range.start())
.next_back()?;
if !avma_range.encompasses(&mapping.avma_range) {
return None;
}
Some(mapping)
}
pub fn check_mmap(&mut self, path_slice: &[u8], mapping_start_avma: u64) {
let filename_is_pe = path_slice.ends_with(b".exe")
|| path_slice.ends_with(b".dll")
|| path_slice.ends_with(b".EXE")
|| path_slice.ends_with(b".DLL");
if !filename_is_pe {
return;
}
let Ok(path) = std::str::from_utf8(path_slice) else {
return;
};
let path = Path::new(path);
if let Some((size, code_id)) = get_pe_mapping_size_and_codeid(path) {
let mapping = SuspectedPeMapping {
path: path.to_owned(),
code_id,
avma_range: AvmaRange::with_start_size(mapping_start_avma, size),
};
self.suspected_pe_mappings
.insert(mapping_start_avma, mapping);
}
}
}
fn get_pe_mapping_size_and_codeid(path: &Path) -> Option<(u64, CodeId)> {
fn inner<T: ImageNtHeaders>(data: &[u8]) -> Option<(u64, CodeId)> {
let file = PeFile::<T>::parse(data).ok()?;
let image_size = file.nt_headers().optional_header().size_of_image();
let timestamp = file
.nt_headers()
.file_header()
.time_date_stamp
.get(object::LittleEndian);
let code_id = CodeId::PeCodeId(PeCodeId {
timestamp,
image_size,
});
Some((image_size as u64, code_id))
}
let file = std::fs::File::open(path).ok()?;
let mmap = unsafe { Mmap::map(&file).ok()? };
match FileKind::parse(&mmap[..]).ok()? {
FileKind::Pe32 => inner::<ImageNtHeaders32>(&mmap),
FileKind::Pe64 => inner::<ImageNtHeaders64>(&mmap),
_ => None,
}
}