1 files changed, 447 insertions, 0 deletions
diff --git a/tests/test_worker_client.py b/tests/test_worker_client.py
new file mode 100644
index 0000000..cd99377
--- /dev/null
+++ b/tests/test_worker_client.py
@@ -0,0 +1,447 @@
+#!/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())