diff options
Diffstat (limited to '')
| -rw-r--r-- | tests/_fixtures.py | 8 | ||||
| -rwxr-xr-x | tests/run.py | 4 | ||||
| -rw-r--r-- | tests/test_codemode_client.py | 140 | ||||
| -rw-r--r-- | tests/test_launch.py | 124 | ||||
| -rw-r--r-- | tests/test_pool.py | 61 | ||||
| -rw-r--r-- | tests/test_project.py | 2 | ||||
| -rw-r--r-- | tests/test_scenarios.py | 34 | ||||
| -rw-r--r-- | tests/test_worker_client.py | 447 |
8 files changed, 233 insertions, 587 deletions
diff --git a/tests/_fixtures.py b/tests/_fixtures.py index 766ecd5..c7a45d3 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -50,7 +50,13 @@ async def build_pristine(binary: str, cache: str, app_factory) -> None: await pilot.pause(0.05) if app._func_index is not None and app._func_index.complete: break - app.program.client.call("idb_save", timeout=600.0) + app.program.client.save_database() + # Textual's headless run_test context does not reliably emit App.Unmount on + # every platform/version; release the Code Mode lease explicitly. + if app.program is not None: + app.program.close() + if app.client is not None: + app.client.close() db = binary + ".i64" if os.path.exists(db): shutil.copy2(db, cache) diff --git a/tests/run.py b/tests/run.py index 4920e54..b38a707 100755 --- a/tests/run.py +++ b/tests/run.py @@ -51,8 +51,8 @@ import time ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) TESTS = os.path.join(ROOT, "tests") -#: The IDA-capable interpreter. The pilot tests need textual AND idapro in one -#: python; the worker python is auto-detected separately by WorkerClient. +#: The IDA-capable interpreter. The pilot tests need textual AND the Code Mode +#: library in one python; the database process is Code Mode's to place. DEFAULT_PY = os.path.expanduser("~/ida-venv/bin/python") #: Both shapes the suites print: "N passed, M failed" and "N checks, M failed". diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py new file mode 100644 index 0000000..0954918 --- /dev/null +++ b/tests/test_codemode_client.py @@ -0,0 +1,140 @@ +"""IDA-free contract tests for the Code Mode client adapter.""" +from __future__ import annotations + +import os +import sys +import tempfile +from dataclasses import dataclass + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import idatui.codemode_client as module # noqa: E402 +from idatui.codemode_client import CodeModeClient, _parse_load_args # noqa: E402 +from idatui.errors import IDAToolError # noqa: E402 + +#: Pure: fakes the DatabaseHandle, never touches IDA or the Code Mode library. +NEEDS_IDA = False + +PASS = FAIL = 0 + + +def check(name: str, condition: bool, detail="") -> None: + global PASS, FAIL + if condition: + PASS += 1 + print(f" ok {name}") + else: + FAIL += 1 + print(f" FAIL {name} {detail}") + + +@dataclass(frozen=True) +class FakeEntry: + pid: int = 123 + backend: str = "gui" + record_id: str = "123-abcdef" + exe_path: str = "" + idb_path: str = "" + + +class FakeHandle: + def __init__(self, path: str) -> None: + self.connected = True + self.entry = FakeEntry(exe_path=path, idb_path=path + ".i64") + self.waited = None + self.saved = 0 + self.closed = False + self.code = "" + self.code_timeout = None + + def wait_autoanalysis(self, timeout=None): + self.waited = timeout + return {"complete": True, "status": "complete"} + + def execute_python(self, code, timeout=None): + self.code = code + self.code_timeout = timeout + return {"result": {"sentinel": 7}, "stdout": "", "stderr": ""} + + def save_database(self): + self.saved += 1 + return {"saved": True, "idb_path": self.entry.idb_path} + + def close(self): + self.connected = False + self.closed = True + + +class FakeDatabaseHandle: + opened = None + kwargs = None + + @classmethod + def open(cls, path, **kwargs): + cls.opened = path + cls.kwargs = kwargs + return FakeHandle(path) + + +def main() -> int: + proc, base, file_type = _parse_load_args("-parm:ARMv7-M -b800000 -TRaw") + check("legacy switches map to typed Code Mode options", + (proc, base, file_type) == ("arm:ARMv7-M", 0x8000000, "Raw"), + (proc, base, file_type)) + try: + _parse_load_args("-parm -zcustom") + except ValueError as exc: + check("arbitrary IDA switches fail loudly", "cannot represent" in str(exc), exc) + else: + check("arbitrary IDA switches fail loudly", False) + + original = module.DatabaseHandle + module.DatabaseHandle = FakeDatabaseHandle + try: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "sample.bin") + with open(path, "wb") as file: + file.write(b"sample") + client = CodeModeClient(path, load_args="-parm:ARMv7-A -b100") + notes = [] + client.connect(timeout=42, progress=notes.append) + handle = client._handle + check("connect delegates database discovery to DatabaseHandle.open", + FakeDatabaseHandle.opened == path and handle is not None) + check("typed loader options cross the dependency boundary", + FakeDatabaseHandle.kwargs["processor"] == "arm:ARMv7-A" + and FakeDatabaseHandle.kwargs["loading_address"] == 0x1000, + FakeDatabaseHandle.kwargs) + check("connect waits for Code Mode autoanalysis", + handle.waited == 42, getattr(handle, "waited", None)) + check("progress distinguishes discovery and backend attachment", + len(notes) == 2 and "gui" in notes[-1], notes) + result = client.invoke("list_funcs", queries=[{"offset": 0, "count": 2}]) + check("invoke returns execute_python's result", result == {"sentinel": 7}, result) + check("operation scripts use the preloaded ida-domain database", + "db.functions.get_all()" in handle.code, handle.code[:200]) + check("health exposes registry identity", + client.health()["record_id"] == "123-abcdef") + client.save_database() + check("save uses the public Code Mode save route", handle.saved == 1) + client.close() + check("close releases only the handle lease", handle.closed) + check("GUI lifetime is never claimed by the client", + client.wait_released(0) is False) + finally: + module.DatabaseHandle = original + + client = CodeModeClient(__file__) + try: + client.invoke("not-an-operation") + except IDAToolError as exc: + check("unknown adapter operations are explicit", exc.tool == "not-an-operation") + else: + check("unknown adapter operations are explicit", False) + + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_launch.py b/tests/test_launch.py index c6ed8f1..d57ee5f 100644 --- a/tests/test_launch.py +++ b/tests/test_launch.py @@ -1,12 +1,14 @@ #!/usr/bin/env python3 -"""The launcher's file handling -- the part that deletes things. +"""The launcher's option handling, and the file handling it must NOT do. -`_sweep_locks` runs automatically when a database fails to open, and it removes -files next to the user's binary. That is exactly the kind of code that must not -be tested by trying it, so it is tested here: which files it takes, which it -must never take, and what it reports. +The old `_sweep_locks` deleted `.id0/.id1/.id2/.nam/.til` next to the user's +binary when a database failed to open. That was only defensible while the TUI +exclusively owned a private worker; under Code Mode a GUI or another client may +own the database, so the sweep is gone. Its tests are replaced by one that keeps +it gone -- deleting a shared database's working files is unrecoverable, and this +is the cheapest guard against someone reintroducing the "helpful" cleanup. -Pure: no IDA, no worker, no Textual. +Pure: no IDA, no Code Mode library, no Textual. """ from __future__ import annotations @@ -20,7 +22,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) #: Read by tests/run.py (--fast skips every NEEDS_IDA file). NEEDS_IDA = False -from idatui.launch import _LOCK_SUFFIXES, _load_args, _sweep_locks # noqa: E402 +import idatui.launch as launch # noqa: E402 +from idatui.launch import _load_args # noqa: E402 PASS = FAIL = 0 @@ -41,100 +44,18 @@ def touch(*paths): fh.write(b"x") -def t_sweeps_the_scratch_files(): - """IDA unpacks a .i64 into .id0/.id1/.id2/.nam/.til while it is open; a - hard-killed worker leaves them and the .i64 then refuses to reopen.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - touch(binary, *[binary + s for s in _LOCK_SUFFIXES]) - n = _sweep_locks(binary) - check("every unpacked scratch file is swept", n == len(_LOCK_SUFFIXES), - f"swept {n} of {len(_LOCK_SUFFIXES)}") - check("none of them survive", - not any(os.path.exists(binary + s) for s in _LOCK_SUFFIXES)) - check("the binary itself is untouched", os.path.exists(binary)) +def t_no_lock_sweeping(): + """The launcher must not delete database working files any more. - -def t_sweeps_by_stem_too(): - """IDA keys the scratch on the full name or the stem depending on how the - database was created, so both are swept.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "prog.elf") - stem = os.path.join(d, "prog") - touch(binary, stem + ".id0", stem + ".nam", binary + ".id1") - n = _sweep_locks(binary) - check("scratch named after the stem is swept too", n == 3, f"n={n}") - check("stem-keyed files are gone", - not os.path.exists(stem + ".id0") - and not os.path.exists(stem + ".nam")) - check("full-name-keyed files are gone", not os.path.exists(binary + ".id1")) - check("the binary itself is untouched", os.path.exists(binary)) - - -def t_never_the_database(): - """The .i64 IS the database. Nothing is saved unless idb_save was called, so - deleting it throws away every rename and comment in the session.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - db = binary + ".i64" - stem_db = os.path.join(d, "echo.i64") - touch(binary, db, binary + ".id0") - _sweep_locks(binary) - check("the .i64 is never swept", os.path.exists(db)) - check("nor the stem-keyed .i64", os.path.exists(stem_db)) - check(".i64 is not in the suffix list", ".i64" not in _LOCK_SUFFIXES, - str(_LOCK_SUFFIXES)) - - -def t_never_the_input_itself(): - """`.til` is both an unpacked-DB suffix and the extension of an IDA type - library, so `ida-tui mylib.til` used to sweep its own argument out of - existence -- irreversibly, on a path that runs automatically when an open - fails. Same for anything named *.id0/*.id1/*.id2/*.nam. + Code Mode's registry locks, health probes and IDA itself arbitrate database + ownership now. A sweep here would delete files out from under a live GUI. """ - for suf in _LOCK_SUFFIXES: - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "mylib" + suf) - touch(binary) - _sweep_locks(binary) - check(f"a binary named *{suf} is not deleted by its own sweep", - os.path.exists(binary), f"{binary} was removed") - - -def t_relative_path_is_still_the_input(): - """The guard compares absolute paths -- a relative argument names the same - file and must be protected the same way.""" - with tempfile.TemporaryDirectory() as d: - cwd = os.getcwd() - try: - os.chdir(d) - touch("mylib.til") - _sweep_locks("mylib.til") - check("a relative path to the input is protected too", - os.path.exists("mylib.til")) - finally: - os.chdir(cwd) - - -def t_missing_files_are_fine(): - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "nothing-here") - touch(binary) - n = _sweep_locks(binary) - check("sweeping with nothing to sweep reports 0", n == 0, f"n={n}") - check("and does not raise", True) - - -def t_leaves_the_neighbours_alone(): - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - other = os.path.join(d, "other.id0") # another binary's scratch - src = os.path.join(d, "echo.c") - touch(binary, other, src, binary + ".id0") - _sweep_locks(binary) - check("another binary's scratch is left alone", os.path.exists(other)) - check("unrelated neighbours are left alone", os.path.exists(src)) - check("our own scratch is still swept", not os.path.exists(binary + ".id0")) + check("_sweep_locks is gone", not hasattr(launch, "_sweep_locks")) + check("the scratch-suffix list is gone", not hasattr(launch, "_LOCK_SUFFIXES")) + src = open(launch.__file__, encoding="utf-8").read() + check("the launcher does not remove files at all", + "os.remove" not in src and "shutil.rmtree" not in src, + "launch.py deletes something again") def t_load_args(): @@ -154,10 +75,7 @@ def t_load_args(): def main() -> int: - for fn in (t_sweeps_the_scratch_files, t_sweeps_by_stem_too, - t_never_the_database, t_never_the_input_itself, - t_relative_path_is_still_the_input, t_missing_files_are_fine, - t_leaves_the_neighbours_alone, t_load_args): + for fn in (t_no_lock_sweeping, t_load_args): print(f"\n{fn.__name__}") try: fn() diff --git a/tests/test_pool.py b/tests/test_pool.py index 94e197b..9fc1446 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 -"""Unit tests for idatui.pool (worker residency: LRU + memory budget). +"""Unit tests for idatui.pool (Code Mode lease residency and LRU budget). -Pure stdlib with a fake client injected, so the eviction policy is testable -without spawning real idalib workers. +A fake client keeps the policy testable without IDA or Textual. python tests/test_pool.py """ @@ -15,7 +14,7 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.pool import WorkerPool # noqa: E402 +from idatui.pool import DatabasePool # noqa: E402 from idatui.project import Project # noqa: E402 PASS = FAIL = 0 @@ -32,11 +31,12 @@ def check(name, cond, detail=""): class FakeClient: - """Stands in for a WorkerClient: records saves/closes, reports fixed memory.""" + """Stands in for a CodeModeClient lease and records saves/closes.""" - def __init__(self, ref, mem=100): + def __init__(self, ref, mem=100, backend="idalib"): self.ref = ref self.mem = mem + self.backend = backend self.saved = 0 self.closed = False self.connected = False @@ -45,10 +45,9 @@ class FakeClient: self.connected = True return self - def call(self, tool, **kw): - if tool == "idb_save": - self.saved += 1 - return {} + def save_database(self): + self.saved += 1 + return {"saved": True} def close(self, grace=None): self.closed = True @@ -76,15 +75,15 @@ def main() -> int: made[ref.label] = c return c - pool = WorkerPool(proj, budget_mb=350, spawn=spawn, + pool = DatabasePool(proj, budget_mb=350, spawn=spawn, mem_fn=lambda c: c.mem) # -- lazy spawn + reuse -------------------------------------------- # a = pool.get("bin0") - check("get() spawns a worker on first use", a is made["bin0"] and a.connected) + check("get() spawns a database lease on first use", a is made["bin0"] and a.connected) check("get() stages the binary first", os.path.isfile(proj.by_label("bin0").staged)) - check("get() reuses the resident worker", pool.get("bin0") is a) + check("get() reuses the resident lease", pool.get("bin0") is a) check("resident() reports it", pool.resident() == ["bin0"], pool.resident()) # -- LRU ordering ---------------------------------------------------- # @@ -100,9 +99,9 @@ def main() -> int: check("exceeding the budget evicts the least-recently-used", pool.evicted == ["bin1"] and not pool.is_resident("bin1"), f"evicted={pool.evicted} resident={pool.resident()}") - check("the just-spawned worker is never the victim", pool.is_resident("bin3")) + check("the just-attached lease is never the victim", pool.is_resident("bin3")) check("eviction saves the database first", made["bin1"].saved == 1) - check("eviction closes the worker", made["bin1"].closed) + check("eviction closes the lease", made["bin1"].closed) check("pool is back within budget", pool.memory_mb() <= pool.budget_mb, f"{pool.memory_mb()}/{pool.budget_mb}") @@ -116,7 +115,7 @@ def main() -> int: # -- pinning ---------------------------------------------------------- # pool.close_all() - pool2 = WorkerPool(proj, budget_mb=250, spawn=spawn, mem_fn=lambda c: c.mem) + pool2 = DatabasePool(proj, budget_mb=250, spawn=spawn, mem_fn=lambda c: c.mem) pool2.get("bin0") pool2.pin("bin0") pool2.get("bin1") @@ -143,7 +142,7 @@ def main() -> int: # -- teardown ----------------------------------------------------------- # pool2.close_all() - check("close_all() closes every worker", + check("close_all() closes every lease", not pool2.resident() and all(c.closed for c in made.values())) check("close_all() clears the active binary", pool2.active is None) @@ -155,8 +154,8 @@ def main() -> int: check("an unknown label raises KeyError", True) # -- default budget comes from the project's memory_pct ------------------- # - pool3 = WorkerPool(proj, spawn=spawn, mem_fn=lambda c: c.mem) - check("default budget is derived, not a fixed worker count", + pool3 = DatabasePool(proj, spawn=spawn, mem_fn=lambda c: c.mem) + check("default budget is derived, not a fixed lease count", pool3.budget_mb >= 256, pool3.budget_mb) # -- prewarm: speculative, and never at the cost of a real binary ------ # @@ -169,7 +168,7 @@ def main() -> int: made2[ref.label] = c return c - pool = WorkerPool(proj, budget_mb=250, spawn=spawn2, + pool = DatabasePool(proj, budget_mb=250, spawn=spawn2, mem_fn=lambda c: c.mem) labels = [r.label for r in proj.refs] a, b, c_ = labels[0], labels[1], labels[2] @@ -188,6 +187,28 @@ def main() -> int: check("prewarm ignores a label outside the project", pool.prewarm("nope") is False) + # Budget eviction releases GUI leases but must not save somebody's open IDA + # implicitly. An explicit save-and-close remains authoritative. + with tempfile.TemporaryDirectory() as tmp: + proj = _mkproject(tmp, n=1) + made_gui = [] + + def spawn_gui(ref, ttl): + client = FakeClient(ref, backend="gui") + made_gui.append(client) + return client + + pool = DatabasePool(proj, spawn=spawn_gui, mem_fn=lambda c: c.mem) + label = proj.refs[0].label + pool.get(label) + pool.evict(label) + check("LRU release does not implicitly save a GUI database", + made_gui[-1].saved == 0) + pool.get(label) + pool.close_all(save=True) + check("explicit close_all(save=True) does save a GUI database", + made_gui[-1].saved == 1) + print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_project.py b/tests/test_project.py index 822002f..690f360 100644 --- a/tests/test_project.py +++ b/tests/test_project.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Unit tests for idatui.project (the multi-binary project model + staging). -Pure stdlib: no IDA, no textual, no worker — runs anywhere in under a second. +IDA-free: exercises staging plus Code Mode ownership checks without opening a database. python tests/test_project.py """ diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py index 9b2cc7e..e19ea11 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -694,18 +694,18 @@ async def s_split_view(c: Ctx): # to count. The bound is loose because the bug was three orders of magnitude # out, not a near miss. _lookups = {"n": 0} - _orig_call = c.prog.client.call + _orig_call = c.prog.client.invoke def _counting(name, *a, **kw): if name == "lookup_funcs": _lookups["n"] += 1 return _orig_call(name, *a, **kw) - c.prog.client.call = _counting + c.prog.client.invoke = _counting try: await _split_view_body(c, app, lst, dec) finally: - c.prog.client.call = _orig_call + c.prog.client.invoke = _orig_call c.check("split view doesn't storm the worker with function lookups", _lookups["n"] < 500, f"{_lookups['n']} lookup_funcs calls") @@ -1586,14 +1586,14 @@ async def s_decomp_nav(c: Ctx): old_ea = dec._line_ea(drow) if old_ea is not None and dsym in old_line: tmp = f"stale_{os.getpid()}" - app.program.client.call("rename", batch={"func": {"addr": hex(dstale), "name": tmp}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(dstale), "name": tmp}}) app.program.bump_names() d2 = len(app._nav) app._follow_decomp(old_line, dsym, old_ea) await c.wait(lambda: len(app._nav) > d2, 25) c.check("decomp follow works with a stale name (ea-marker fallback)", app._cur.ea == dstale, f"cur={app._cur.ea:#x} want={dstale:#x}") - app.program.client.call("rename", batch={"func": {"addr": hex(dstale), "name": dsym}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(dstale), "name": dsym}}) app.program.bump_names() @@ -1674,7 +1674,7 @@ async def s_rename(c: Ctx): c.check("rename updates the function name", app._func_index.by_addr(dtarget).name == newname, app._func_index.by_addr(dtarget).name) - rr = app.program.client.call("rename", batch={"func": {"addr": hex(dtarget), "name": dsym}}) + rr = app.program.client.invoke("rename", batch={"func": {"addr": hex(dtarget), "name": dsym}}) c.check("rename reverted cleanly", rr.get("summary", {}).get("ok", 0) == 1, str(rr.get("summary"))) # goto label refuse @@ -1724,7 +1724,7 @@ async def s_rename(c: Ctx): and any(cnote in t for t in dec._texts), 25) c.check("comment appears in the pseudocode after ';'", any(cnote in t for t in dec._texts), "comment not shown") - app.program.client.call("set_comments", items=[{"addr": hex(cea), "comment": ""}]) + app.program.client.invoke("set_comments", items=[{"addr": hex(cea), "comment": ""}]) else: c.check("found a pseudocode line to comment", False, "no marker line") @@ -1766,7 +1766,7 @@ async def s_comment_func(c: Ctx): la is not None and lb is not None and lb > la and a not in dec._texts[lb], f"la={la} lb={lb}") - app.program.client.call("set_comments", items=[{"addr": hex(fn.addr), "comment": ""}]) + app.program.client.invoke("set_comments", items=[{"addr": hex(fn.addr), "comment": ""}]) @scenario("retype") @@ -2031,7 +2031,7 @@ async def s_rename_history(c: Ctx): await c.pause(0.15) c.check("caller pseudocode shows renamed callee after 'back'", any(hnew in tx for tx in dec._texts), "pseudocode still stale") - app.program.client.call("rename", batch={"func": {"addr": hex(htarget), "name": hsym}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(htarget), "name": hsym}}) @scenario("region_define") @@ -2250,7 +2250,7 @@ async def s_listing_name_addr(c: Ctx): finally: # revert: drop the label and restore raw bytes at A try: - c.prog.client.call("rename", batch={"data": {"addr": hex(A + 1), "new": ""}}) + c.prog.client.invoke("rename", batch={"data": {"addr": hex(A + 1), "new": ""}}) except Exception: # noqa: BLE001 pass c.prog.undefine(A, size=8) @@ -2327,7 +2327,7 @@ async def s_listing_struct_expand(c: Ctx): c.check("found a data address for the struct test", False) return try: - c.prog.client.call( + c.prog.client.invoke( "declare_type", decls=["struct TuiExpandS { int a; char b[4]; short c; };"]) c.prog.make_data(A, "TuiExpandS") @@ -3206,7 +3206,7 @@ async def s_graph_rename(c: Ctx): f"resolve({new}) -> {got if got is None else hex(got)} want {ea:#x}") # revert, so the suite stays idempotent if got is not None: - app.program.client.call( + app.program.client.invoke( "rename", batch={"data": {"addr": hex(ea), "new": ""}}) app.program.bump_names() @@ -3259,7 +3259,7 @@ async def run(binary, only=None): async def _run_on(binary, only=None): - # Own idalib worker: opens the binary in-process over a unix socket. + # Code Mode attaches a registered GUI or starts/reuses a managed worker. app = IdaTui(open_path=binary, keepalive=False) async with app.run_test(size=(140, 44)) as pilot: c = Ctx(app, pilot) @@ -3279,6 +3279,14 @@ async def _run_on(binary, only=None): print(f"── {name} ({asyncio.get_event_loop().time() - _t0:.1f}s) CRASHED") c.check("scenario did not crash", False, f"{type(e).__name__}: {e}") traceback.print_exc() + # Headless run_test does not reliably emit App.Unmount; explicitly release + # the lease. Then wait through the managed worker's final-lease grace and + # IDB close so Windows can remove this suite's TemporaryDirectory safely. + if app.program is not None: + app.program.close() + if app.client is not None: + app.client.close() + await asyncio.to_thread(app.client.wait_released, 45.0) def main(argv): 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()) |