Skip to main content

callgraph_verify/
disasm.rs

1//! Recursive-descent function disassembly over real PE bytes (iced-x86).
2//!
3//! Given a function entry VA, we decode instructions and follow the function's
4//! *intra-function* control flow (conditional branches + their fallthrough,
5//! unconditional `JMP rel`) until every reachable path terminates at a `RET` /
6//! interrupt / undecodable byte. Along the way we harvest, purely from the
7//! decoded operands, every fact that lets us catch a rebuild hallucination:
8//!
9//!   * `direct_calls`  — targets of `CALL rel16/rel32` (the call graph proper)
10//!   * `tail_calls`    — `JMP rel` leaving the body (a tail call into another fn)
11//!   * `indirect_calls`/`indirect_jumps` — counts of register/memory dispatch
12//!   * `vcall_slots`   — displacements of `CALL/JMP [reg+0xNN]` (vtable slots)
13//!   * `iat_calls`     — absolute targets of `CALL [0x00a28xxx]` (IAT imports)
14//!   * `field_offsets` — displacements of `[reg+0xNN]` accesses (struct fields),
15//!     excluding ESP/EBP (those are stack locals/args)
16//!   * `globals`       — absolute data addresses referenced as `[0x00bXXXXX]`
17//!   * `imm_addrs`     — in-image immediates (`MOV [esi],0x00a7e38c` vtable VAs,
18//!     function pointers, string-literal addresses)
19//!
20//! This is *structural* ground truth: no decompiler in the loop.
21
22use crate::pe::Pe;
23use iced_x86::{Decoder, DecoderOptions, FlowControl, Instruction, Mnemonic, OpKind, Register};
24use std::collections::{BTreeMap, BTreeSet, HashMap};
25use std::sync::{Arc, Mutex, OnceLock};
26
27type AnalyzeCacheKey = (u64, u64, bool);
28type AnalyzeCache = Mutex<HashMap<AnalyzeCacheKey, Arc<Result<CallGraph, String>>>>;
29
30/// Upper bound on how far past the entry we follow control flow. Real-world
31/// largest functions are a few KB; 64 KB stops a mis-decoded `JMP` from walking
32/// the whole section.
33const MAX_FN_SPAN: u64 = 0x10000;
34/// `[reg+disp]` with disp below this is treated as a struct-field access; larger
35/// displacements are almost certainly absolute addresses or noise.
36const MAX_FIELD_OFFSET: u64 = 0x8000;
37
38/// The recovered "kind" of a single call argument, harvested from the instruction that set it
39/// up (a register load or a `PUSH`). Coarse on purpose — only the distinctions a faithful port
40/// must preserve, and only those a short local scan can recover without false positives.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum ArgKind {
43    /// An immediate constant (`PUSH 0x10`, `MOV ECX,0`, `XOR ECX,ECX`).
44    Imm(u64),
45    /// A struct-field load `[reg+disp]` off a non-stack base (`PUSH [ESI+0xc]`).
46    Field(u64),
47    /// An absolute in-image data reference `[0x00bXXXXX]` (a global).
48    Global(u64),
49    /// Address of a stack local (`LEA ECX,[EBP-8]` / `LEA EAX,[ESP+4]`) — an out-param / temp.
50    LocalAddr,
51    /// Anything not confidently classified (a register value, computed expression, prior result).
52    Other,
53}
54
55/// One direct call site with the argument set-up recovered by a short backward scan of its basic
56/// block. `ecx`/`edx` are the fastcall/thiscall register args (the receiver lands in `ecx`); `stack`
57/// is the pushed args in source (left-to-right) order. `None`/empty means "not recovered in-block".
58#[derive(Debug, Clone, Default)]
59pub struct CallSite {
60    pub target: u64,
61    pub ecx: Option<ArgKind>,
62    pub edx: Option<ArgKind>,
63    pub stack: Vec<ArgKind>,
64}
65
66#[derive(Debug, Default, Clone)]
67pub struct CallGraph {
68    pub entry: u64,
69    /// Direct `CALL` targets → number of call sites hitting that target.
70    pub direct_calls: BTreeMap<u64, u32>,
71    /// Fall-through addresses of decoded direct CALL instructions. These are
72    /// guest return addresses pushed by the instruction, not arbitrary body
73    /// immediates; operand validation may accept them when a port models the
74    /// architectural return address explicitly.
75    pub call_continuations: BTreeSet<u64>,
76    /// Per-call-site recovered argument shapes (one entry per direct `CALL` instruction).
77    pub call_sites: Vec<CallSite>,
78    /// Direct `JMP rel` targets that leave the walked body (tail calls).
79    pub tail_calls: BTreeSet<u64>,
80    pub indirect_calls: u32,
81    pub indirect_jumps: u32,
82    /// vtable slot displacements seen in `CALL/JMP [reg+0xNN]`.
83    pub vcall_slots: BTreeSet<u64>,
84    /// IAT import slot addresses seen in `CALL [0x00a28xxx]`.
85    pub iat_calls: BTreeSet<u64>,
86    /// Struct-field displacements off a non-stack base register.
87    pub field_offsets: BTreeSet<u64>,
88    /// Absolute data addresses referenced as memory operands (globals).
89    pub globals: BTreeSet<u64>,
90    /// In-image immediate values (vtable VAs, fn pointers, string addresses).
91    pub imm_addrs: BTreeSet<u64>,
92    /// Conditional-branch polarity (by condition class): signed (`JL/JLE/JG/JGE`),
93    /// unsigned (`JB/JBE/JA/JAE`), equality (`JE/JNE`). Other conditions
94    /// (`JP/JS/JO/...`) are not counted. Ground truth for the polarity gap:
95    /// the gate checks CFG *structure* but never `==` vs `!=` or signed vs
96    /// unsigned, so an inverted mirror passes. These counts let ports and
97    /// tests pin which class the binary actually uses.
98    pub signed_branches: u32,
99    pub unsigned_branches: u32,
100    pub equality_branches: u32,
101    /// Highest instruction-end VA reached — the function's apparent extent.
102    pub max_end: u64,
103    pub decode_errors: u32,
104    /// First few decode failures as `(ip, raw_bytes)`. Decoders hitting
105    /// invalid opcodes is normal at padding/alignment/data boundaries, but a
106    /// long run inside what is supposed to be a function usually means the
107    /// entry VA is wrong, or the function crosses into embedded data. The
108    /// hex bytes here are enough to tell which.
109    pub decode_error_samples: Vec<(u64, Vec<u8>)>,
110}
111
112impl CallGraph {
113    pub fn call_targets(&self) -> BTreeSet<u64> {
114        self.direct_calls.keys().copied().collect()
115    }
116
117    /// Every in-image address this function references — call targets, tail-call
118    /// JMPs, IAT slots, global data refs and baked immediates (vtable VAs / fn
119    /// pointers / string addresses). Used to detect inlining: if callee C's whole
120    /// footprint also appears in caller F's footprint, the binary pasted C's body
121    /// into F, so a Rust call F→C is faithful (un-inlined), not hallucinated.
122    pub fn footprint(&self) -> BTreeSet<u64> {
123        let mut s: BTreeSet<u64> = self.direct_calls.keys().copied().collect();
124        s.extend(&self.tail_calls);
125        s.extend(&self.iat_calls);
126        s.extend(&self.globals);
127        s.extend(&self.imm_addrs);
128        s
129    }
130
131    /// Addresses the operand gate may accept when a port spells binary facts
132    /// directly, including architectural return PCs pushed by direct CALLs.
133    pub fn operand_references(&self) -> BTreeSet<u64> {
134        let mut s: BTreeSet<u64> = self.direct_calls.keys().copied().collect();
135        s.extend(&self.call_continuations);
136        s.extend(&self.tail_calls);
137        s.extend(&self.iat_calls);
138        s.extend(&self.globals);
139        s.extend(&self.imm_addrs);
140        s
141    }
142}
143
144/// Harvest every memory/immediate operand fact from one decoded instruction.
145fn scan_operands(pe: &Pe, instr: &Instruction, cg: &mut CallGraph) {
146    for i in 0..instr.op_count() {
147        match instr.op_kind(i) {
148            OpKind::Memory => {
149                let base = instr.memory_base();
150                let index = instr.memory_index();
151                let disp = instr.memory_displacement64();
152                if base == Register::None && index == Register::None {
153                    // absolute `[0x...]` — a global/static data reference.
154                    // Skip absolute LEA: it materializes an address but does
155                    // not dereference it (kept out by the existing policy).
156                    if pe.section_of(disp).is_some() && instr.mnemonic() != Mnemonic::Lea {
157                        cg.globals.insert(disp);
158                    }
159                } else if pe.section_of(disp).is_some() {
160                    // Indexed `[base + index*scale + disp]` (or `[reg+disp]`)
161                    // off a large in-image displacement — a table base, not a
162                    // struct field (e.g. `[eax*4+0x18F01]` in a real LE binary). Record it
163                    // as a global so bodies may bake the true address without
164                    // tripping HALLUCINATED-ADDR. Small offsets miss
165                    // `section_of` and fall through to the field branch below.
166                    // LEA address materialization belongs here because its
167                    // result may be dereferenced by a later instruction.
168                    cg.globals.insert(disp);
169                } else if base != Register::None
170                    && base != Register::ESP
171                    && base != Register::EBP
172                    && index == Register::None
173                    && disp > 0
174                    && disp < MAX_FIELD_OFFSET
175                    // LEA computes an address (pointer math / &field), it doesn't
176                    // dereference a struct field — skip it to avoid e.g. the
177                    // `LEA EDX,[EAX+1]` of a strlen loop looking like field 0x1.
178                    && instr.mnemonic() != Mnemonic::Lea
179                {
180                    // `[reg+0xNN]` off a non-stack base — a struct field access.
181                    cg.field_offsets.insert(disp);
182                }
183            }
184            OpKind::Immediate8
185            | OpKind::Immediate8to16
186            | OpKind::Immediate8to32
187            | OpKind::Immediate16
188            | OpKind::Immediate32 => {
189                let v = instr.immediate(i);
190                // In-image immediate = a baked address (vtable VA, fn ptr, string).
191                if pe.section_of(v).is_some() {
192                    cg.imm_addrs.insert(v);
193                }
194            }
195            _ => {}
196        }
197    }
198}
199
200/// Classify operand `i` of `instr` as an argument kind (immediate / field / global / other).
201fn classify_operand(pe: &Pe, instr: &Instruction, i: u32) -> ArgKind {
202    match instr.op_kind(i) {
203        OpKind::Immediate8
204        | OpKind::Immediate8to16
205        | OpKind::Immediate8to32
206        | OpKind::Immediate16
207        | OpKind::Immediate32 => ArgKind::Imm(instr.immediate(i)),
208        OpKind::Memory => {
209            let base = instr.memory_base();
210            let index = instr.memory_index();
211            let disp = instr.memory_displacement64();
212            if base == Register::None && index == Register::None {
213                if pe.section_of(disp).is_some() {
214                    ArgKind::Global(disp)
215                } else {
216                    ArgKind::Other
217                }
218            } else if base != Register::ESP
219                && base != Register::EBP
220                && base != Register::None
221                && index == Register::None
222                && disp > 0
223                && disp < MAX_FIELD_OFFSET
224            {
225                ArgKind::Field(disp)
226            } else {
227                ArgKind::Other
228            }
229        }
230        _ => ArgKind::Other,
231    }
232}
233
234/// Classify the *source* of a register-setup instruction (`MOV reg,src` / `LEA reg,src` /
235/// `XOR reg,reg`). Returns the arg kind the register holds going into the following `CALL`.
236fn classify_reg_src(pe: &Pe, instr: &Instruction) -> ArgKind {
237    match instr.mnemonic() {
238        Mnemonic::Lea => {
239            let base = instr.memory_base();
240            if base == Register::EBP || base == Register::ESP {
241                ArgKind::LocalAddr
242            } else {
243                ArgKind::Other // lea of a field address — not distinguished in v1
244            }
245        }
246        Mnemonic::Xor => ArgKind::Imm(0), // only called when op0==op1 (zero idiom)
247        Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx => classify_operand(pe, instr, 1),
248        _ => ArgKind::Other,
249    }
250}
251
252/// True if `instr` writes register `reg` as its destination (op0) via a recognized setup form.
253fn writes_reg(instr: &Instruction, reg: Register) -> bool {
254    if instr.op0_kind() != OpKind::Register || instr.op0_register() != reg {
255        return false;
256    }
257    match instr.mnemonic() {
258        Mnemonic::Mov | Mnemonic::Lea | Mnemonic::Movzx | Mnemonic::Movsx => true,
259        // XOR reg,reg is the zero idiom only when both operands are the same register.
260        Mnemonic::Xor => instr.op1_kind() == OpKind::Register && instr.op1_register() == reg,
261        _ => false,
262    }
263}
264
265/// Recover the argument shape of a `CALL` from the instructions preceding it in its basic block
266/// (`buf` excludes the CALL itself). Scans backward: the first `PUSH`es encountered are the
267/// left-to-right stack args (cdecl/stdcall push right-to-left, so backward order = source order),
268/// and the first writes to ECX/EDX are the fastcall/thiscall register args.
269pub(crate) fn recover_call_site(pe: &Pe, buf: &[Instruction], target: u64) -> CallSite {
270    let mut cs = CallSite {
271        target,
272        ..Default::default()
273    };
274    for instr in buf.iter().rev().take(64) {
275        // Stop at a prologue / frame boundary: an explicit ESP adjust (`sub esp,N`) or any
276        // FS-relative access (SEH frame set-up). Argument pushes never precede these, so this
277        // keeps callee-saved / SEH-cookie pushes out of the recovered argument list.
278        if is_frame_boundary(instr) {
279            break;
280        }
281        if instr.mnemonic() == Mnemonic::Push {
282            cs.stack.push(classify_operand(pe, instr, 0));
283        } else if cs.ecx.is_none() && writes_reg(instr, Register::ECX) {
284            cs.ecx = Some(classify_reg_src(pe, instr));
285        } else if cs.edx.is_none() && writes_reg(instr, Register::EDX) {
286            cs.edx = Some(classify_reg_src(pe, instr));
287        }
288    }
289    cs
290}
291
292/// A prologue / stack-frame boundary that argument set-up never precedes: an `ADD/SUB ESP,imm`
293/// or any instruction touching FS-relative memory (the SEH frame).
294fn is_frame_boundary(instr: &Instruction) -> bool {
295    if matches!(instr.mnemonic(), Mnemonic::Add | Mnemonic::Sub)
296        && instr.op0_kind() == OpKind::Register
297        && instr.op0_register() == Register::ESP
298    {
299        return true;
300    }
301    (0..instr.op_count())
302        .any(|i| instr.op_kind(i) == OpKind::Memory && instr.memory_segment() == Register::FS)
303}
304
305/// If `addr` is a one-instruction thunk (`JMP target`), return `target`. Callers often invoke a
306/// function through such a thunk (a tail-jump stub, or `/OPT:ICF`-folded forwarder), so the byte
307/// call edge points at the thunk, not the real callee — following it one level lets the call-graph
308/// check see the true target.
309pub fn resolve_thunk(pe: &Pe, addr: u64) -> Option<u64> {
310    let bytes = pe.read(addr, 16)?;
311    let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
312    if !dec.can_decode() {
313        return None;
314    }
315    let instr = dec.decode();
316    if instr.flow_control() == FlowControl::UnconditionalBranch {
317        let t = instr.near_branch_target();
318        if t != 0 {
319            return Some(t);
320        }
321    }
322    None
323}
324
325/// Does `instr` *read* `reg` (as a source register or a memory base/index)? Excludes the
326/// destination register operand (op0) so a pure write like `MOV ECX,EAX` doesn't count as a read.
327fn reads_reg(instr: &Instruction, reg: Register) -> bool {
328    for i in 0..instr.op_count() {
329        match instr.op_kind(i) {
330            OpKind::Memory => {
331                if instr.memory_base() == reg || instr.memory_index() == reg {
332                    return true;
333                }
334            }
335            OpKind::Register if i > 0 && instr.op_register(i) == reg => {
336                return true;
337            }
338            _ => {}
339        }
340    }
341    false
342}
343
344/// Heuristically determine whether the function at `addr` takes register parameters in ECX / EDX
345/// (i.e. is thiscall/fastcall). Scans the entry block for a READ of the register before any write:
346/// thiscall reads ECX (the `this`), fastcall also reads EDX, cdecl/stdcall read neither. Used to
347/// align a caller's recovered register args to the callee's actual parameters (so a stray EDX in a
348/// caller is not mistaken for an argument to a thiscall callee). Conservative: under-detection just
349/// means a position is skipped, never a false mismatch.
350pub fn reg_params(pe: &Pe, addr: u64) -> (bool, bool) {
351    let Some(bytes) = pe.read(addr, 32 * 32) else {
352        return (false, false);
353    };
354    let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
355    let mut instr = Instruction::default();
356    let (mut ecx_written, mut edx_written) = (false, false);
357    let (mut ecx_param, mut edx_param) = (false, false);
358    let mut n = 0;
359    while dec.can_decode() && n < 24 {
360        dec.decode_out(&mut instr);
361        if instr.is_invalid() {
362            break;
363        }
364        if !ecx_written && reads_reg(&instr, Register::ECX) {
365            ecx_param = true;
366        }
367        if !edx_written && reads_reg(&instr, Register::EDX) {
368            edx_param = true;
369        }
370        if instr.op0_kind() == OpKind::Register {
371            match instr.op0_register() {
372                Register::ECX => ecx_written = true,
373                Register::EDX => edx_written = true,
374                _ => {}
375            }
376        }
377        // A leaf that immediately forwards its incoming registers to another call is still
378        // thiscall/fastcall even though iced-x86 cannot represent the implicit register use on
379        // the CALL itself. Preserve untouched ECX/EDX on the first instruction so constructors
380        // such as BSScissorTriShape (which forwards both registers to its base ctor) align stack
381        // arguments correctly. Later calls retain the caller-saved rule below.
382        if n == 0
383            && matches!(instr.flow_control(), FlowControl::Call)
384            && instr.near_branch_target() != instr.next_ip()
385        {
386            if !ecx_written {
387                ecx_param = true;
388            }
389            if !edx_written {
390                edx_param = true;
391            }
392        }
393        // Per the x86 cdecl/fastcall/thiscall conventions, ECX and EDX are caller-saved:
394        // any CALL in the entry block clobbers them, so a later READ of ECX/EDX is not
395        // evidence that the callee takes that register as a parameter. Mark them written
396        // across the CALL — without this, a stray `MOV EAX, [ECX]` after a prior CALL
397        // would falsely set `ecx_param = true` and pollute ARG-SHAPE diffs.
398        if matches!(instr.flow_control(), FlowControl::Call) {
399            ecx_written = true;
400            edx_written = true;
401        }
402        if !matches!(instr.flow_control(), FlowControl::Next | FlowControl::Call) {
403            break;
404        }
405        n += 1;
406    }
407    (ecx_param, edx_param)
408}
409
410/// Walk the function at `entry` and return its facts, sourced purely from bytes.
411pub fn analyze(pe: &Pe, entry: u64) -> Result<CallGraph, String> {
412    analyze_opt(pe, entry, true)
413}
414
415/// [`analyze`] with control over LE jump-table translation. With
416/// `translate_le = false` a table's pre-fixup base/entries are read raw
417/// (the pre-0.1.8 behaviour: LE arms are invisible). The gate uses both:
418/// the `false` set drives `MISSING` (no new demands), the `true` set drives
419/// acceptance (spelling an arm is faithful).
420pub fn analyze_opt(pe: &Pe, entry: u64, translate_le: bool) -> Result<CallGraph, String> {
421    // The macro asks for the same entry through several checks (call diff,
422    // operand facts, tail-call expansion, and inlining detection). Keep the
423    // decoded result per loaded image and mode so those consumers share the
424    // expensive recursive-descent walk. Pe IDs are assigned at parse time,
425    // so same-layout binaries with different bytes cannot alias.
426    static CACHE: OnceLock<AnalyzeCache> = OnceLock::new();
427    let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
428    let key = (pe.id(), entry, translate_le);
429    if let Some(cached) = cache
430        .lock()
431        .unwrap_or_else(|e| e.into_inner())
432        .get(&key)
433        .cloned()
434    {
435        return (*cached).clone();
436    }
437
438    let result = analyze_uncached(pe, entry, translate_le);
439    let mut entries = cache.lock().unwrap_or_else(|e| e.into_inner());
440    // This is a process-local build cache, not a general purpose disassembly
441    // database. Bound it for long-lived library users that analyze many
442    // unrelated images/functions.
443    if entries.len() >= 8192 {
444        entries.clear();
445    }
446    entries.insert(key, Arc::new(result.clone()));
447    result
448}
449
450fn analyze_uncached(pe: &Pe, entry: u64, translate_le: bool) -> Result<CallGraph, String> {
451    if !pe.is_executable_va(entry) {
452        return Err(format!(
453            "{entry:#010x} is not inside any PE section of this image\n{}",
454            pe.describe_sections(entry)
455        ));
456    }
457    let mut cg = CallGraph {
458        entry,
459        ..Default::default()
460    };
461    let lo = entry;
462    let hi = entry.saturating_add(MAX_FN_SPAN);
463
464    let mut work = vec![entry];
465    let mut seen: BTreeSet<u64> = BTreeSet::new();
466    // Per-branch-target seed: when we encounter `Jcc target` / `JMP target` /
467    // indirect-jump case, the args set up before the branch (PUSHes, MOV-to-ECX,
468    // etc.) also apply to the target block's first CALL — the predecessor and
469    // target share the same register/stack state at the branch instruction. The
470    // outer loop resets `block_buf` when it pops a new `start`, so without this
471    // seed map the target's calls would be processed with an empty `block_buf`
472    // and `recover_call_site` would undercount arguments. (BUG #8 fix.)
473    //
474    // Multiple predecessors may push the same target — `entry().or_insert_with`
475    // keeps the FIRST push, since the args are equivalent along either path up
476    // to the branch instruction itself (the differing branches only affect the
477    // value of flags, which `recover_call_site` doesn't read).
478    let mut branch_seed: std::collections::HashMap<u64, Vec<Instruction>> =
479        std::collections::HashMap::new();
480
481    while let Some(start) = work.pop() {
482        if !seen.insert(start) {
483            continue;
484        }
485        let Some(bytes) = pe.read(start, 32 * 64) else {
486            continue;
487        };
488        let mut dec = Decoder::with_ip(32, bytes, start, DecoderOptions::NONE);
489        let mut instr = Instruction::default();
490        // Straight-line instruction buffer for the current block, used to recover the argument
491        // set-up preceding each CALL. Reset at block boundaries (branches) and after each call.
492        // For non-entry `start`s, seed from the predecessor branch (see `branch_seed`).
493        let mut block_buf: Vec<Instruction> = branch_seed.remove(&start).unwrap_or_default();
494
495        loop {
496            if !dec.can_decode() {
497                break;
498            }
499            let ip = dec.ip();
500            if ip != start && seen.contains(&ip) {
501                break; // block already covered by another path
502            }
503            if ip > start {
504                seen.insert(ip);
505            }
506            dec.decode_out(&mut instr);
507            if instr.is_invalid() {
508                cg.decode_errors += 1;
509                // Keep up to 5 samples so the user can tell at a glance
510                // whether the failures are padding (single 0x00/0xCC/0x90),
511                // an aligned jump table (4 distinct bytes), or something
512                // structural (random bytes that decode to nothing — entry
513                // VA is probably wrong).
514                if cg.decode_error_samples.len() < 5 {
515                    let bytes = pe.read(ip, 8).unwrap_or(&[]).to_vec();
516                    cg.decode_error_samples.push((ip, bytes));
517                }
518                break;
519            }
520            cg.max_end = cg.max_end.max(instr.next_ip());
521            scan_operands(pe, &instr, &mut cg);
522            block_buf.push(instr);
523
524            match instr.flow_control() {
525                FlowControl::Call => {
526                    let t = instr.near_branch_target();
527                    // Every decoded near CALL pushes next_ip, even when a
528                    // call-next/get-PC pattern has no recoverable callee.
529                    cg.call_continuations.insert(instr.next_ip());
530                    // `call $+5` (rel32 == 0, target == next_ip) is either a
531                    // genuine get-PC/fall-through or a load-time relocation
532                    // site (LE 08 self-relative fixup) whose on-disk field is
533                    // a zero placeholder. Prefer the fixup's runtime target
534                    // (operand sits at ip+1 for `E8 rel32`).
535                    let fixed = if t == instr.next_ip() {
536                        pe.fixup_rel32_target(instr.ip() + 1)
537                    } else {
538                        None
539                    };
540                    let dest = match fixed {
541                        Some(target) => Some(target),
542                        None if t != 0 && t != instr.next_ip() => Some(t),
543                        None => None,
544                    };
545                    if let Some(target) = dest {
546                        if target != 0 {
547                            *cg.direct_calls.entry(target).or_insert(0) += 1;
548                            let prior = &block_buf[..block_buf.len() - 1]; // exclude the CALL itself
549                            let cs = recover_call_site(pe, prior, target);
550                            cg.call_sites.push(cs);
551                        }
552                    }
553                    block_buf.clear(); // args after a call start fresh
554                }
555                FlowControl::IndirectCall => {
556                    cg.indirect_calls += 1;
557                    record_indirect_target(&instr, pe, &mut cg);
558                    block_buf.clear(); // an indirect call's args must not leak into the next call
559                }
560                FlowControl::ConditionalBranch => {
561                    // Polarity ground truth (one count per branch instruction;
562                    // each ip is visited once via `seen`, so no double-count).
563                    match instr.mnemonic() {
564                        Mnemonic::Jl | Mnemonic::Jle | Mnemonic::Jg | Mnemonic::Jge => {
565                            cg.signed_branches += 1;
566                        }
567                        Mnemonic::Jb | Mnemonic::Jbe | Mnemonic::Ja | Mnemonic::Jae => {
568                            cg.unsigned_branches += 1;
569                        }
570                        Mnemonic::Je | Mnemonic::Jne => {
571                            cg.equality_branches += 1;
572                        }
573                        _ => {}
574                    }
575                    let t = instr.near_branch_target();
576                    if (lo..hi).contains(&t) {
577                        branch_seed.entry(t).or_insert_with(|| block_buf.clone());
578                        work.push(t);
579                    } else {
580                        cg.tail_calls.insert(t);
581                    }
582                    // Don't clear `block_buf` here: the args set up before the
583                    // conditional apply to BOTH the taken and the fallthrough
584                    // paths (they share register/stack state at the branch).
585                    // Clearing would make the fallthrough's first CALL recover
586                    // an empty argument set. The buffer is reset by the next
587                    // CALL we encounter (Call/IndirectCall arms above clear
588                    // post-call) or by an explicit boundary.
589                }
590                FlowControl::UnconditionalBranch => {
591                    let t = instr.near_branch_target();
592                    if (lo..hi).contains(&t) {
593                        // Single-JMP entry stub: a function whose entire
594                        // entry block is one unconditional JMP (sibling-call
595                        // `JMP tail_target`, ICF forwarder, `jmp` over no
596                        // bytes) transfers control elsewhere instead of
597                        // containing its own body. Following it would merge
598                        // the callee's body (and globals) into the caller,
599                        // hiding the tail edge the gate checks. Scoped to
600                        // exactly one instruction so a `JMP` that skips an
601                        // inline jump table after a prologue (multi-instr
602                        // entry block) keeps the follow behavior below.
603                        if start == lo && block_buf.len() == 1 {
604                            cg.tail_calls.insert(t);
605                        } else {
606                            branch_seed.entry(t).or_insert_with(|| block_buf.clone());
607                            work.push(t);
608                        }
609                    } else {
610                        cg.tail_calls.insert(t);
611                    }
612                    break;
613                }
614                FlowControl::IndirectBranch => {
615                    cg.indirect_jumps += 1;
616                    record_indirect_target(&instr, pe, &mut cg);
617                    let mut cases: Vec<u64> = Vec::new();
618                    // Indirect-jump target recovery, by operand shape:
619                    // * SIB-indexed `[base+idx*s+TABLE]` (MSVC-style): the
620                    //   existing follow_jump_table read below. Unchanged.
621                    // * Unindexed `[reg+TABLE]` with NO base register, i.e.
622                    //   an absolute `[disp]` (IAT thunks, import jumps):
623                    //   NEVER followed. The disp names data, not code, and
624                    //   the fallthrough past a non-returning jump is dead
625                    //   code -- decoding it attributes neighboring
626                    //   functions' calls here (verified: 9KB/20-call blowup
627                    //   on a real PE at 0x006ff480, whose IAT tail-jmps sit
628                    //   after cmp/jcc).
629                    // * Unindexed `[reg+TABLE]` WITH a base register: Watcom
630                    //   switch lowering (`shl reg,2` then `jmp [reg+TABLE]`,
631                    //   e.g. an LE binary at 0x602c0). The table entries are
632                    //   pre-fixup bytes on LE, so the read below yields
633                    //   nothing; instead continue decoding at the fallthrough
634                    //   address -- the case bodies tile the bytes right after
635                    //   the dispatch (verified: 602c0's five hidden 7d6ae
636                    //   calls, previously mis-reported as HALLUCINATED).
637                    //   Gated on switch evidence (a compare/test plus a
638                    //   conditional branch in this block) so vtable tail-jumps
639                    //   like `jmp [ecx+0x44]` keep today's behavior.
640                    let has_index = instr.memory_index() != Register::None;
641                    // The base must be produced locally (e.g. the scaled
642                    // index): an incoming register like a vtable `this`
643                    // is never a switch table.
644                    let mut base_written = false;
645                    if instr.memory_base() != Register::None {
646                        // CMP/TEST only write flags, not their operands --
647                        // without this carve-out the evidence check itself
648                        // (`cmp reg,...`) would trivially satisfy "written".
649                        for b in block_buf.iter() {
650                            match b.mnemonic() {
651                                Mnemonic::Cmp | Mnemonic::Test => continue,
652                                _ => {}
653                            }
654                            if b.op0_kind() == OpKind::Register
655                                && b.op0_register() == instr.memory_base()
656                            {
657                                base_written = true;
658                                break;
659                            }
660                        }
661                    }
662                    let unindexed_switch =
663                        !has_index && base_written && is_switch_dispatch(&block_buf);
664                    follow_jump_table(
665                        &instr,
666                        pe,
667                        lo,
668                        hi,
669                        unindexed_switch,
670                        translate_le,
671                        &mut cases,
672                    );
673                    // GCC -O0 switch lowering keeps the scale in a separate
674                    // `shl`, adds the table base as an immediate, loads the
675                    // entry, then jumps through a register:
676                    //   shl $2,%eax; add $TABLE,%eax; mov (%eax),%eax; jmp *%eax
677                    // The JMP itself has no memory operand, so the readers
678                    // above see nothing (MinGW `c_callgraph_od` classify:
679                    // 1 indirect jump, 0 cases, 0 calls). Recover the table
680                    // base from the `add` and read it like any other table,
681                    // gated on the same switch evidence so a vtable
682                    // `mov eax,[ecx+slot]; jmp eax` (no bound check, `this`
683                    // never `add`ed with an in-image base) stays untouched.
684                    if cases.is_empty()
685                        && instr.op_count() > 0
686                        && instr.op_kind(0) == OpKind::Register
687                        && is_switch_dispatch(&block_buf)
688                    {
689                        if let Some(mut table) =
690                            register_table_base(&block_buf, pe, instr.op_register(0))
691                        {
692                            // LE bases are pre-fixup too; translate only when
693                            // the raw base is not already a valid section VA.
694                            if translate_le && table != 0 && pe.section_of(table).is_none() {
695                                if let Some(post) = pe.fixup_post_image(table) {
696                                    if pe.section_of(post).is_some() {
697                                        table = post;
698                                    }
699                                }
700                            }
701                            read_table_entries(pe, table, lo, hi, translate_le, &mut cases);
702                        }
703                    }
704                    for c in &cases {
705                        branch_seed.entry(*c).or_insert_with(|| block_buf.clone());
706                        work.push(*c);
707                    }
708                    if unindexed_switch {
709                        let ft = instr.next_ip();
710                        if (lo..hi).contains(&ft) {
711                            branch_seed.entry(ft).or_insert_with(|| block_buf.clone());
712                            work.push(ft);
713                        }
714                    }
715                    break;
716                }
717                FlowControl::Return | FlowControl::Exception => break,
718                // Most `int n` instructions are returning DOS/DPMI API calls,
719                // so keep decoding at next_ip(). DOS INT 21h/AH=4Ch terminates
720                // the process, though; following that fallthrough would merge
721                // the next adjacent function into this one.
722                FlowControl::Interrupt => {
723                    let dos_terminate = instr.op_count() == 1
724                        && instr.op0_kind() == OpKind::Immediate8
725                        && instr.immediate8() == 0x21
726                        && block_buf
727                            .get(block_buf.len().saturating_sub(2))
728                            .is_some_and(|prior| {
729                                prior.mnemonic() == Mnemonic::Mov
730                                    && prior.op0_kind() == OpKind::Register
731                                    && prior.op0_register() == Register::AX
732                                    && prior.op1_kind() == OpKind::Immediate16
733                                    && prior.immediate16() & 0xff00 == 0x4c00
734                            });
735                    if dos_terminate {
736                        break;
737                    }
738                }
739                FlowControl::Next | FlowControl::XbeginXabortXend => {}
740            }
741        }
742    }
743    Ok(cg)
744}
745
746/// Addresses a tail-call target references (its globals, baked immediates,
747/// and IAT slots), unioned over every `JMP`-out edge of `cg`.
748///
749/// A tail-call `JMP G` transfers control to G, so G's data references execute
750/// as part of THIS function. The call-graph diff (`diff_body_calls`) already
751/// counts a tail target's *callees* as acceptable spellings; this is the
752/// address-side mirror for the HALLUCINATED-ADDR check: a port that
753/// duplicates a shared tail body (e.g. a predicate tail-jumped into from two
754/// origins) bakes the tail's addresses faithfully. Only ever ADDS acceptance
755/// — a target that fails to analyze is skipped, never an error.
756pub fn tail_referenced(pe: &Pe, cg: &CallGraph) -> BTreeSet<u64> {
757    let mut out = BTreeSet::new();
758    for t in &cg.tail_calls {
759        if let Ok(tcg) = analyze(pe, *t) {
760            out.extend(&tcg.globals);
761            out.extend(&tcg.imm_addrs);
762            out.extend(&tcg.iat_calls);
763        }
764    }
765    out
766}
767
768/// Follow an MSVC-style `switch` jump table: `JMP dword ptr [idx*4 + TABLE_VA]`.
769/// `TABLE_VA` (the operand displacement, an in-image .rdata address) holds an
770/// array of u32 case addresses. We read entries until one falls outside the
771/// function `[lo, hi)` span -- switch cases always point back into the function,
772/// so that's a reliable terminator. Recovers the case bodies (and their `CALL`s)
773/// that a plain recursive descent, unable to follow an indirect jump, would miss.
774///
775/// `allow_unindexed` covers Watcom-style lowering where the scale lives in a
776/// separate `shl` and the dispatch is a plain `[reg+TABLE]` with no SIB index
777/// (e.g. an LE binary at `0x602c0`: `jmp dword ptr cs:[eax+0x502D8]`). Only set it
778/// when the block already looks like a switch (see `is_switch_dispatch`):
779/// without that evidence a vtable tail-jump would read arbitrary words as
780/// "cases". On LE binaries the table entries are pre-fixup bytes, so the
781/// read yields no in-range targets and the fallthrough arm below carries
782/// the recovery instead.
783pub(crate) fn follow_jump_table(
784    instr: &Instruction,
785    pe: &Pe,
786    lo: u64,
787    hi: u64,
788    allow_unindexed: bool,
789    translate_le: bool,
790    work: &mut Vec<u64>,
791) {
792    for i in 0..instr.op_count() {
793        if instr.op_kind(i) != OpKind::Memory {
794            continue;
795        }
796        let mut disp = instr.memory_displacement64();
797        // LE stores the table base pre-fixup (e.g. `jmp [ecx*4+0x549A9]`);
798        // translate it to the runtime VA before the section check and the
799        // page-map read. A LE absolute operand is always a link-time value, so
800        // a fixup-site displacement is translated whenever its post image is a
801        // valid section VA (a `0` terminator/offset is excluded). PE has no
802        // fixup tables, so `fixup_post_image` is None and this is a no-op.
803        if translate_le && disp != 0 {
804            if let Some(post) = pe.fixup_post_image(disp) {
805                if pe.section_of(post).is_some() {
806                    disp = post;
807                }
808            }
809        }
810        if instr.memory_index() == Register::None && !allow_unindexed {
811            return;
812        }
813        if pe.section_of(disp).is_none() {
814            return;
815        }
816        read_table_entries(pe, disp, lo, hi, translate_le, work);
817        return;
818    }
819}
820
821/// Read a u32 VA jump table at `table_va`, pushing entries while they point
822/// back into the function `[lo, hi)` (the terminator shared by every table
823/// shape: switch cases always target the function's own bodies).
824fn read_table_entries(
825    pe: &Pe,
826    table_va: u64,
827    lo: u64,
828    hi: u64,
829    translate_le: bool,
830    work: &mut Vec<u64>,
831) {
832    let mut k = 0u64;
833    while k < 1024 {
834        let Some(b) = pe.read(table_va + k * 4, 4) else {
835            break;
836        };
837        if b.len() < 4 {
838            break;
839        }
840        let raw = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as u64;
841        if raw == 0 {
842            break; // zero terminator (and pre=0 fixup sites are not cases)
843        }
844        // LE table entries are pre-fixup case addresses; translate each to
845        // its runtime VA *only* when the raw value is out of the function
846        // span and the post image is inside it. In-span entries (PE, or
847        // already-post values) pass through unchanged.
848        let mut target = raw;
849        if translate_le && !(lo..hi).contains(&target) {
850            if let Some(post) = pe.fixup_post_image(raw) {
851                if (lo..hi).contains(&post) {
852                    target = post;
853                }
854            }
855        }
856        if !(lo..hi).contains(&target) {
857            break; // first entry outside the function ends the table
858        }
859        work.push(target);
860        k += 1;
861    }
862}
863
864/// Table base for a register-indirect switch dispatch (`jmp *%reg`).
865/// Finds the entry load `mov %reg,[%base]` nearest the jump, then the
866/// `add %base,$TABLE` / `lea %base,[...]` / `mov %base,$TABLE` that
867/// established it with an in-image table address. Returns `None` when the
868/// jump register comes straight from memory with no locally-built table
869/// (e.g. a vtable `mov %eax,[%ecx+slot]; jmp *%eax`, whose `%ecx` is an
870/// incoming `this` never `add`ed with an in-image base).
871fn register_table_base(block: &[Instruction], pe: &Pe, jump_reg: Register) -> Option<u64> {
872    // Nearest load into the jump register: `mov %reg,[%base]`.
873    let mut load_base: Option<Register> = None;
874    for b in block.iter().rev() {
875        if matches!(
876            b.mnemonic(),
877            Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx
878        ) && b.op0_kind() == OpKind::Register
879            && b.op0_register() == jump_reg
880            && b.op1_kind() == OpKind::Memory
881        {
882            load_base = Some(b.memory_base());
883            break;
884        }
885    }
886    let base = load_base?;
887    // The table address built into the base register further up.
888    for b in block.iter().rev() {
889        match b.mnemonic() {
890            Mnemonic::Add => {
891                if b.op0_kind() == OpKind::Register && b.op0_register() == base {
892                    for i in 0..b.op_count() {
893                        if matches!(
894                            b.op_kind(i),
895                            OpKind::Immediate8
896                                | OpKind::Immediate8to16
897                                | OpKind::Immediate8to32
898                                | OpKind::Immediate16
899                                | OpKind::Immediate32
900                        ) && pe.section_of(b.immediate(i)).is_some()
901                        {
902                            return Some(b.immediate(i));
903                        }
904                    }
905                }
906            }
907            Mnemonic::Lea => {
908                if b.op0_kind() == OpKind::Register
909                    && b.op0_register() == base
910                    && b.op1_kind() == OpKind::Memory
911                    && pe.section_of(b.memory_displacement64()).is_some()
912                {
913                    return Some(b.memory_displacement64());
914                }
915            }
916            Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx
917                if b.op0_kind() == OpKind::Register && b.op0_register() == base =>
918            {
919                for i in 0..b.op_count() {
920                    if matches!(
921                        b.op_kind(i),
922                        OpKind::Immediate8
923                            | OpKind::Immediate8to16
924                            | OpKind::Immediate8to32
925                            | OpKind::Immediate16
926                            | OpKind::Immediate32
927                    ) && pe.section_of(b.immediate(i)).is_some()
928                    {
929                        return Some(b.immediate(i));
930                    }
931                }
932            }
933            _ => {}
934        }
935    }
936    None
937}
938
939/// Switch-dispatch evidence: a compare/test plus a conditional branch in
940/// the current straight-line block (e.g. `cmp edx,6; ja default` ahead of
941/// `jmp [reg+TABLE]`). Used to license fallthrough decoding past an
942/// indirect jump only where a switch table is plausible, so vtable
943/// tail-jumps (no bound check) never gain fallthrough coverage.
944fn is_switch_dispatch(block: &[Instruction]) -> bool {
945    let mut has_cmp = false;
946    let mut has_jcc = false;
947    for i in block {
948        match i.mnemonic() {
949            Mnemonic::Cmp | Mnemonic::Test => has_cmp = true,
950            _ => {}
951        }
952        if matches!(i.flow_control(), FlowControl::ConditionalBranch) {
953            has_jcc = true;
954        }
955    }
956    has_cmp && has_jcc
957}
958
959/// Classify an indirect `CALL`/`JMP [...]` target: `[reg+disp]` is a vtable slot
960/// (record the displacement); `[abs]` in the IAT range is an import (record addr).
961fn record_indirect_target(instr: &Instruction, pe: &Pe, cg: &mut CallGraph) {
962    // The branch target operand is the (single) memory operand, if any.
963    for i in 0..instr.op_count() {
964        if instr.op_kind(i) == OpKind::Memory {
965            let base = instr.memory_base();
966            let disp = instr.memory_displacement64();
967            if base == Register::None && instr.memory_index() == Register::None {
968                if pe.section_of(disp).is_some() {
969                    cg.iat_calls.insert(disp); // CALL [0x00a28xxx]
970                }
971            } else if base != Register::None && disp < MAX_FIELD_OFFSET {
972                cg.vcall_slots.insert(disp); // CALL [reg+slot]
973            }
974            return;
975        }
976    }
977}