use std::collections::BTreeSet;
use crate::{
analysis::cfg::{flow::Flow, CFG},
VA,
};
pub fn is_thunk(cfg: &CFG, va: VA) -> bool {
if let Some(succs) = cfg.flows.flows_by_src.get(&va) {
if succs.len() != 1 {
false
} else {
matches!(succs[0], Flow::UnconditionalJump(_))
}
} else {
false
}
}
pub fn find_thunks<'a, T>(cfg: &CFG, functions: T) -> BTreeSet<VA>
where
T: Iterator<Item = &'a VA>,
{
let mut thunks: BTreeSet<VA> = Default::default();
for &function in functions {
if is_thunk(cfg, function) {
thunks.insert(function);
}
}
thunks
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{analysis::cfg::InstructionIndex, rsrc::*};
use anyhow::Result;
#[test]
fn nop() -> Result<()> {
let buf = get_buf(Rsrc::NOP);
let pe = crate::loader::pe::PE::from_bytes(&buf)?;
let mut insns: InstructionIndex = Default::default();
for &ep in crate::analysis::pe::entrypoints::find_pe_entrypoint(&pe)?.iter() {
insns.build_index(&pe.module, ep)?;
}
for &exp in crate::analysis::pe::exports::find_pe_exports(&pe)?.iter() {
insns.build_index(&pe.module, exp)?;
}
insns.build_index(&pe.module, 0x4027F4)?;
let cfg = CFG::from_instructions(&pe.module, insns)?;
let thunks = find_thunks(&cfg, [0x405F42, 0x401000].iter());
assert!(thunks.contains(&0x405F42));
Ok(())
}
}