#!/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())