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 | |
| 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')
| -rw-r--r-- | tests/_fixtures.py | 38 | ||||
| -rwxr-xr-x | tests/run.py | 4 | ||||
| -rw-r--r-- | tests/test_blob_ui.py | 109 | ||||
| -rw-r--r-- | tests/test_codemode_client.py | 162 | ||||
| -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_thumb_ui.py | 95 | ||||
| -rw-r--r-- | tests/test_trace_ui.py | 12 | ||||
| -rw-r--r-- | tests/test_worker_client.py | 447 |
11 files changed, 436 insertions, 652 deletions
diff --git a/tests/_fixtures.py b/tests/_fixtures.py index 766ecd5..0b72856 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -37,6 +37,36 @@ def cache_is_fresh(binary: str) -> bool: return os.path.exists(c) and os.path.getmtime(c) >= os.path.getmtime(binary) +#: Generated targets live here so their pristine caches survive between runs. +#: Gitignored; safe to delete (the next run rebuilds both). +SYNTHETIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".synthetic") + + +def synthetic(name: str, build) -> str: + """A generated binary at a STABLE path, rebuilt only when its bytes change. + + Generated targets used to be written into a fresh TemporaryDirectory on + every run, which quietly defeated the whole pristine-cache scheme: a new + path with new bytes every time means auto-analysis is paid in full, every + run, forever. `test_blob_ui`'s 64KB blob cost ~40s a run that way. + + ``build()`` must be DETERMINISTIC and return bytes. That is also what makes + the suites reproducible: a blob built from os.urandom can, by luck, contain + something IDA reads as a function, and then a test asserting "no functions" + fails for reasons no one can reproduce. + """ + os.makedirs(SYNTHETIC_DIR, exist_ok=True) + path = os.path.join(SYNTHETIC_DIR, name) + data = build() + if not os.path.exists(path) or open(path, "rb").read() != data: + with open(path, "wb") as fh: # content changed -> cache is stale + fh.write(data) + for stale in (cache_path(path), path + ".i64"): + if os.path.exists(stale): + os.remove(stale) + return path + + async def build_pristine(binary: str, cache: str, app_factory) -> None: """Analyse ``binary`` once and keep the database as a golden copy. @@ -50,7 +80,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_blob_ui.py b/tests/test_blob_ui.py index f60e914..1055e5a 100644 --- a/tests/test_blob_ui.py +++ b/tests/test_blob_ui.py @@ -17,10 +17,13 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, Static # noqa: E402 from idatui.app import ConfirmScreen, IdaTui, ListingView # noqa: E402 +from _fixtures import staged, synthetic # noqa: E402 +from idatui._sync import settle # noqa: E402 PASS = FAIL = 0 @@ -46,28 +49,63 @@ async def wait(pred, pilot, t=240.0): return False -async def run() -> int: - with tempfile.TemporaryDirectory() as tmp: - # Random bytes so IDA finds no functions... but with REAL AArch64 - # instructions planted at a known offset. Whether arbitrary random bytes - # happen to decode is chance, and a test that depends on chance tells you - # nothing on the run where it fails. - data = bytearray(os.urandom(64 * 1024)) +#: File offset of the planted instruction run -> ea 0x4000 + PLANTED. +PLANTED = 0x40 + + +def _blob_bytes() -> bytes: + """A DETERMINISTIC pseudo-random blob with real AArch64 instructions planted. + + Seeded, not os.urandom: the bytes must be identical every run or the + pristine-database cache can never apply (this suite used to pay ~40s of + auto-analysis per run because the content, and the path, changed each time). + Determinism also removes a genuine flake -- whether 64KB of chance bytes + contains something IDA reads as a function is luck, and "and really has no + functions" is asserted below. + """ + import random + data = bytearray(random.Random(0xB10BCAFE).randbytes(64 * 1024)) # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 # spelling of a nop (0xE1A00000) is NOT decodable there and made this # test fail for a reason that had nothing to do with what it checks. - planted = 0x40 # file offset -> ea 0x4040 - for k, insn in enumerate((0xD503201F, # nop - 0xD503201F, # nop - 0xD65F03C0)): # ret <- the run must stop here - data[planted + k * 4:planted + k * 4 + 4] = insn.to_bytes(4, "little") - blob = os.path.join(tmp, "rnd.bin") - with open(blob, "wb") as f: - f.write(bytes(data)) + for k, insn in enumerate((0xD503201F, # nop + 0xD503201F, # nop + 0xD65F03C0)): # ret <- the run must stop here + data[PLANTED + k * 4:PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") + return bytes(data) + + +#: -parm puts IDA in AArch64 mode (the ARM32 spelling of a nop is not decodable +#: there); -b400 sets the image base. The cached database must be built with the +#: SAME switches, so both go through one factory. +BLOB_ARGS = "-parm -b400" + +def _blob_app(path): + return IdaTui(open_path=path, keepalive=False, load_args=BLOB_ARGS) + + +def head_at(lst, ea): + """The listing row for ``ea`` off the LIVE model, or None. + + Always re-reads ``lst.model``: an edit may rebuild the model, and holding + the old object shows pre-edit rows -- which looks exactly like the edit + silently failing. + """ + m = lst.model + if m is None: + return None + i = m.index_of_ea(ea) + return m.get(i) if i is not None and i >= 0 else None + + +async def run() -> int: + blob_src = synthetic("rnd.bin", _blob_bytes) + # staged() analyses once ever and copies the result in on later runs. + async with staged(blob_src, _blob_app) as blob: # Skip the dialog by answering up front; this test is about what # happens AFTER a described blob turns out to contain nothing. - app = IdaTui(open_path=blob, keepalive=False, load_args="-parm -b400") + app = _blob_app(blob) async with app.run_test(size=(140, 44)) as pilot: ok = await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -126,7 +164,7 @@ async def run() -> int: i >= 0 and m.get(i).ea == 0x4021, f"row={i} ea={m.get(i).ea if i >= 0 else None}") - target = 0x4000 + planted # a NOP we put there ourselves + target = 0x4000 + PLANTED # a NOP we put there ourselves lst.cursor = m.index_of_ea(target) lst._scroll_cursor_into_view() await pilot.pause(0.1) @@ -134,12 +172,18 @@ async def run() -> int: lst._cursor_ea() == target, f"{lst._cursor_ea():#x} want {target:#x}") await pilot.press("c") - await pilot.pause(2.0) - # Defining an item REBUILDS the listing model, so re-read it from the - # view: holding the old object shows the pre-edit rows and looks - # exactly like the edit silently failing. + # settle(), not a fixed sleep AND not a bare predicate: an edit can + # look done for a moment and then be replaced when a queued listing + # rebuild lands, so the gate has to be "the row is code AND the app + # has stopped working". settle() is the same helper the app's own + # RPC layer uses, so tests and driver agree on what "done" means. + await settle(app, lambda: (lambda h: h is not None and h.kind == "code")( + head_at(lst, target)), timeout=30) + # Re-read the model: defining an item rebuilds it, and holding the + # old object shows pre-edit rows -- which looks exactly like the + # edit silently failing. m = lst.model - h = m.get(m.index_of_ea(target)) + h = head_at(lst, target) check("`c` on a chosen byte carves an instruction there", h is not None and h.kind == "code", f"kind={h.kind if h else None} text={h.text if h else None!r}") @@ -184,9 +228,17 @@ async def run() -> int: for ch in "note": await pilot.press(ch) await pilot.press("enter") - await wait(lambda: lst.model is not old and lst.model is not None, - pilot, 30) - await pilot.pause(0.4) + # Wait for the COMMENT ITSELF to show up, not for the model object to + # be replaced: a comment now re-renders the listing in place (the + # walk is kept), so `model is not old` never becomes true and this + # burned its full 30s timeout on every run -- after which the check + # below passed vacuously, because nothing had happened at all. + # The prompt closing plus quiescence is the real end of the edit. + # (The listing re-renders its text lazily, so the comment is not + # necessarily visible in model rows the moment the worker returns -- + # which is why this waits for the app, not for the text.) + await settle(app, lambda: not app.query_one("#comment", Input).display, + timeout=30) check("commenting leaves the view where it was", lst.model.get(round(lst.scroll_offset.y)).ea == ctop and lst._cursor_ea() == ccur, @@ -208,7 +260,12 @@ async def run() -> int: check("scrolled somewhere with rows above us", round(lst.scroll_offset.y) > 0, f"top={lst.scroll_offset.y}") await pilot.press("c") - await pilot.pause(2.5) + # No predicate here on purpose: this spot is random data, so the + # carve may legitimately produce nothing and "the row became code" + # would never hold (it timed out for 30s and then passed anyway). + # What is being checked is that the VIEW did not move, so the gate + # is simply "the app has finished reacting". + await settle(app, timeout=30) m2 = lst.model top_after = m2.get(round(lst.scroll_offset.y)).ea check("carving leaves the scroll position where it was", diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py new file mode 100644 index 0000000..2f70ba3 --- /dev/null +++ b/tests/test_codemode_client.py @@ -0,0 +1,162 @@ +"""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 _open_kwargs_are_real(sent: dict): + """(ok, detail) for the kwargs the adapter passes to DatabaseHandle.open. + + Skips (passes) when ida_codemode is not installed, so the file stays pure. + """ + try: + import inspect + from ida_codemode.client import DatabaseHandle as Real + except ImportError: + return True, "ida_codemode not installed - signature not checked" + accepted = set(inspect.signature(Real.open).parameters) + unknown = sorted(set(sent) - accepted) + return not unknown, f"open() rejects {unknown}" + + +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["image_base"] == 0x1000, + FakeDatabaseHandle.kwargs) + # A fake that swallows **kwargs cannot catch a keyword the real + # library does not have -- which is exactly how this port shipped + # `loading_address` (the real name is `image_base`) and would have + # raised TypeError on the very first connect. Check the names we + # send against the real signature whenever it is importable. + check("every open() keyword exists in the real library", + *_open_kwargs_are_real(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_thumb_ui.py b/tests/test_thumb_ui.py index 92a19fb..fed08a7 100644 --- a/tests/test_thumb_ui.py +++ b/tests/test_thumb_ui.py @@ -15,12 +15,15 @@ Needs IDA. ~40s. NEEDS_IDA = True import asyncio import os +import shutil import sys +import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from textual.widgets import Static # noqa: E402 +from idatui._sync import settle # noqa: E402 from idatui.app import DecompView, IdaTui, ListingView # noqa: E402 PASS = FAIL = 0 @@ -38,6 +41,36 @@ def check(name, ok, detail=""): print(f" FAIL {name} {detail}") + +#: Every phase gets its OWN copy of the fixture. +#: +#: This suite used to delete <BIN>.i64 and reopen the SAME path for each phase. +#: That was safe when the TUI owned a private worker that died with it; under +#: Code Mode the database is leased and the previous phase's worker can still +#: hold it through its lease grace, so the delete raced a live owner and the +#: next open never produced a listing (the crash this fixed). Separate paths +#: cannot collide, and nothing has to wait for anyone else to let go. +_SCRATCH = [] + + +def fresh_copy(src: str, tag: str) -> str: + d = tempfile.mkdtemp(prefix=f"idatui-thumb-{tag}-") + _SCRATCH.append(d) + dst = os.path.join(d, os.path.basename(src)) + shutil.copy2(src, dst) + return dst + + +def drop_scratch() -> None: + for d in _SCRATCH: + shutil.rmtree(d, ignore_errors=True) + _SCRATCH.clear() + + +def status_of(app) -> str: + return str(app.query_one("#status", Static).render()) + + async def wait(pred, pilot, t=240.0): for _ in range(int(t / 0.05)): await pilot.pause(0.05) @@ -52,13 +85,8 @@ async def wait(pred, pilot, t=240.0): async def run() -> int: # A fresh database every time: the T flag and the segment's addressing mode # are SAVED in the .i64, so a previous run would answer the question for us. - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm") + app = IdaTui(open_path=fresh_copy(BIN, "arm"), keepalive=False, + load_args="-parm") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -84,10 +112,14 @@ async def run() -> int: m1 = lst.model await pilot.press("t") - await wait(lambda: lst.model is not m1 and lst.model is not None, pilot, 60) - await pilot.pause(0.5) + # The mode switch announces itself; wait for THAT, plus quiescence. + # `lst.model is not m1` used to be the gate, but an item edit now keeps + # the listing's walk instead of rebuilding it, so the model object is + # never replaced -- every one of these waits sat out its full 60s and + # the suite still "passed", four times over. + await settle(app, lambda: "Thumb" in status_of(app), timeout=60) - status = str(app.query_one("#status", Static).render()) + status = status_of(app) check("the status says it switched to Thumb", "Thumb" in status, status[:90]) # Thumb doesn't exist in AArch64, and -parm on a headerless blob gives a # 64-bit segment, so setting T alone would change nothing and look broken. @@ -114,9 +146,8 @@ async def run() -> int: m2 = lst.model lst.cursor = lst.model.index_of_ea(0) await pilot.press("t") - await wait(lambda: lst.model is not m2 and lst.model is not None, pilot, 60) - await pilot.pause(0.5) - status = str(app.query_one("#status", Static).render()) + await settle(app, lambda: "ARM @" in status_of(app), timeout=60) + status = status_of(app) check("`t` toggles back to ARM", "ARM @" in status, status[:80]) # -- and the reason a carved function wouldn't decompile ---------------- # @@ -126,12 +157,8 @@ async def run() -> int: # disassembly that F5 can never turn into pseudocode. The database's bitness # is fixed at load and cannot be corrected afterwards, so the only honest # thing is to say so. - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm") # 64-bit + app = IdaTui(open_path=fresh_copy(BIN, "arm64"), keepalive=False, + load_args="-parm") # 64-bit async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -143,9 +170,8 @@ async def run() -> int: await pilot.pause(0.3) m = lst.model await pilot.press("t") - await wait(lambda: lst.model is not m and lst.model is not None, pilot, 60) - await pilot.pause(0.5) - status = str(app.query_one("#status", Static).render()) + await settle(app, lambda: "64-bit" in status_of(app), timeout=60) + status = status_of(app) check("a 64-bit database warns that Hex-Rays won't decompile", "64-bit" in status and "decompile" in status, status[:120]) check("and names the fix", "ARMv7-A" in status, status[:120]) @@ -158,9 +184,10 @@ async def run() -> int: await pilot.pause(0.2) mp = lst.model await pilot.press("p") - await wait(lambda: lst.model is not mp and lst.model is not None, pilot, 60) - await wait(lambda: app._func_index is not None - and len(app._func_index) > 0, pilot, 60) + # The function appearing in the index IS the signal; the model identity + # never was one. + await settle(app, lambda: app._func_index is not None + and len(app._func_index) > 0, timeout=60) await pilot.press("tab") await wait(lambda: "cannot decompile" in str(app.query_one("#status", Static).render()), pilot, 90) @@ -177,12 +204,8 @@ async def run() -> int: len(status) < 110, f"{len(status)} chars: {status[:130]}") # -- the whole point: a 32-bit database decompiles ---------------------- # - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm:ARMv7-A") + app = IdaTui(open_path=fresh_copy(BIN, "armv7a"), keepalive=False, + load_args="-parm:ARMv7-A") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -216,12 +239,8 @@ async def run() -> int: if not os.path.isfile(vec): check("the cortexm fixture exists", False, vec) else: - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(vec + ext) - except OSError: - pass - app = IdaTui(open_path=vec, keepalive=False, load_args="-parm:ARMv7-M") + app = IdaTui(open_path=fresh_copy(vec, "cortexm"), keepalive=False, + load_args="-parm:ARMv7-M") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -253,6 +272,8 @@ async def run() -> int: check("the result survives the reload AND the reindex", "3 Thumb entries" in status, status[:90]) + drop_scratch() + print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_trace_ui.py b/tests/test_trace_ui.py index bb016eb..5844c90 100644 --- a/tests/test_trace_ui.py +++ b/tests/test_trace_ui.py @@ -20,6 +20,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, OptionList, Static # noqa: E402 from _fixtures import staged # noqa: E402 +from idatui._sync import settle # noqa: E402 from idatui.app import (DecompView, IdaTui, ListingView, # noqa: E402 RegWriteScreen, TraceDock) @@ -109,13 +110,20 @@ async def run() -> int: lst.focus() await pilot.pause(0.4) await pilot.press("]") - await wait(lambda: app._t == 1, pilot, 20) + # `app._t` is assigned the moment the key is handled, so it is NOT a + # signal that the VIEW has followed -- the navigation it kicks off + # runs in a worker. Waiting on it and then reading the cursor was a + # race that the (slower) Code Mode backend loses. Gate on the thing + # the check is about. + await settle(app, lambda: app._t == 1 and lst._cursor_ea() == t.ip(1), + timeout=20) check("] steps forward one instruction", app._t == 1, f"t={app._t}") check("the code view follows the trace", lst._cursor_ea() == t.ip(1), f"{lst._cursor_ea()} vs {t.ip(1)}") await pilot.press("[") - await wait(lambda: app._t == 0, pilot, 20) + await settle(app, lambda: app._t == 0 and lst._cursor_ea() == t.ip(0), + timeout=20) check("[ steps backward", app._t == 0, f"t={app._t}") await pilot.press("[") await pilot.pause(0.4) 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()) |
