idakit 0.2.0

//! A non-trivial consumer of the materialized ctree, used as a stress test and a benchmark.
//!
//! For every function in the database it runs a crude intra-procedural taint pass
//! (call-return sources -> local def/use fixpoint -> dangerous-call sinks) and times the
//! work in four separate phases:
//!
//!   decompile   -- Hex-Rays, kernel thread, serial and unavoidable
//!   extract     -- `cfunc.ctree()`, the facade DFS + Rust rebuild, kernel thread
//!   resolve     -- turning `Obj(address)` callees into names, kernel thread (needs `&Database`)
//!   analyze     -- the pure taint pass over the Send image (no kernel access)
//!
//! The split answers the live question: how much of the work is the serial kernel
//! floor (decompile + extract + resolve) versus the part that could ever be fanned
//! out (analyze)? If `analyze` is a rounding error, parallelism buys nothing and we
//! stay sequential. The run also exercises every node kind, so it is where API
//! friction shows up.
//!
//! Run (release matters for the numbers):
//!   `cargo run -p idakit --release --example taint -- <db.i64>`
//! Cap the sweep with `TAINT_LIMIT=2000` while iterating.

use std::collections::{HashMap, HashSet};
use std::time::{Duration, Instant};

use idakit::decompiler::ctree::{Ctree, ExpressionId, LocalId, NodeRef};
use idakit::prelude::*;

/// Calls whose return value introduces taint (matched as a substring of the name).
const SOURCES: &[&str] = &["recv", "read", "fgets", "getenv", "scanf", "gets"];
/// Calls whose arguments are dangerous to feed tainted data into.
const SINKS: &[&str] = &[
    "memcpy", "memmove", "strcpy", "strcat", "sprintf", "system", "malloc", "alloca", "exec",
];

fn matches(name: &str, set: &[&str]) -> bool {
    set.iter().any(|n| name.contains(n))
}

/// A function lifted into a Send-able analysis input.
///
/// The callee-name map is the telling part, since the ctree's `Call` carries only an
/// `Obj(address)`/`Helper`, so the names every analysis actually keys on have to be resolved
/// *here*, on the kernel thread, and folded in. The bare tree cannot answer "what does this
/// call?" off-thread.
struct TaintInput {
    tree: Ctree,
    /// Call expression -> resolved callee name (only the ones that resolve to a symbol).
    callees: HashMap<ExpressionId, String>,
}

/// Resolves a call's callee to a name, if it is a direct symbol or a decompiler helper.
///
/// Indirect calls (through a variable or computed pointer) stay unresolved. The symbol name
/// rides on the `Obj` node, so this needs no kernel access.
fn callee_name(tree: &Ctree, callee: ExpressionId) -> Option<String> {
    let kind = tree.kind(callee);
    if let Some((_, name)) = kind.as_obj() {
        return name.map(str::to_owned);
    }
    kind.as_helper().map(str::to_owned)
}

/// Builds the callee-name map for a tree.
///
/// The tree carries callee names directly, so this is pure: no `Database` access, no
/// kernel-thread name resolution.
fn resolve_callees(tree: &Ctree) -> HashMap<ExpressionId, String> {
    let mut map = HashMap::new();
    for (id, callee, _) in tree.calls() {
        if let Some(name) = callee_name(tree, callee) {
            map.insert(id, name);
        }
    }
    map
}

/// Whether `e`'s subtree reads a tainted local or calls a source directly.
///
/// Flow-insensitive and deliberately crude, doing real work proportional to tree size.
fn expression_tainted(img: &TaintInput, e: ExpressionId, tainted: &HashSet<u32>) -> bool {
    img.tree
        .expression_descendants(NodeRef::Expression(e))
        .any(|id| match img.tree.kind(id).as_var() {
            Some(LocalId(i)) => tainted.contains(&i),
            None => img.callees.get(&id).is_some_and(|n| matches(n, SOURCES)),
        })
}

/// The pure phase returns the number of source->sink flows found. No `Database` access,
/// so this is exactly the work that could move to a worker thread.
fn analyze(img: &TaintInput) -> usize {
    // Collect `Var(i) = rhs` definitions once.
    let defs: Vec<(u32, ExpressionId)> = img
        .tree
        .assigns()
        .filter_map(|(_, _, x, y)| {
            let LocalId(i) = img.tree.kind(x).as_var()?;
            Some((i, y))
        })
        .collect();

    // Fixpoint: a local is tainted once any of its defining RHS is tainted.
    let mut tainted: HashSet<u32> = HashSet::new();
    loop {
        let before = tainted.len();
        for &(lv, rhs) in &defs {
            if !tainted.contains(&lv) && expression_tainted(img, rhs, &tainted) {
                tainted.insert(lv);
            }
        }
        if tainted.len() == before {
            break;
        }
    }

    // Sinks: a tainted argument to a dangerous call is a flow.
    img.tree
        .calls()
        .filter(|(id, _, args)| {
            img.callees.get(id).is_some_and(|n| matches(n, SINKS))
                && args.iter().any(|a| expression_tainted(img, *a, &tainted))
        })
        .count()
}

#[derive(Default)]
struct Totals {
    decompile: Duration,
    extract: Duration,
    resolve: Duration,
    analyze: Duration,
    funcs: usize,
    decompile_failed: usize,
    extract_failed: usize,
    flows: usize,
    nodes: u64,
}

fn run(idb: &mut Database, db: &str) -> Result<(), Error> {
    idb.open(db).call()?;

    let limit = std::env::var("TAINT_LIMIT")
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(usize::MAX);

    let eas: Vec<Address> = idb.functions().map(|f| f.address()).take(limit).collect();
    println!("[taint] sweeping {} functions", eas.len());

    let mut t = Totals::default();
    let wall = Instant::now();

    for (i, &address) in eas.iter().enumerate() {
        let started = Instant::now();
        let Ok(cf) = idb.decompile(address) else {
            t.decompile_failed += 1;
            continue;
        };
        t.decompile += started.elapsed();

        let started = Instant::now();
        let Ok(tree) = cf.ctree() else {
            t.extract_failed += 1;
            continue;
        };
        t.extract += started.elapsed();

        let started = Instant::now();
        let callees = resolve_callees(&tree);
        t.resolve += started.elapsed();

        t.nodes += (tree.expressions().count() + tree.statements().count()) as u64;
        let img = TaintInput { tree, callees };

        let started = Instant::now();
        t.flows += analyze(&img);
        t.analyze += started.elapsed();

        t.funcs += 1;
        if (i + 1) % 5000 == 0 {
            println!("[taint] {} / {} ...", i + 1, eas.len());
        }
    }

    report(&t, wall.elapsed());
    idb.close(false);
    Ok(())
}

fn report(t: &Totals, wall: Duration) {
    let kernel = t.decompile + t.extract + t.resolve;
    let pct = |d: Duration| 100.0 * d.as_secs_f64() / wall.as_secs_f64().max(f64::EPSILON);
    let per = |d: Duration| {
        if t.funcs == 0 {
            0.0
        } else {
            d.as_secs_f64() * 1e6 / t.funcs as f64
        }
    };

    println!("\n=== taint sweep ===");
    println!(
        "functions analyzed: {}  (decompile-failed {}, extract-failed {})",
        t.funcs, t.decompile_failed, t.extract_failed
    );
    println!("ctree nodes total:  {}", t.nodes);
    println!("source->sink flows: {}", t.flows);
    println!("wall:               {:.2}s", wall.as_secs_f64());
    println!(
        "  decompile  {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
        t.decompile.as_secs_f64(),
        pct(t.decompile),
        per(t.decompile)
    );
    println!(
        "  extract    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
        t.extract.as_secs_f64(),
        pct(t.extract),
        per(t.extract)
    );
    println!(
        "  resolve    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn",
        t.resolve.as_secs_f64(),
        pct(t.resolve),
        per(t.resolve)
    );
    println!(
        "  analyze    {:>7.2}s  {:>5.1}%   {:>7.1} us/fn   (the only parallelizable part)",
        t.analyze.as_secs_f64(),
        pct(t.analyze),
        per(t.analyze)
    );
    println!(
        "kernel-bound (decompile+extract+resolve): {:.1}% of wall -- the serial floor",
        pct(kernel)
    );
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let db = std::env::args()
        .nth(1)
        .expect("usage: taint <db.i64>  (set TAINT_LIMIT to cap the sweep)");

    Ida::run(move |ida| -> Result<(), Error> { ida.call(move |idb| run(idb, &db))? })??;

    Ok(())
}