diff options
| author | blasty <blasty@local> | 2026-07-09 11:45:45 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-07-09 11:45:45 +0200 |
| commit | 5bfcfa8f01d6f0e0d7535639eaa341fcce9fd9b8 (patch) | |
| tree | 9794c017a5d1ba9f112ff1f37c40cff737836b0a | |
| parent | stress paging on 10k-func binary; document scale constraints (diff) | |
| download | ida-tui-5bfcfa8f01d6f0e0d7535639eaa341fcce9fd9b8.tar.gz ida-tui-5bfcfa8f01d6f0e0d7535639eaa341fcce9fd9b8.tar.xz ida-tui-5bfcfa8f01d6f0e0d7535639eaa341fcce9fd9b8.zip | |
domain/paging layer: FunctionIndex, block-cached DisasmModel, decompile, resolve
- Program: model registry + small prefetch pool; cache invalidation hooks
- FunctionIndex: lazy pagination (advance by len, clamp page=500), filter globs,
viewport windows; full 10,092-func enumerate matches survey in ~3.1s
- DisasmModel: block-cached windowed disasm (256/block), forward+back prefetch,
cached instruction totals; deep window revisit 196ms -> 0ms (cache proven)
- Decompilation: truncation-aware, hard-failure surfaced as data (monster funcs)
- resolve(): int/hex/symbol via lookup_funcs (string-array query shape)
- tests/test_domain.py: 17/17 on both ls (~290 funcs) and libcrypto (10k+52k-insn)
- smoke_client no longer assumes 'main' exists (works on libraries)
- doc: idle_ttl worker self-exit finding ("not reachable" needs re-open)
| -rw-r--r-- | docs/PAGING_FINDINGS.md | 10 | ||||
| -rw-r--r-- | idatui/__init__.py | 20 | ||||
| -rw-r--r-- | idatui/domain.py | 384 | ||||
| -rw-r--r-- | tests/smoke_client.py | 16 | ||||
| -rw-r--r-- | tests/test_domain.py | 155 |
5 files changed, 582 insertions, 3 deletions
diff --git a/docs/PAGING_FINDINGS.md b/docs/PAGING_FINDINGS.md index 1c7abff..33cb2e6 100644 --- a/docs/PAGING_FINDINGS.md +++ b/docs/PAGING_FINDINGS.md @@ -93,6 +93,16 @@ fall back to the disassembly view or show an error panel. Normal decompile bodies are server-truncated with a `[N chars total]` marker (still to be solved for full-body display — see Phase 2). +## Workers self-exit when idle (`idle_ttl_sec`, default 600s) + +A headless idalib worker self-exits after ~10 min idle. Its session then lingers +in `idb_list` but calls fail with `"Worker for session <id> is not reachable"` +(distinct from `"Session not found"`). Recovery differs: this needs a fresh +`idb_open` of the same `input_path` (new worker), not just a db re-resolve. The +TUI should (a) raise `idle_ttl_sec` when opening, and/or (b) detect +"not reachable" and transparently re-open. Currently surfaced as a clean +`IDAToolError`; auto-reopen is a TODO for the app layer. + ## Writable path requirement (operational) `idb_open` writes the `.i64` next to the input binary, so the path must be diff --git a/idatui/__init__.py b/idatui/__init__.py index 450d64f..8dd371c 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -11,8 +11,28 @@ from .client import ( IDASessionError, Session, ) +from .domain import ( + Program, + FunctionIndex, + DisasmModel, + Func, + Line, + Ref, + Decompilation, + LIST_PAGE, + DISASM_BLOCK, +) __all__ = [ + "Program", + "FunctionIndex", + "DisasmModel", + "Func", + "Line", + "Ref", + "Decompilation", + "LIST_PAGE", + "DISASM_BLOCK", "IDAClient", "IDAError", "IDAConnectionError", diff --git a/idatui/domain.py b/idatui/domain.py new file mode 100644 index 0000000..d19a8a3 --- /dev/null +++ b/idatui/domain.py @@ -0,0 +1,384 @@ +"""Domain / paging layer: address-centric models over the raw MCP client. + +This is where the "millions of lines" problem is solved, so the TUI widgets only +ever see a viewport-sized slice. Every hard-won constraint from +``docs/PAGING_FINDINGS.md`` is encoded here: + +* Per-call caps are silent (over the cap the server returns 10, not a clamp), so + we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. +* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. +* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly + is **block-cached** (revisits are free) and **prefetches** the next block on a + background thread (the client is concurrency-safe). +* ``include_total`` scans the whole function (~200ms on monsters); totals are + fetched once and cached. +* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is + null); that is surfaced as data, not an exception. + +Everything here is synchronous and thread-safe. The TUI runs these calls from +Textual worker threads; the internal prefetch pool is separate and small. +""" + +from __future__ import annotations + +import re +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Callable, Iterable + +from .client import IDAClient, IDAToolError + +# Clamps derived from measured caps (list ~700, disasm ~500). Margin included. +LIST_PAGE = 500 +DISASM_BLOCK = 256 # instructions per cached/fetched block (<= disasm cap) + +_TRUNC_RE = re.compile(r"\[(\d+) chars total\]\s*$") + + +# --------------------------------------------------------------------------- # +# Value models +# --------------------------------------------------------------------------- # +def _as_int(v) -> int: + if isinstance(v, int): + return v + return int(v, 16) if isinstance(v, str) and v.startswith("0x") else int(v, 16) + + +@dataclass(frozen=True) +class Func: + addr: int + name: str + size: int + + @classmethod + def from_raw(cls, d: dict) -> "Func": + return cls(addr=_as_int(d["addr"]), name=d["name"], size=_as_int(d["size"])) + + +@dataclass(frozen=True) +class Line: + """One rendered disassembly line. ``ea`` is the instruction address.""" + + ea: int + text: str + label: str | None = None + + @classmethod + def from_raw(cls, d: dict) -> "Line": + return cls( + ea=_as_int(d["addr"]), + text=d.get("instruction", ""), + label=d.get("label"), + ) + + +@dataclass +class Ref: + addr: int + name: str + string: str | None = None + + +@dataclass +class Decompilation: + ea: int + code: str | None + failed: bool + error: str | None + truncated: bool + total_chars: int | None + refs: list[Ref] = field(default_factory=list) + + +# --------------------------------------------------------------------------- # +# Query-payload unwrapping (shape: {"result":[{"data":[...],"next_offset":N}]}) +# --------------------------------------------------------------------------- # +def _query_data(payload) -> list: + res = payload.get("result", payload) if isinstance(payload, dict) else payload + if isinstance(res, list): + res = res[0] if res else {} + return res.get("data", []) if isinstance(res, dict) else [] + + +# --------------------------------------------------------------------------- # +# Function index: lazy, clamped, cached, filterable +# --------------------------------------------------------------------------- # +class FunctionIndex: + """A lazily-paginated, cached view of the function list. + + Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by + ``next_offset``). A single index instance corresponds to one server-side + ``filter`` glob (``None`` = all functions). + """ + + def __init__(self, program: "Program", filter: str | None = None): + self._prog = program + self.filter = filter + self._funcs: list[Func] = [] + self._by_addr: dict[int, Func] = {} + self._done = False + self._lock = threading.Lock() + + def _load_next_page(self) -> int: + with self._lock: + if self._done: + return 0 + offset = len(self._funcs) + query: dict = {"offset": offset, "count": LIST_PAGE} + if self.filter: + query["filter"] = self.filter + data = _query_data(self._prog.client.call("list_funcs", queries=[query])) + added = 0 + with self._lock: + for d in data: + f = Func.from_raw(d) + if f.addr not in self._by_addr: + self._by_addr[f.addr] = f + self._funcs.append(f) + added += 1 + if len(data) < LIST_PAGE: + self._done = True + return len(data) + + def ensure(self, n: int) -> None: + """Ensure at least ``n`` functions are loaded (or all, if fewer exist).""" + while not self._done and len(self._funcs) < n: + if self._load_next_page() == 0: + break + + def load_all(self, progress: Callable[[int], None] | None = None) -> None: + """Fully enumerate (background-friendly). ~3s for 10k funcs.""" + while not self._done: + self._load_next_page() + if progress: + progress(len(self._funcs)) + + @property + def complete(self) -> bool: + return self._done + + def __len__(self) -> int: + return len(self._funcs) + + def loaded(self) -> int: + return len(self._funcs) + + def get(self, i: int) -> Func | None: + self.ensure(i + 1) + with self._lock: + return self._funcs[i] if 0 <= i < len(self._funcs) else None + + def window(self, start: int, count: int) -> list[Func]: + """Viewport slice ``[start, start+count)`` (loads as needed).""" + self.ensure(start + count) + with self._lock: + return list(self._funcs[start : start + count]) + + def by_addr(self, ea: int) -> Func | None: + with self._lock: + return self._by_addr.get(ea) + + +# --------------------------------------------------------------------------- # +# Disassembly model: block-cached windowed listing for ONE function +# --------------------------------------------------------------------------- # +class DisasmModel: + """Windowed, block-cached disassembly of a single function. + + Because ``disasm offset=N`` is O(N) and uncacheable server-side, we fetch and + cache fixed ``DISASM_BLOCK``-sized blocks; a viewport read slices from cached + blocks and returns instantly on revisit. The block just past the viewport is + prefetched on a background thread. + """ + + BLOCK = DISASM_BLOCK + + def __init__(self, program: "Program", ea: int, name: str | None = None): + self._prog = program + self.ea = ea + self.name = name + self._blocks: dict[int, list[Line]] = {} + self._total: int | None = None + self._lock = threading.Lock() + self._inflight: set[int] = set() + + def total(self) -> int: + """Total instruction count (fetched once; ~200ms on huge funcs).""" + if self._total is not None: + return self._total + payload = self._prog.client.call( + "disasm", addr=hex(self.ea), max_instructions=1, include_total=True + ) + total = payload.get("total_instructions") + if total is None: + total = payload.get("instruction_count", 0) + with self._lock: + self._total = int(total) + return self._total + + def _fetch_block(self, b: int) -> list[Line]: + payload = self._prog.client.call( + "disasm", addr=hex(self.ea), offset=b * self.BLOCK, + max_instructions=self.BLOCK, + ) + raw = payload.get("asm", {}).get("lines", []) if isinstance(payload, dict) else [] + lines = [Line.from_raw(r) for r in raw] + with self._lock: + self._blocks[b] = lines + self._inflight.discard(b) + return lines + + def _get_block(self, b: int) -> list[Line]: + with self._lock: + hit = self._blocks.get(b) + if hit is not None: + return hit + return self._fetch_block(b) + + def _prefetch_block(self, b: int) -> None: + if b < 0: + return + with self._lock: + if b in self._blocks or b in self._inflight: + return + self._inflight.add(b) + self._prog.submit(self._fetch_block, b) + + def lines(self, start: int, count: int, prefetch: bool = True) -> list[Line]: + """Return rendered lines for viewport ``[start, start+count)``.""" + if count <= 0 or start < 0: + return [] + end = start + count + b0, b1 = start // self.BLOCK, (end - 1) // self.BLOCK + out: list[Line] = [] + for b in range(b0, b1 + 1): + block = self._get_block(b) + lo = start - b * self.BLOCK if b == b0 else 0 + hi = end - b * self.BLOCK if b == b1 else self.BLOCK + out.extend(block[max(lo, 0):hi]) + if prefetch: + self._prefetch_block(b1 + 1) # forward scroll + self._prefetch_block(b0 - 1) # backward scroll + return out + + def cached_blocks(self) -> int: + with self._lock: + return len(self._blocks) + + def invalidate(self) -> None: + with self._lock: + self._blocks.clear() + self._total = None + + +# --------------------------------------------------------------------------- # +# Program: top-level handle, model registry, prefetch pool +# --------------------------------------------------------------------------- # +class Program: + """The bound analysis session: models, caches, and a small prefetch pool.""" + + def __init__(self, client: IDAClient, prefetch_workers: int = 2): + self.client = client + self._pool = ThreadPoolExecutor( + max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" + ) + self._indices: dict[str | None, FunctionIndex] = {} + self._disasm: dict[int, DisasmModel] = {} + self._decomp: dict[int, Decompilation] = {} + self._lock = threading.Lock() + + # -- prefetch plumbing ------------------------------------------------- # + def submit(self, fn, *args) -> None: + try: + self._pool.submit(fn, *args) + except RuntimeError: + pass # pool shut down + + def close(self) -> None: + self._pool.shutdown(wait=False, cancel_futures=True) + + # -- functions --------------------------------------------------------- # + def functions(self, filter: str | None = None) -> FunctionIndex: + with self._lock: + idx = self._indices.get(filter) + if idx is None: + idx = FunctionIndex(self, filter) + self._indices[filter] = idx + return idx + + # -- disassembly ------------------------------------------------------- # + def disasm(self, ea: int, name: str | None = None) -> DisasmModel: + with self._lock: + m = self._disasm.get(ea) + if m is None: + m = DisasmModel(self, ea, name) + self._disasm[ea] = m + return m + + # -- decompilation ----------------------------------------------------- # + def decompile(self, ea: int, refresh: bool = False) -> Decompilation: + if not refresh: + with self._lock: + hit = self._decomp.get(ea) + if hit is not None: + return hit + payload = self.client.call("decompile", addr=hex(ea)) + result = _parse_decompilation(ea, payload) + with self._lock: + self._decomp[ea] = result + return result + + # -- address resolution ------------------------------------------------ # + def resolve(self, target: int | str) -> int: + """Resolve an int/hex-string/symbol name to an address (ea).""" + if isinstance(target, int): + return target + s = target.strip() + if s.startswith("0x") or s.startswith("0X"): + return int(s, 16) + if re.fullmatch(r"[0-9a-fA-F]+", s): + return int(s, 16) + # Symbol name -> ask the server. lookup_funcs takes an array of strings + # and returns [{"query":..., "fn": {addr,name,size} | null, "error":...}]. + try: + payload = self.client.call("lookup_funcs", queries=[s]) + except IDAToolError as e: + raise KeyError(f"cannot resolve {target!r}: {e}") from e + res = payload.get("result", []) if isinstance(payload, dict) else [] + fn = res[0].get("fn") if res and isinstance(res[0], dict) else None + if not fn: + raise KeyError(f"cannot resolve {target!r}") + return _as_int(fn["addr"]) + + # -- invalidation (after edits) --------------------------------------- # + def invalidate(self, ea: int) -> None: + """Drop caches for a function after a rename/comment/patch/etc.""" + with self._lock: + self._decomp.pop(ea, None) + m = self._disasm.get(ea) + if m is not None: + m.invalidate() + + def invalidate_functions(self) -> None: + """Drop the function index caches (after rename/define/undefine).""" + with self._lock: + self._indices.clear() + + +def _parse_decompilation(ea: int, payload) -> Decompilation: + if not isinstance(payload, dict): + return Decompilation(ea, None, True, "unexpected payload", False, None) + code = payload.get("code") + error = payload.get("error") + if not code: + return Decompilation(ea, None, True, error or "decompilation failed", + False, None) + m = _TRUNC_RE.search(code) + truncated = m is not None + total_chars = int(m.group(1)) if m else len(code) + refs = [ + Ref(addr=_as_int(r["addr"]), name=r.get("name", ""), string=r.get("string")) + for r in payload.get("refs", []) if isinstance(r, dict) and "addr" in r + ] + return Decompilation(ea, code, False, error, truncated, total_chars, refs) diff --git a/tests/smoke_client.py b/tests/smoke_client.py index 647b589..5466284 100644 --- a/tests/smoke_client.py +++ b/tests/smoke_client.py @@ -26,6 +26,13 @@ from idatui.client import ( # noqa: E402 PASS, FAIL = 0, 0 +def _query_data(payload): + res = payload.get("result", payload) if isinstance(payload, dict) else payload + if isinstance(res, list): + res = res[0] if res else {} + return res.get("data", []) if isinstance(res, dict) else [] + + def check(name, cond, detail=""): global PASS, FAIL if cond: @@ -82,8 +89,11 @@ def main(argv): check("health is dict", isinstance(h, dict), type(h).__name__) funcs = ida.call("list_funcs", queries=[{"count": 5}]) check("list_funcs shape", isinstance(funcs, (list, dict)), type(funcs).__name__) - dis = ida.call("disasm", addr="main", max_instructions=10) - lines = dis.get("asm", {}).get("lines") if isinstance(dis, dict) else None + # Pick a real function (don't assume 'main' exists — libraries have none). + some = ida.call("list_funcs", queries=[{"filter": "sub_*", "count": 1}]) + target = _query_data(some)[0]["addr"] + dis = ida.call("disasm", addr=target, max_instructions=10) + lines = (dis.get("asm") or {}).get("lines") if isinstance(dis, dict) else None check("disasm main has lines", bool(lines), str(dis)[:120]) check("disasm line carries addr", bool(lines and "addr" in lines[0]), str(lines[0]) if lines else "no lines") @@ -136,7 +146,7 @@ def main(argv): print("[concurrency]") def worker(i): - return ida.call("disasm", addr="main", max_instructions=5) + return ida.call("disasm", addr=target, max_instructions=5) with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex: results = list(ex.map(worker, range(40))) diff --git a/tests/test_domain.py b/tests/test_domain.py new file mode 100644 index 0000000..6875b66 --- /dev/null +++ b/tests/test_domain.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Validate the domain/paging layer against a real, large binary. + +Run with a session open on a big binary (e.g. libcrypto.so.3, 10k funcs): + + python3 tests/test_domain.py --db <session_id> + +Checks: correct full pagination (advance by len, not next_offset), viewport +slicing across block boundaries, window caching (revisit is instant), prefetch +warming, cached instruction totals, decompile success + hard-failure handling, +and address resolution. +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from idatui.client import IDAClient # noqa: E402 +from idatui.domain import DISASM_BLOCK, LIST_PAGE, Program # noqa: E402 + +URL = "http://127.0.0.1:8745/mcp" +PASS = FAIL = 0 + + +def check(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f" ok {name}") + else: + FAIL += 1 + print(f" FAIL {name} {detail}") + + +def ms(fn): + t = time.time() + r = fn() + return (time.time() - t) * 1e3, r + + +def main(argv): + db = None + it = iter(argv) + for a in it: + if a == "--db": + db = next(it) + c = IDAClient(URL, db=db) + c.connect() + if db is None: + c.set_db(c.resolve_db()) + prog = Program(c) + module = c.health().get("module") + total_funcs = c.call("survey_binary").get("statistics", {}).get("total_functions") + print(f"db={c.db} module={module} survey_total_functions={total_funcs}") + + # ---- function index: full enumeration correctness ------------------ # + print("\n[function index]") + idx = prog.functions() + dt, _ = ms(lambda: idx.load_all()) + n = len(idx) + check("enumerated all functions (matches survey)", n == total_funcs, + f"got {n} vs survey {total_funcs}") + print(f" loaded {n} funcs in {dt:.0f}ms ({dt / max(n,1):.3f}ms/func), " + f"page={LIST_PAGE}") + # no duplicates, monotonic-ish uniqueness by addr + addrs = [idx.get(i).addr for i in range(min(n, 3000))] + check("no duplicate addrs in first 3000", len(addrs) == len(set(addrs))) + # viewport slice (adapt to binary size) + wstart = min(1000, max(n - 50, 0)) + wlen = min(50, n - wstart) + w = idx.window(wstart, wlen) + check(f"window({wstart},{wlen}) returns {wlen}", len(w) == wlen, str(len(w))) + + # ---- filtered index uses server-side glob -------------------------- # + print("\n[filtered index]") + sub = prog.functions(filter="sub_*") + sub.ensure(10) + check("filter sub_* yields sub_ names", all(f.name.startswith("sub_") for f in sub.window(0, 10)), + str([f.name for f in sub.window(0, 5)])) + + # ---- pick the fattest function for disasm stress ------------------- # + print("\n[disasm windowing]") + fattest = max((idx.get(i) for i in range(n)), key=lambda f: f.size) + dm = prog.disasm(fattest.addr, fattest.name) + dt, total = ms(dm.total) + check("total() returns positive count", total > 0, str(total)) + dt2, total2 = ms(dm.total) + check("total() cached (2nd call ~instant)", dt2 < dt / 2 + 1, f"{dt:.0f}ms -> {dt2:.1f}ms") + print(f" fattest {fattest.name}: {total} insns, total() {dt:.0f}ms then {dt2:.1f}ms") + + # viewport across a block boundary + start = DISASM_BLOCK - 5 + win = dm.lines(start, 60, prefetch=False) + check("viewport spans block boundary, right length", + len(win) == min(60, max(total - start, 0)), f"got {len(win)}") + # addresses strictly increasing and contiguous slice + eas = [ln.ea for ln in win] + check("viewport addrs strictly increasing", all(b > a for a, b in zip(eas, eas[1:])), + str(eas[:4])) + + # deep window: first slow (O(offset)), revisit instant (cached) + if total > DISASM_BLOCK * 4: + deep = (total // DISASM_BLOCK - 1) * DISASM_BLOCK + dt_cold, a = ms(lambda: dm.lines(deep, 60, prefetch=False)) + dt_warm, b = ms(lambda: dm.lines(deep, 60, prefetch=False)) + check("deep window revisit is cached/instant", dt_warm < dt_cold / 2 + 1, + f"cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms") + check("cached window identical", [l.ea for l in a] == [l.ea for l in b]) + print(f" deep@{deep}: cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms") + + # prefetch warms the next block + print("\n[prefetch]") + dm2 = prog.disasm(fattest.addr + 0) # same model (cached by ea) + fresh = prog.disasm(idx.get(0).addr, idx.get(0).name) + fresh.lines(0, 60, prefetch=True) # should prefetch block 1 + time.sleep(0.3) + check("prefetch warmed a neighbor block", fresh.cached_blocks() >= 2, + f"cached_blocks={fresh.cached_blocks()}") + + # ---- decompile: success + hard-failure ----------------------------- # + print("\n[decompile]") + # a small function likely decompiles + small = min((idx.get(i) for i in range(n)), key=lambda f: f.size if f.size > 4 else 1 << 30) + d_small = prog.decompile(small.addr) + check("small func decompiles or fails cleanly", isinstance(d_small.failed, bool)) + dt_c, _ = ms(lambda: prog.decompile(small.addr)) + check("decompile cached (2nd ~instant)", dt_c < 5, f"{dt_c:.1f}ms") + # the monster should hard-fail as a soft error (not raise) + d_big = prog.decompile(fattest.addr) + check("monster decompile handled (failed flag, no raise)", + d_big.failed or d_big.code is not None, + f"failed={d_big.failed} err={d_big.error}") + print(f" small {small.name}: failed={d_small.failed} " + f"trunc={d_small.truncated} chars={d_small.total_chars}") + print(f" monster {fattest.name}: failed={d_big.failed} err={d_big.error}") + + # ---- resolve ------------------------------------------------------- # + print("\n[resolve]") + check("resolve hex", prog.resolve(hex(fattest.addr)) == fattest.addr) + check("resolve int passthrough", prog.resolve(fattest.addr) == fattest.addr) + named = next((idx.get(i) for i in range(n) if not idx.get(i).name.startswith("sub_")), None) + if named: + try: + r = prog.resolve(named.name) + check("resolve symbol name", r == named.addr, f"{hex(r)} vs {hex(named.addr)} ({named.name})") + except KeyError as e: + check("resolve symbol name", False, str(e)) + + prog.close() + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) |
