diff options
| author | user <user@clank> | 2026-08-07 15:14:30 +0200 |
|---|---|---|
| committer | user <user@clank> | 2026-08-07 15:14:30 +0200 |
| commit | 72fce7da1a1fd6527e389ffeb0f951157523589a (patch) | |
| tree | 3f08a83f0d99f3d3fc7069b7be00d235bab82817 /tests/test_worker_client.py | |
| parent | Stop tracking 157MB of core dumps, and ignore them (diff) | |
| parent | docs: upstream findings for the ida-codemode maintainers (diff) | |
| download | ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.tar.gz ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.tar.xz ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.zip | |
Merge the IDA Code Mode port
Replaces the private idalib worker (idatui/worker.py + worker_client.py, with
server/patch_server.py injecting tools into ida-pro-mcp) with an ordinary
client of ida_codemode.client.DatabaseHandle. A database open by an IDA GUI is
reused; otherwise Code Mode starts or shares a managed idalib worker. The TUI
no longer owns an IDA process, and closing it releases only its lease.
Based on Duncan Ogilvie's port, rebased onto ~150 commits of local work it
predated. The rebase itself was mechanical; landing it was not. Nine defects
had to be fixed before the feature set was whole again, none of which the
patch's own tests could catch:
- DatabaseHandle.open() takes image_base, not loading_address: every
connect() would have raised TypeError on the first call
- five operations our tree had grown were simply missing (flowchart, so the
graph view was dead; op_format/pc_nums/pc_num_format, so 'o'/'O' were;
survey_binary)
- set_comments wrote only the disassembly comment, so comments never
appeared in pseudocode
- xref_query returned rows in raw IDA order, and 'follow the call' silently
followed the fall-through instead
- rename accepted one edit per category, so bulk symbol import was dead
- decompile ran decomp_map's full per-column ctree sweep to fill in a
per-line address anchor
- heads shipped without operand extents or the digest protocol
- the package became unimportable without ida_codemode installed, which
killed the offline test suites
Verified against the pre-codemode tag rather than against assumptions: the
full suite is 788 passed / 0 failed, and the pilot's 301 checks match the old
backend exactly. Performance is within 2x on the listing hot path and faster
on decompile, disasm and connect, after fixing two runtime costs that are
documented for upstream in docs/CODEMODE_UPSTREAM.md.
Test runtime came down from ~9m20s to 115s along the way -- not by removing
checks, but by removing four kinds of waiting-on-a-guess that were also
hiding real failures.
Diffstat (limited to 'tests/test_worker_client.py')
| -rw-r--r-- | tests/test_worker_client.py | 447 |
1 files changed, 0 insertions, 447 deletions
diff --git a/tests/test_worker_client.py b/tests/test_worker_client.py deleted file mode 100644 index cd99377..0000000 --- a/tests/test_worker_client.py +++ /dev/null @@ -1,447 +0,0 @@ -#!/usr/bin/env python3 -"""WorkerClient: spawn, transport, failure reporting, shutdown. - -This is the layer between the app and idalib, and it had no tests -- which is -awkward, because it is where the failures are silent and expensive. A worker -that dies during startup, a socket that drops mid-call, two UI threads sharing -one socket: none of those look like a bug from the outside, they look like the -TUI hanging or showing stale data. - -None of it needs IDA. The client spawns whatever ``_WORKER_PY`` points at, so -these tests point it at a fake that speaks the same length-prefixed pickle -protocol and can be told to misbehave on demand. -""" -from __future__ import annotations - -import os -import sys -import threading -import time - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -#: pure: a fake worker over a unix socket, no idalib anywhere. -#: Read by tests/run.py (--fast skips every NEEDS_IDA file). -NEEDS_IDA = False - -from idatui import worker_client as wc # noqa: E402 -from idatui.errors import IDAConnectionError, IDAToolError # noqa: E402 - -PASS = FAIL = 0 - - -def check(name, ok, detail=""): - global PASS, FAIL - if ok: - PASS += 1 - print(f" ok {name}") - else: - FAIL += 1 - print(f" FAIL {name} {detail}") - - -# --------------------------------------------------------------------------- # -# A worker that isn't IDA -# --------------------------------------------------------------------------- # -#: Speaks the real protocol (idatui.worker.send/recv) and implements a handful -#: of tools whose only job is to be predictable, plus the misbehaviours we need: -#: dying at startup, dropping the socket mid-conversation, taking its time. -FAKE_WORKER = r''' -import os, socket, sys, time -sys.path.insert(0, %(repo)r) -from idatui.worker import send, recv - -sock_path, binary = sys.argv[1], sys.argv[2] -mode = os.environ.get("FAKE_MODE", "ok") - -if mode == "die": - # A startup crash, the way the real worker reports one. - print("IDA Pro: thank you for using it") # banner noise, must be skipped - print("WORKER-FATAL: could not open database: it is wedged") - sys.stdout.flush() - sys.exit(3) -if mode == "hang": - time.sleep(60) # never binds: connect() must time out - sys.exit(0) - -srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) -if os.path.exists(sock_path): - os.unlink(sock_path) -srv.bind(sock_path) -srv.listen(1) -conn, _ = srv.accept() -served = 0 -while True: - msg = recv(conn) - if msg is None: - break - tool, args = msg - if tool == "__shutdown__": - # The real worker closes its database here; a clean exit is the signal - # the client waits for rather than killing us. - open(sock_path + ".clean", "w").write("shutdown") - break - served += 1 - if tool == "drop": - conn.close() # vanish mid-conversation - break - if tool == "boom": - send(conn, (False, "the tool exploded")) - continue - if tool == "slow": - time.sleep(float(args.get("secs", 0.2))) - send(conn, (True, {"tool": tool, "args": args, "n": served})) - continue - send(conn, (True, {"tool": tool, "args": args, "n": served, - "binary": os.path.basename(binary)})) -sys.exit(0) -''' - - -def _install_fake(tmpdir: str) -> str: - repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - path = os.path.join(tmpdir, "fake_worker.py") - with open(path, "w", encoding="utf-8") as fh: - fh.write(FAKE_WORKER % {"repo": repo}) - wc._WORKER_PY = path - return path - - -def client(tmpdir, **kw): - """A client wired to the fake worker, running under THIS interpreter. - - ``python=`` matters: the real constructor probes three interpreters for - ``import ida_pro_mcp`` and that is both slow and beside the point here. - """ - binary = os.path.join(tmpdir, "target.bin") - if not os.path.exists(binary): - with open(binary, "wb") as fh: - fh.write(b"\x7fELF" + b"\0" * 60) - return wc.WorkerClient(binary, python=sys.executable, **kw) - - -# --------------------------------------------------------------------------- # -def t_roundtrip(tmp): - c = client(tmp) - try: - c.connect(timeout=30) - r = c.call("survey_binary", depth=2) - check("a call round-trips through the socket", - r["tool"] == "survey_binary" and r["args"] == {"depth": 2}, str(r)) - check("the worker got the binary path we asked for", - r["binary"] == "target.bin", str(r)) - check("pid is exposed for memory accounting", isinstance(c.pid, int)) - r2 = c.call("second") - check("the connection is reused, not respawned per call", - r2["n"] == 2, f"n={r2['n']}") - finally: - c.close(grace=5) - - -def t_envelope(tmp): - """domain.decompile reads result.structuredContent -- keep that shape.""" - c = client(tmp) - try: - c.connect(timeout=30) - env = c.call_envelope("decompile", addr="0x1000") - inner = env["result"]["structuredContent"] - check("call_envelope wraps the payload the way domain.py unwraps it", - inner["tool"] == "decompile" and inner["args"] == {"addr": "0x1000"}, - str(env)) - finally: - c.close(grace=5) - - -def t_tool_error(tmp): - c = client(tmp) - try: - c.connect(timeout=30) - try: - c.call("boom") - check("a failing tool raises IDAToolError", False, "no exception") - except IDAToolError as e: - check("a failing tool raises IDAToolError", True) - check("the error names the tool", e.tool == "boom", f"tool={e.tool!r}") - check("and carries the worker's message", - "exploded" in e.message, e.message) - # A tool error is not a transport error: the connection must survive it, - # or one bad decompile would tear down the session. - r = c.call("after") - check("the connection survives a tool error", r["tool"] == "after", str(r)) - finally: - c.close(grace=5) - - -def t_dropped_socket(tmp): - """The reconnect trigger. The app catches IDAConnectionError and reconnects; - if a drop raised something else, or left _sock set, it would instead surface - as a crash or as every later call failing.""" - c = client(tmp) - try: - c.connect(timeout=30) - try: - c.call("drop") - check("a dropped socket raises IDAConnectionError", False, - "no exception") - except IDAConnectionError: - check("a dropped socket raises IDAConnectionError", True) - except Exception as e: # noqa: BLE001 - check("a dropped socket raises IDAConnectionError", False, - f"got {type(e).__name__}: {e}") - check("the dead socket is cleared, so a retry can reconnect", - c._sock is None) - finally: - c.close(grace=5) - - -def t_startup_crash(tmp): - """A worker that dies before binding must say WHY. - - 'worker exited (code 3)' on its own is indistinguishable from a bug in the - app; the real cause is the last meaningful line of its log, and the licence - banner must not be mistaken for it. - """ - os.environ["FAKE_MODE"] = "die" - try: - c = client(tmp) - try: - c.connect(timeout=30) - check("a worker that exits during startup raises", False, - "connect() returned") - except IDAConnectionError as e: - msg = str(e) - check("a worker that exits during startup raises", True) - check("the exit code is reported", "code 3" in msg, msg) - check("the WORKER-FATAL line is surfaced", - "wedged" in msg, msg) - check("the licence banner is not mistaken for the error", - "thank you" not in msg.lower(), msg) - check("the full log path is offered", ".log" in msg, msg) - finally: - os.environ.pop("FAKE_MODE", None) - - -def t_connect_timeout(tmp): - """A worker that never binds must give up, not block the UI forever.""" - os.environ["FAKE_MODE"] = "hang" - try: - c = client(tmp) - t0 = time.time() - try: - c.connect(timeout=0.4) - check("connect() gives up on a worker that never binds", False, - "returned") - except IDAConnectionError as e: - took = time.time() - t0 - check("connect() gives up on a worker that never binds", True) - check("it honours the timeout it was given", took < 10, f"{took:.1f}s") - check("and says so", "in time" in str(e), str(e)) - finally: - # grace=0: it is sleeping by design, don't wait it out. - c.close(grace=0) - finally: - os.environ.pop("FAKE_MODE", None) - - -def t_progress(tmp): - """connect() reports progress while analysis runs -- that callback is the - only thing on screen during a long open.""" - os.environ["FAKE_MODE"] = "hang" - seen = [] - try: - c = client(tmp) - try: - c.connect(timeout=0.6, progress=seen.append) - except IDAConnectionError: - pass - finally: - c.close(grace=0) - finally: - os.environ.pop("FAKE_MODE", None) - check("connect() reports progress while waiting", bool(seen), - f"{len(seen)} callbacks") - check("progress names the binary being analysed", - any("target.bin" in s for s in seen), str(seen[:1])) - - -def t_serialized(tmp): - """One socket, many UI threads. - - The app fires calls from several worker threads over one client. The frames - are length-prefixed pickle with no request ids, so if two calls interleaved - on the wire each would read the other's reply -- silently, as wrong data - rather than an error. The lock is the only thing preventing that, so this - checks every thread gets its own answer back. - """ - c = client(tmp) - try: - c.connect(timeout=30) - out, errs = {}, [] - - def go(i): - try: - out[i] = c.call("slow", secs=0.05, tag=i) - except Exception as e: # noqa: BLE001 - errs.append(e) - - threads = [threading.Thread(target=go, args=(i,)) for i in range(8)] - t0 = time.time() - for t in threads: - t.start() - for t in threads: - t.join(30) - took = time.time() - t0 - check("concurrent calls all completed", len(out) == 8 and not errs, - f"{len(out)} results, errors={errs[:1]}") - check("each thread got ITS OWN reply, not another's", - all(out[i]["args"]["tag"] == i for i in out), - str({i: out[i]["args"].get("tag") for i in sorted(out)})) - check("calls were serialized, not interleaved", - took >= 8 * 0.05, f"{took:.2f}s for 8 x 0.05s") - check("the worker saw every call exactly once", - sorted(r["n"] for r in out.values()) == list(range(1, 9)), - str(sorted(r["n"] for r in out.values()))) - finally: - c.close(grace=5) - - -def t_clean_shutdown(tmp): - """close() must let the worker close its database. - - A hard kill leaves the .i64 unpacked into .id0/.id1/... and the database - then fails to reopen. So close() sends __shutdown__ and WAITS; only a truly - stuck worker gets signalled. - """ - c = client(tmp) - c.connect(timeout=30) - sock_path = c._sock_path - proc = c._proc - c.close(grace=15) - check("close() sends __shutdown__ rather than killing", - os.path.exists(sock_path + ".clean")) - check("and waits for the worker to exit on its own", - proc.poll() == 0, f"returncode={proc.poll()}") - - -def t_call_after_close(tmp): - """A closed client must stay closed. - - call() reconnects when _sock is None, which is what makes a dropped socket - recoverable -- but it made an explicitly CLOSED client resurrect too, and - spawn a whole new idalib worker to serve one stray call. Teardown and - binary-switch both close while @work threads are in flight, so quitting - during a decompile left a fresh process re-opening the .i64 we had just - released. (Verified before the fix: pid 1066961 -> 1066962.) - """ - c = client(tmp) - c.connect(timeout=30) - pid = c.pid - c.close(grace=5) - try: - c.call("zombie") - check("a call after close() does not resurrect the worker", False, - f"call succeeded; pid {pid} -> {c.pid}") - except IDAConnectionError as e: - check("a call after close() does not resurrect the worker", True) - check("and says the client was closed", "closed" in str(e), str(e)) - check("no second worker was spawned", c.pid == pid, f"{pid} -> {c.pid}") - # ... but an explicit reconnect still revives it: that is how the app - # recovers from a worker that segfaulted. - c.connect(timeout=30) - r = c.call("revived") - check("connect() revives a closed client", r["tool"] == "revived", str(r)) - c.close(grace=5) - - -def t_worker_python_override(tmp): - """$IDATUI_WORKER_PYTHON wins, and the answer is cached. - - Without the override the constructor probes interpreters with a subprocess - each, which is why the override exists at all. - """ - wc._worker_python_cache = None - os.environ["IDATUI_WORKER_PYTHON"] = sys.executable - try: - got = wc._find_worker_python() - check("$IDATUI_WORKER_PYTHON is honoured", got == sys.executable, got) - finally: - os.environ.pop("IDATUI_WORKER_PYTHON", None) - wc._worker_python_cache = None - missing = "/nonexistent/python-that-is-not-there" - os.environ["IDATUI_WORKER_PYTHON"] = missing - try: - got = wc._find_worker_python() - check("an override that doesn't exist falls back instead of crashing", - got != missing and os.path.exists(got), got) - finally: - os.environ.pop("IDATUI_WORKER_PYTHON", None) - wc._worker_python_cache = None - - -def t_session_shims(tmp): - """The single-DB worker still has to answer the session questions the app - inherited from the old multi-session HTTP client.""" - c = client(tmp) - try: - c.connect(timeout=30) - sess = c.list_sessions() - check("list_sessions describes the one open database", - len(sess) == 1 and sess[0].filename == "target.bin" - and sess[0].is_active, str(sess)) - c.set_db("chosen") - check("set_db/resolve_db round-trip", c.resolve_db() == "chosen") - ka = c.keepalive() - ka.start() - ka.stop() - check("keepalive is a no-op the app can still drive", - ka.beats == 0 and ka.failures == 0) - h = c.health() - check("health answers even though the fake has no server_health tool", - isinstance(h, dict) and h, str(h)) - finally: - c.close(grace=5) - - -def t_log_tail(tmp): - """_log_tail picks the real error out of IDA's noise.""" - c = client(tmp) - with open(c._log_path, "w", encoding="utf-8") as fh: - fh.write("Thank you for using IDA\n" - "Licensed to: somebody\n" - "[MCP] registering tools\n" - "WORKER-FATAL: Failed to open database\n") - check("_log_tail surfaces WORKER-FATAL over the banner", - c._log_tail() == "Failed to open database", repr(c._log_tail())) - with open(c._log_path, "w", encoding="utf-8") as fh: - fh.write("Thank you for using IDA\nsomething odd happened\n") - tail = c._log_tail() - check("without a FATAL line it skips the banner and keeps the rest", - "odd happened" in tail and "Thank you" not in tail, repr(tail)) - os.unlink(c._log_path) - check("a missing log is reported, not raised", - "no worker log" in c._log_tail(), repr(c._log_tail())) - - -def main() -> int: - import tempfile - tests = [t_roundtrip, t_envelope, t_tool_error, t_dropped_socket, - t_startup_crash, t_connect_timeout, t_progress, t_serialized, - t_clean_shutdown, t_call_after_close, t_worker_python_override, - t_session_shims, t_log_tail] - with tempfile.TemporaryDirectory(prefix="idatui-wc-") as tmp: - _install_fake(tmp) - for fn in tests: - print(f"\n{fn.__name__}") - try: - fn(tmp) - except Exception as e: # noqa: BLE001 -- isolate one test's crash - import traceback - check(f"{fn.__name__} did not crash", False, - f"{type(e).__name__}: {e}") - traceback.print_exc() - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main()) |
