taint/
taint.rs

1//! A non-trivial consumer of the materialized ctree, used as a stress test and a benchmark.
2//!
3//! For every function in the database it runs a crude intra-procedural taint pass
4//! (call-return sources -> local def/use fixpoint -> dangerous-call sinks) and times the
5//! work in four separate phases:
6//!
7//!   decompile   -- Hex-Rays, kernel thread, serial and unavoidable
8//!   extract     -- `cfunc.ctree()`, the facade DFS + Rust rebuild, kernel thread
9//!   resolve     -- turning `Obj(address)` callees into names, kernel thread (needs `&Database`)
10//!   analyze     -- the pure taint pass over the Send image (no kernel access)
11//!
12//! The split answers the live question: how much of the work is the serial kernel
13//! floor (decompile + extract + resolve) versus the part that could ever be fanned
14//! out (analyze)? If `analyze` is a rounding error, parallelism buys nothing and we
15//! stay sequential. The run also exercises every node kind, so it is where API
16//! friction shows up.
17//!
18//! Run (release matters for the numbers):
19//!   `cargo run -p idakit --release --example taint -- <db.i64>`
20//! Cap the sweep with `TAINT_LIMIT=2000` while iterating.
21
22use std::collections::{HashMap, HashSet};
23use std::time::{Duration, Instant};
24
25use idakit::decompiler::ctree::{Ctree, ExpressionId, LocalId, NodeRef};
26use idakit::prelude::*;
27
28/// Calls whose return value introduces taint (matched as a substring of the name).
29const SOURCES: &[&str] = &["recv", "read", "fgets", "getenv", "scanf", "gets"];
30/// Calls whose arguments are dangerous to feed tainted data into.
31const SINKS: &[&str] = &[
32    "memcpy", "memmove", "strcpy", "strcat", "sprintf", "system", "malloc", "alloca", "exec",
33];
34
35fn matches(name: &str, set: &[&str]) -> bool {
36    set.iter().any(|n| name.contains(n))
37}
38
39/// A function lifted into a Send-able analysis input.
40///
41/// The callee-name map is the telling part, since the ctree's `Call` carries only an
42/// `Obj(address)`/`Helper`, so the names every analysis actually keys on have to be resolved
43/// *here*, on the kernel thread, and folded in. The bare tree cannot answer "what does this
44/// call?" off-thread.
45struct TaintInput {
46    tree: Ctree,
47    /// Call expression -> resolved callee name (only the ones that resolve to a symbol).
48    callees: HashMap<ExpressionId, String>,
49}
50
51/// Resolves a call's callee to a name, if it is a direct symbol or a decompiler helper.
52///
53/// Indirect calls (through a variable or computed pointer) stay unresolved. The symbol name
54/// rides on the `Obj` node, so this needs no kernel access.
55fn callee_name(tree: &Ctree, callee: ExpressionId) -> Option<String> {
56    let kind = tree.kind(callee);
57    if let Some((_, name)) = kind.as_obj() {
58        return name.map(str::to_owned);
59    }
60    kind.as_helper().map(str::to_owned)
61}
62
63/// Builds the callee-name map for a tree.
64///
65/// The tree carries callee names directly, so this is pure: no `Database` access, no
66/// kernel-thread name resolution.
67fn resolve_callees(tree: &Ctree) -> HashMap<ExpressionId, String> {
68    let mut map = HashMap::new();
69    for (id, callee, _) in tree.calls() {
70        if let Some(name) = callee_name(tree, callee) {
71            map.insert(id, name);
72        }
73    }
74    map
75}
76
77/// Whether `e`'s subtree reads a tainted local or calls a source directly.
78///
79/// Flow-insensitive and deliberately crude, doing real work proportional to tree size.
80fn expression_tainted(img: &TaintInput, e: ExpressionId, tainted: &HashSet<u32>) -> bool {
81    img.tree
82        .expression_descendants(NodeRef::Expression(e))
83        .any(|id| match img.tree.kind(id).as_var() {
84            Some(LocalId(i)) => tainted.contains(&i),
85            None => img.callees.get(&id).is_some_and(|n| matches(n, SOURCES)),
86        })
87}
88
89/// The pure phase returns the number of source->sink flows found. No `Database` access,
90/// so this is exactly the work that could move to a worker thread.
91fn analyze(img: &TaintInput) -> usize {
92    // Collect `Var(i) = rhs` definitions once.
93    let defs: Vec<(u32, ExpressionId)> = img
94        .tree
95        .assigns()
96        .filter_map(|(_, _, x, y)| {
97            let LocalId(i) = img.tree.kind(x).as_var()?;
98            Some((i, y))
99        })
100        .collect();
101
102    // Fixpoint: a local is tainted once any of its defining RHS is tainted.
103    let mut tainted: HashSet<u32> = HashSet::new();
104    loop {
105        let before = tainted.len();
106        for &(lv, rhs) in &defs {
107            if !tainted.contains(&lv) && expression_tainted(img, rhs, &tainted) {
108                tainted.insert(lv);
109            }
110        }
111        if tainted.len() == before {
112            break;
113        }
114    }
115
116    // Sinks: a tainted argument to a dangerous call is a flow.
117    img.tree
118        .calls()
119        .filter(|(id, _, args)| {
120            img.callees.get(id).is_some_and(|n| matches(n, SINKS))
121                && args.iter().any(|a| expression_tainted(img, *a, &tainted))
122        })
123        .count()
124}
125
126#[derive(Default)]
127struct Totals {
128    decompile: Duration,
129    extract: Duration,
130    resolve: Duration,
131    analyze: Duration,
132    funcs: usize,
133    decompile_failed: usize,
134    extract_failed: usize,
135    flows: usize,
136    nodes: u64,
137}
138
139fn run(idb: &mut Database, db: &str) -> Result<(), Error> {
140    idb.open(db).call()?;
141
142    let limit = std::env::var("TAINT_LIMIT")
143        .ok()
144        .and_then(|s| s.parse::<usize>().ok())
145        .unwrap_or(usize::MAX);
146
147    let eas: Vec<Address> = idb.functions().map(|f| f.address()).take(limit).collect();
148    println!("[taint] sweeping {} functions", eas.len());
149
150    let mut t = Totals::default();
151    let wall = Instant::now();
152
153    for (i, &address) in eas.iter().enumerate() {
154        let started = Instant::now();
155        let Ok(cf) = idb.decompile(address) else {
156            t.decompile_failed += 1;
157            continue;
158        };
159        t.decompile += started.elapsed();
160
161        let started = Instant::now();
162        let Ok(tree) = cf.ctree() else {
163            t.extract_failed += 1;
164            continue;
165        };
166        t.extract += started.elapsed();
167
168        let started = Instant::now();
169        let callees = resolve_callees(&tree);
170        t.resolve += started.elapsed();
171
172        t.nodes += (tree.expressions().count() + tree.statements().count()) as u64;
173        let img = TaintInput { tree, callees };
174
175        let started = Instant::now();
176        t.flows += analyze(&img);
177        t.analyze += started.elapsed();
178
179        t.funcs += 1;
180        if (i + 1) % 5000 == 0 {
181            println!("[taint] {} / {} ...", i + 1, eas.len());
182        }
183    }
184
185    report(&t, wall.elapsed());
186    idb.close(false);
187    Ok(())
188}
189
190fn report(t: &Totals, wall: Duration) {
191    let kernel = t.decompile + t.extract + t.resolve;
192    let pct = |d: Duration| 100.0 * d.as_secs_f64() / wall.as_secs_f64().max(f64::EPSILON);
193    let per = |d: Duration| {
194        if t.funcs == 0 {
195            0.0
196        } else {
197            d.as_secs_f64() * 1e6 / t.funcs as f64
198        }
199    };
200
201    println!("\n=== taint sweep ===");
202    println!(
203        "functions analyzed: {}  (decompile-failed {}, extract-failed {})",
204        t.funcs, t.decompile_failed, t.extract_failed
205    );
206    println!("ctree nodes total:  {}", t.nodes);
207    println!("source->sink flows: {}", t.flows);
208    println!("wall:               {:.2}s", wall.as_secs_f64());
209    println!(
210        "  decompile  {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
211        t.decompile.as_secs_f64(),
212        pct(t.decompile),
213        per(t.decompile)
214    );
215    println!(
216        "  extract    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
217        t.extract.as_secs_f64(),
218        pct(t.extract),
219        per(t.extract)
220    );
221    println!(
222        "  resolve    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
223        t.resolve.as_secs_f64(),
224        pct(t.resolve),
225        per(t.resolve)
226    );
227    println!(
228        "  analyze    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn   (the only parallelizable part)",
229        t.analyze.as_secs_f64(),
230        pct(t.analyze),
231        per(t.analyze)
232    );
233    println!(
234        "kernel-bound (decompile+extract+resolve): {:.1}% of wall -- the serial floor",
235        pct(kernel)
236    );
237}
238
239fn main() -> Result<(), Box<dyn std::error::Error>> {
240    let db = std::env::args()
241        .nth(1)
242        .expect("usage: taint <db.i64>  (set TAINT_LIMIT to cap the sweep)");
243
244    Ida::run(move |ida| -> Result<(), Error> { ida.call(move |idb| run(idb, &db))? })??;
245
246    Ok(())
247}
taint/taint.rs

taint/
taint.rs
