diff options
| author | blasty <blasty@local> | 2026-08-07 12:39:54 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-08-07 12:40:14 +0200 |
| commit | c9208de05d8583b677117fe43c9d3567e89eb2ce (patch) | |
| tree | 8f7b7487d9939be31b7c2b1a7932ee3a16c7403d /tests/test_worker_client.py | |
| parent | Stop tracking 157MB of core dumps, and ignore them (diff) | |
| download | ida-tui-c9208de05d8583b677117fe43c9d3567e89eb2ce.tar.gz ida-tui-c9208de05d8583b677117fe43c9d3567e89eb2ce.tar.xz ida-tui-c9208de05d8583b677117fe43c9d3567e89eb2ce.zip | |
Rebase MISTER EXO's ida-codemode port onto the current tree
Mechanical part of the port: the 27-file patch was cut against a base ~148
commits behind us, so it did not apply. Resolved 11 conflicts (all of them
diff drift, not semantic clashes) and the three file deletions:
- app.py: the patch re-inserted _do_rename/_do_name_addr/_seek_split etc. as
"theirs" because our tree moved them to edit_ctl.py/trace_ctl.py. Kept ours
and applied the real intent (WorkerClient->CodeModeClient, .call->.invoke,
_open_worker_client->_open_database_client) at their current homes.
- domain.py: kept Head as a NamedTuple -- the patch reverted it to a frozen
dataclass, which the perf work measured at 2.9us vs 1.9us per row on a
quarter-million-row walk. Dropped _fetch_output (no download_url under Code
Mode) and its now-dead urllib/json imports.
- pane.py: the patch's deletion swallowed our zellij support along with the
worker-reaping block it meant to remove. Kept zellij, removed the reaping.
- test_scenarios.py: the idb_save->save_database teardown hunk belongs to
tests/_fixtures.py in our tree; applied it there and kept our pc_num_format
scenario that the drift landed on.
Three defects in the patch itself, fixed here:
- It made "import idatui" hard-require ida_codemode, so every offline suite
died at import -- including the pure ones (graph/index/trace) that are the
house rule for "tests/run.py --fast". The import is now deferred and gated
on the binding, which is also what lets the port's own contract tests
inject a fake DatabaseHandle.
- project.stage() inlined an ida_codemode.registry import and treated "library
not installed" as "someone owns this database", which broke IDA-free project
staging. Ownership lookup moved to codemode_client.database_owner().
- tests/test_codemode_client.py had no NEEDS_IDA marker, which tests/run.py
rejects outright.
Offline suite: 301 passed, 0 failed. Against master's 344 the whole delta is
accounted for: -40 worker_client (module deleted), -18 launch sweep checks
(behaviour deliberately removed) +3 guarding that it stays removed, +2 pool
(GUI-save semantics), +13 new codemode_client contract tests.
NOT yet done, and the port is not functional without it: the adapter is
missing five operations our tree grew since the patch's base (flowchart,
op_format, pc_nums, pc_num_format, survey_binary) and its "heads" predates
back-walking and digest/expect.
Diffstat (limited to '')
| -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()) |
