use crate::pe::Pe;
use iced_x86::{Decoder, DecoderOptions, FlowControl, Instruction, Mnemonic, OpKind, Register};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::sync::{Arc, Mutex, OnceLock};
type AnalyzeCacheKey = (u64, u64, bool);
type AnalyzeCache = Mutex<HashMap<AnalyzeCacheKey, Arc<Result<CallGraph, String>>>>;
const MAX_FN_SPAN: u64 = 0x10000;
const MAX_FIELD_OFFSET: u64 = 0x8000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ArgKind {
Imm(u64),
Field(u64),
Global(u64),
LocalAddr,
Other,
}
#[derive(Debug, Clone, Default)]
pub struct CallSite {
pub target: u64,
pub ecx: Option<ArgKind>,
pub edx: Option<ArgKind>,
pub stack: Vec<ArgKind>,
}
#[derive(Debug, Default, Clone)]
pub struct CallGraph {
pub entry: u64,
pub direct_calls: BTreeMap<u64, u32>,
pub call_sites: Vec<CallSite>,
pub tail_calls: BTreeSet<u64>,
pub indirect_calls: u32,
pub indirect_jumps: u32,
pub vcall_slots: BTreeSet<u64>,
pub iat_calls: BTreeSet<u64>,
pub field_offsets: BTreeSet<u64>,
pub globals: BTreeSet<u64>,
pub imm_addrs: BTreeSet<u64>,
pub signed_branches: u32,
pub unsigned_branches: u32,
pub equality_branches: u32,
pub max_end: u64,
pub decode_errors: u32,
pub decode_error_samples: Vec<(u64, Vec<u8>)>,
}
impl CallGraph {
pub fn call_targets(&self) -> BTreeSet<u64> {
self.direct_calls.keys().copied().collect()
}
pub fn footprint(&self) -> BTreeSet<u64> {
let mut s: BTreeSet<u64> = self.direct_calls.keys().copied().collect();
s.extend(&self.tail_calls);
s.extend(&self.iat_calls);
s.extend(&self.globals);
s.extend(&self.imm_addrs);
s
}
}
fn scan_operands(pe: &Pe, instr: &Instruction, cg: &mut CallGraph) {
for i in 0..instr.op_count() {
match instr.op_kind(i) {
OpKind::Memory => {
let base = instr.memory_base();
let index = instr.memory_index();
let disp = instr.memory_displacement64();
if base == Register::None && index == Register::None {
if pe.section_of(disp).is_some() && instr.mnemonic() != Mnemonic::Lea {
cg.globals.insert(disp);
}
} else if pe.section_of(disp).is_some() && instr.mnemonic() != Mnemonic::Lea {
cg.globals.insert(disp);
} else if base != Register::None
&& base != Register::ESP
&& base != Register::EBP
&& index == Register::None
&& disp > 0
&& disp < MAX_FIELD_OFFSET
&& instr.mnemonic() != Mnemonic::Lea
{
cg.field_offsets.insert(disp);
}
}
OpKind::Immediate8
| OpKind::Immediate8to16
| OpKind::Immediate8to32
| OpKind::Immediate16
| OpKind::Immediate32 => {
let v = instr.immediate(i);
if pe.section_of(v).is_some() {
cg.imm_addrs.insert(v);
}
}
_ => {}
}
}
}
fn classify_operand(pe: &Pe, instr: &Instruction, i: u32) -> ArgKind {
match instr.op_kind(i) {
OpKind::Immediate8
| OpKind::Immediate8to16
| OpKind::Immediate8to32
| OpKind::Immediate16
| OpKind::Immediate32 => ArgKind::Imm(instr.immediate(i)),
OpKind::Memory => {
let base = instr.memory_base();
let index = instr.memory_index();
let disp = instr.memory_displacement64();
if base == Register::None && index == Register::None {
if pe.section_of(disp).is_some() {
ArgKind::Global(disp)
} else {
ArgKind::Other
}
} else if base != Register::ESP
&& base != Register::EBP
&& base != Register::None
&& index == Register::None
&& disp > 0
&& disp < MAX_FIELD_OFFSET
{
ArgKind::Field(disp)
} else {
ArgKind::Other
}
}
_ => ArgKind::Other,
}
}
fn classify_reg_src(pe: &Pe, instr: &Instruction) -> ArgKind {
match instr.mnemonic() {
Mnemonic::Lea => {
let base = instr.memory_base();
if base == Register::EBP || base == Register::ESP {
ArgKind::LocalAddr
} else {
ArgKind::Other }
}
Mnemonic::Xor => ArgKind::Imm(0), Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx => classify_operand(pe, instr, 1),
_ => ArgKind::Other,
}
}
fn writes_reg(instr: &Instruction, reg: Register) -> bool {
if instr.op0_kind() != OpKind::Register || instr.op0_register() != reg {
return false;
}
match instr.mnemonic() {
Mnemonic::Mov | Mnemonic::Lea | Mnemonic::Movzx | Mnemonic::Movsx => true,
Mnemonic::Xor => instr.op1_kind() == OpKind::Register && instr.op1_register() == reg,
_ => false,
}
}
pub(crate) fn recover_call_site(pe: &Pe, buf: &[Instruction], target: u64) -> CallSite {
let mut cs = CallSite {
target,
..Default::default()
};
for instr in buf.iter().rev().take(64) {
if is_frame_boundary(instr) {
break;
}
if instr.mnemonic() == Mnemonic::Push {
cs.stack.push(classify_operand(pe, instr, 0));
} else if cs.ecx.is_none() && writes_reg(instr, Register::ECX) {
cs.ecx = Some(classify_reg_src(pe, instr));
} else if cs.edx.is_none() && writes_reg(instr, Register::EDX) {
cs.edx = Some(classify_reg_src(pe, instr));
}
}
cs
}
fn is_frame_boundary(instr: &Instruction) -> bool {
if matches!(instr.mnemonic(), Mnemonic::Add | Mnemonic::Sub)
&& instr.op0_kind() == OpKind::Register
&& instr.op0_register() == Register::ESP
{
return true;
}
(0..instr.op_count())
.any(|i| instr.op_kind(i) == OpKind::Memory && instr.memory_segment() == Register::FS)
}
pub fn resolve_thunk(pe: &Pe, addr: u64) -> Option<u64> {
let bytes = pe.read(addr, 16)?;
let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
if !dec.can_decode() {
return None;
}
let instr = dec.decode();
if instr.flow_control() == FlowControl::UnconditionalBranch {
let t = instr.near_branch_target();
if t != 0 {
return Some(t);
}
}
None
}
fn reads_reg(instr: &Instruction, reg: Register) -> bool {
for i in 0..instr.op_count() {
match instr.op_kind(i) {
OpKind::Memory => {
if instr.memory_base() == reg || instr.memory_index() == reg {
return true;
}
}
OpKind::Register if i > 0 && instr.op_register(i) == reg => {
return true;
}
_ => {}
}
}
false
}
pub fn reg_params(pe: &Pe, addr: u64) -> (bool, bool) {
let Some(bytes) = pe.read(addr, 32 * 32) else {
return (false, false);
};
let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
let mut instr = Instruction::default();
let (mut ecx_written, mut edx_written) = (false, false);
let (mut ecx_param, mut edx_param) = (false, false);
let mut n = 0;
while dec.can_decode() && n < 24 {
dec.decode_out(&mut instr);
if instr.is_invalid() {
break;
}
if !ecx_written && reads_reg(&instr, Register::ECX) {
ecx_param = true;
}
if !edx_written && reads_reg(&instr, Register::EDX) {
edx_param = true;
}
if instr.op0_kind() == OpKind::Register {
match instr.op0_register() {
Register::ECX => ecx_written = true,
Register::EDX => edx_written = true,
_ => {}
}
}
if n == 0
&& matches!(instr.flow_control(), FlowControl::Call)
&& instr.near_branch_target() != instr.next_ip()
{
if !ecx_written {
ecx_param = true;
}
if !edx_written {
edx_param = true;
}
}
if matches!(instr.flow_control(), FlowControl::Call) {
ecx_written = true;
edx_written = true;
}
if !matches!(instr.flow_control(), FlowControl::Next | FlowControl::Call) {
break;
}
n += 1;
}
(ecx_param, edx_param)
}
pub fn analyze(pe: &Pe, entry: u64) -> Result<CallGraph, String> {
analyze_opt(pe, entry, true)
}
pub fn analyze_opt(pe: &Pe, entry: u64, translate_le: bool) -> Result<CallGraph, String> {
static CACHE: OnceLock<AnalyzeCache> = OnceLock::new();
let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
let key = (pe.id(), entry, translate_le);
if let Some(cached) = cache
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
.cloned()
{
return (*cached).clone();
}
let result = analyze_uncached(pe, entry, translate_le);
let mut entries = cache.lock().unwrap_or_else(|e| e.into_inner());
if entries.len() >= 8192 {
entries.clear();
}
entries.insert(key, Arc::new(result.clone()));
result
}
fn analyze_uncached(pe: &Pe, entry: u64, translate_le: bool) -> Result<CallGraph, String> {
if !pe.is_executable_va(entry) {
return Err(format!(
"{entry:#010x} is not inside any PE section of this image\n{}",
pe.describe_sections(entry)
));
}
let mut cg = CallGraph {
entry,
..Default::default()
};
let lo = entry;
let hi = entry.saturating_add(MAX_FN_SPAN);
let mut work = vec![entry];
let mut seen: BTreeSet<u64> = BTreeSet::new();
let mut branch_seed: std::collections::HashMap<u64, Vec<Instruction>> =
std::collections::HashMap::new();
while let Some(start) = work.pop() {
if !seen.insert(start) {
continue;
}
let Some(bytes) = pe.read(start, 32 * 64) else {
continue;
};
let mut dec = Decoder::with_ip(32, bytes, start, DecoderOptions::NONE);
let mut instr = Instruction::default();
let mut block_buf: Vec<Instruction> = branch_seed.remove(&start).unwrap_or_default();
loop {
if !dec.can_decode() {
break;
}
let ip = dec.ip();
if ip != start && seen.contains(&ip) {
break; }
if ip > start {
seen.insert(ip);
}
dec.decode_out(&mut instr);
if instr.is_invalid() {
cg.decode_errors += 1;
if cg.decode_error_samples.len() < 5 {
let bytes = pe.read(ip, 8).unwrap_or(&[]).to_vec();
cg.decode_error_samples.push((ip, bytes));
}
break;
}
cg.max_end = cg.max_end.max(instr.next_ip());
scan_operands(pe, &instr, &mut cg);
block_buf.push(instr);
match instr.flow_control() {
FlowControl::Call => {
let t = instr.near_branch_target();
let fixed = if t == instr.next_ip() {
pe.fixup_rel32_target(instr.ip() + 1)
} else {
None
};
let dest = match fixed {
Some(target) => Some(target),
None if t != 0 && t != instr.next_ip() => Some(t),
None => None,
};
if let Some(target) = dest {
if target != 0 {
*cg.direct_calls.entry(target).or_insert(0) += 1;
let prior = &block_buf[..block_buf.len() - 1]; let cs = recover_call_site(pe, prior, target);
cg.call_sites.push(cs);
}
}
block_buf.clear(); }
FlowControl::IndirectCall => {
cg.indirect_calls += 1;
record_indirect_target(&instr, pe, &mut cg);
block_buf.clear(); }
FlowControl::ConditionalBranch => {
match instr.mnemonic() {
Mnemonic::Jl | Mnemonic::Jle | Mnemonic::Jg | Mnemonic::Jge => {
cg.signed_branches += 1;
}
Mnemonic::Jb | Mnemonic::Jbe | Mnemonic::Ja | Mnemonic::Jae => {
cg.unsigned_branches += 1;
}
Mnemonic::Je | Mnemonic::Jne => {
cg.equality_branches += 1;
}
_ => {}
}
let t = instr.near_branch_target();
if (lo..hi).contains(&t) {
branch_seed.entry(t).or_insert_with(|| block_buf.clone());
work.push(t);
} else {
cg.tail_calls.insert(t);
}
}
FlowControl::UnconditionalBranch => {
let t = instr.near_branch_target();
if (lo..hi).contains(&t) {
if start == lo && block_buf.len() == 1 {
cg.tail_calls.insert(t);
} else {
branch_seed.entry(t).or_insert_with(|| block_buf.clone());
work.push(t);
}
} else {
cg.tail_calls.insert(t);
}
break;
}
FlowControl::IndirectBranch => {
cg.indirect_jumps += 1;
record_indirect_target(&instr, pe, &mut cg);
let mut cases: Vec<u64> = Vec::new();
let has_index = instr.memory_index() != Register::None;
let mut base_written = false;
if instr.memory_base() != Register::None {
for b in block_buf.iter() {
match b.mnemonic() {
Mnemonic::Cmp | Mnemonic::Test => continue,
_ => {}
}
if b.op0_kind() == OpKind::Register
&& b.op0_register() == instr.memory_base()
{
base_written = true;
break;
}
}
}
let unindexed_switch =
!has_index && base_written && is_switch_dispatch(&block_buf);
follow_jump_table(
&instr,
pe,
lo,
hi,
unindexed_switch,
translate_le,
&mut cases,
);
if cases.is_empty()
&& instr.op_count() > 0
&& instr.op_kind(0) == OpKind::Register
&& is_switch_dispatch(&block_buf)
{
if let Some(mut table) =
register_table_base(&block_buf, pe, instr.op_register(0))
{
if translate_le && table != 0 && pe.section_of(table).is_none() {
if let Some(post) = pe.fixup_post_image(table) {
if pe.section_of(post).is_some() {
table = post;
}
}
}
read_table_entries(pe, table, lo, hi, translate_le, &mut cases);
}
}
for c in &cases {
branch_seed.entry(*c).or_insert_with(|| block_buf.clone());
work.push(*c);
}
if unindexed_switch {
let ft = instr.next_ip();
if (lo..hi).contains(&ft) {
branch_seed.entry(ft).or_insert_with(|| block_buf.clone());
work.push(ft);
}
}
break;
}
FlowControl::Return | FlowControl::Exception => break,
FlowControl::Interrupt => {
let dos_terminate = instr.op_count() == 1
&& instr.op0_kind() == OpKind::Immediate8
&& instr.immediate8() == 0x21
&& block_buf
.get(block_buf.len().saturating_sub(2))
.is_some_and(|prior| {
prior.mnemonic() == Mnemonic::Mov
&& prior.op0_kind() == OpKind::Register
&& prior.op0_register() == Register::AX
&& prior.op1_kind() == OpKind::Immediate16
&& prior.immediate16() & 0xff00 == 0x4c00
});
if dos_terminate {
break;
}
}
FlowControl::Next | FlowControl::XbeginXabortXend => {}
}
}
}
Ok(cg)
}
pub fn tail_referenced(pe: &Pe, cg: &CallGraph) -> BTreeSet<u64> {
let mut out = BTreeSet::new();
for t in &cg.tail_calls {
if let Ok(tcg) = analyze(pe, *t) {
out.extend(&tcg.globals);
out.extend(&tcg.imm_addrs);
out.extend(&tcg.iat_calls);
}
}
out
}
pub(crate) fn follow_jump_table(
instr: &Instruction,
pe: &Pe,
lo: u64,
hi: u64,
allow_unindexed: bool,
translate_le: bool,
work: &mut Vec<u64>,
) {
for i in 0..instr.op_count() {
if instr.op_kind(i) != OpKind::Memory {
continue;
}
let mut disp = instr.memory_displacement64();
if translate_le && disp != 0 {
if let Some(post) = pe.fixup_post_image(disp) {
if pe.section_of(post).is_some() {
disp = post;
}
}
}
if instr.memory_index() == Register::None && !allow_unindexed {
return;
}
if pe.section_of(disp).is_none() {
return;
}
read_table_entries(pe, disp, lo, hi, translate_le, work);
return;
}
}
fn read_table_entries(
pe: &Pe,
table_va: u64,
lo: u64,
hi: u64,
translate_le: bool,
work: &mut Vec<u64>,
) {
let mut k = 0u64;
while k < 1024 {
let Some(b) = pe.read(table_va + k * 4, 4) else {
break;
};
if b.len() < 4 {
break;
}
let raw = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as u64;
if raw == 0 {
break; }
let mut target = raw;
if translate_le && !(lo..hi).contains(&target) {
if let Some(post) = pe.fixup_post_image(raw) {
if (lo..hi).contains(&post) {
target = post;
}
}
}
if !(lo..hi).contains(&target) {
break; }
work.push(target);
k += 1;
}
}
fn register_table_base(block: &[Instruction], pe: &Pe, jump_reg: Register) -> Option<u64> {
let mut load_base: Option<Register> = None;
for b in block.iter().rev() {
if matches!(
b.mnemonic(),
Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx
) && b.op0_kind() == OpKind::Register
&& b.op0_register() == jump_reg
&& b.op1_kind() == OpKind::Memory
{
load_base = Some(b.memory_base());
break;
}
}
let base = load_base?;
for b in block.iter().rev() {
match b.mnemonic() {
Mnemonic::Add => {
if b.op0_kind() == OpKind::Register && b.op0_register() == base {
for i in 0..b.op_count() {
if matches!(
b.op_kind(i),
OpKind::Immediate8
| OpKind::Immediate8to16
| OpKind::Immediate8to32
| OpKind::Immediate16
| OpKind::Immediate32
) && pe.section_of(b.immediate(i)).is_some()
{
return Some(b.immediate(i));
}
}
}
}
Mnemonic::Lea => {
if b.op0_kind() == OpKind::Register
&& b.op0_register() == base
&& b.op1_kind() == OpKind::Memory
&& pe.section_of(b.memory_displacement64()).is_some()
{
return Some(b.memory_displacement64());
}
}
Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx
if b.op0_kind() == OpKind::Register && b.op0_register() == base =>
{
for i in 0..b.op_count() {
if matches!(
b.op_kind(i),
OpKind::Immediate8
| OpKind::Immediate8to16
| OpKind::Immediate8to32
| OpKind::Immediate16
| OpKind::Immediate32
) && pe.section_of(b.immediate(i)).is_some()
{
return Some(b.immediate(i));
}
}
}
_ => {}
}
}
None
}
fn is_switch_dispatch(block: &[Instruction]) -> bool {
let mut has_cmp = false;
let mut has_jcc = false;
for i in block {
match i.mnemonic() {
Mnemonic::Cmp | Mnemonic::Test => has_cmp = true,
_ => {}
}
if matches!(i.flow_control(), FlowControl::ConditionalBranch) {
has_jcc = true;
}
}
has_cmp && has_jcc
}
fn record_indirect_target(instr: &Instruction, pe: &Pe, cg: &mut CallGraph) {
for i in 0..instr.op_count() {
if instr.op_kind(i) == OpKind::Memory {
let base = instr.memory_base();
let disp = instr.memory_displacement64();
if base == Register::None && instr.memory_index() == Register::None {
if pe.section_of(disp).is_some() {
cg.iat_calls.insert(disp); }
} else if base != Register::None && disp < MAX_FIELD_OFFSET {
cg.vcall_slots.insert(disp); }
return;
}
}
}