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