aboutsummaryrefslogtreecommitdiffstats
path: root/idatui/worker.py
diff options
context:
space:
mode:
authoruser <user@clank>2026-08-07 15:14:30 +0200
committeruser <user@clank>2026-08-07 15:14:30 +0200
commit72fce7da1a1fd6527e389ffeb0f951157523589a (patch)
tree3f08a83f0d99f3d3fc7069b7be00d235bab82817 /idatui/worker.py
parentStop tracking 157MB of core dumps, and ignore them (diff)
parentdocs: upstream findings for the ida-codemode maintainers (diff)
downloadida-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 'idatui/worker.py')
-rw-r--r--idatui/worker.py301
1 files changed, 0 insertions, 301 deletions
diff --git a/idatui/worker.py b/idatui/worker.py
deleted file mode 100644
index 556e69a..0000000
--- a/idatui/worker.py
+++ /dev/null
@@ -1,301 +0,0 @@
-"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor.
-
-Opens ONE database in-process (on the main thread, as idalib requires) and
-serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed
-pickle. Same tool implementations as the MCP path (we call
-``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are
-byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no
-supervisor / HTTP / JSON / 50KB-truncation machinery.
-
- python -m idatui.worker <sock_path> <binary_path>
-
-The socket only appears once the database is open + analyzed, so a client can
-poll ``connect()`` to know when the worker is ready. Requests are served
-serially on the main thread (idalib is single-threaded; every tool runs inline
-through its own execute_sync, which is a no-op on the main thread).
-
-Protocol (both directions length-prefixed: 4-byte big-endian len + pickle):
- request = (tool_name: str, kwargs: dict)
- response = (ok: bool, result_or_error)
- tool_name == "__shutdown__" ends the worker.
-"""
-from __future__ import annotations
-
-import os
-import pickle
-import socket
-import struct
-import sys
-import threading
-import time
-import uuid
-
-#: Seconds a single tool call may run before it is cancelled. 0 disables the
-#: deadline entirely.
-TOOL_TIMEOUT_SEC = float(os.environ.get("IDATUI_TOOL_TIMEOUT_SEC") or 60)
-
-# ida-pro-mcp enforces its own tool deadline by installing a `sys.setprofile`
-# hook for the duration of every call, so that a pure-python loop inside a tool
-# body can be interrupted. That hook runs a python function on EVERY python call
-# and return -- and our tools are exactly the call-heavy kind: `heads` renders
-# hundreds of items per request and measured 92us/row with the hook against
-# 28us/row without it. A 3.3x tax on the whole backend to bound loops that are
-# already bounded by their `count` argument.
-#
-# So: turn the upstream mechanism off and re-arm the half that does the real
-# work ourselves (see _Deadline). ida_kernwin.set_cancelled() is what actually
-# frees the IDA main thread -- decompile, auto_wait, find_bytes and friends poll
-# user_cancelled() and bail within a poll cycle -- and it costs nothing until it
-# fires.
-os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0"
-
-
-class _Deadline:
- """A single watchdog thread that cancels a tool call which overruns.
-
- Arming is two attribute writes, because it is on the path of every call the
- TUI makes (a scroll is dozens of them). The watchdog polls instead of being
- signalled for the same reason: waking a thread per call costs more than the
- 0.25s of granularity it buys on a 60s deadline.
- """
-
- TICK = 0.25
-
- def __init__(self, seconds: float) -> None:
- import ida_kernwin
- self._kernwin = ida_kernwin
- self.seconds = seconds
- self._until: float | None = None
- t = threading.Thread(target=self._run, name="idatui-deadline",
- daemon=True)
- t.start()
-
- def _run(self) -> None:
- while True:
- time.sleep(self.TICK)
- until = self._until
- if until is not None and time.monotonic() >= until:
- self._until = None
- # THREAD_SAFE in the IDA SDK; upstream fires it off a Timer too.
- self._kernwin.set_cancelled()
-
- def arm(self) -> None:
- # Clear unconditionally: the flag is sticky, and one left set would make
- # every later user_cancelled() true forever.
- self._kernwin.clr_cancelled()
- self._until = time.monotonic() + self.seconds
-
- def disarm(self) -> None:
- self._until = None
-
-
-# --------------------------------------------------------------------------- #
-# framing
-# --------------------------------------------------------------------------- #
-def _recvn(sock: socket.socket, n: int) -> bytes | None:
- buf = bytearray()
- while len(buf) < n:
- chunk = sock.recv(n - len(buf))
- if not chunk:
- return None
- buf += chunk
- return bytes(buf)
-
-
-def send(sock: socket.socket, obj) -> None:
- data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
- sock.sendall(struct.pack(">I", len(data)) + data)
-
-
-def recv(sock: socket.socket):
- hdr = _recvn(sock, 4)
- if hdr is None:
- return None
- (n,) = struct.unpack(">I", hdr)
- body = _recvn(sock, n)
- return None if body is None else pickle.loads(body)
-
-
-# --------------------------------------------------------------------------- #
-# worker
-# --------------------------------------------------------------------------- #
-def _ensure_tools_injected() -> None:
- """Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...)
- into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient
- (nothing else has to inject these tools first). Must run BEFORE
- ida_pro_mcp.ida_mcp is imported (the injected code lives in api_types.py)."""
- import importlib.util
- repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
- patch = os.path.join(repo, "server", "patch_server.py")
- if not os.path.exists(patch):
- return
- try:
- spec = importlib.util.spec_from_file_location("_idatui_patch", patch)
- mod = importlib.util.module_from_spec(spec)
- spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types
- mod.main()
- except Exception as e: # noqa: BLE001 -- tools may already be present
- sys.stderr.write(f"idatui: tool injection skipped: {e}\n")
-
-
-def _has_database(binpath: str) -> bool:
- """Whether IDA already has a database for ``binpath``.
-
- IDA names it ``<file>.i64`` (keeping the extension), but a database made
- from ``foo.bin`` can also appear as ``foo.i64`` depending on how it was
- created — check both, because guessing wrong here means re-passing load
- switches to an existing database, which fails the open.
- """
- return (os.path.exists(binpath + ".i64")
- or os.path.exists(os.path.splitext(binpath)[0] + ".i64"))
-
-
-def _open_and_register(binpath: str, load_args: str = ""):
- """Open the DB (main thread) then import ida-pro-mcp so every @tool registers
- against this live database. Returns (tools_dict, module_name, save_fn).
-
- ``load_args`` is passed to IDA as command-line switches, which is the only
- way to tell it how to read a headerless blob: a raw firmware image has no
- format to detect, so without ``-p<processor>`` it loads as metapc at 0 and
- finds nothing. Ignored once a database exists — the .i64 already records how
- it was loaded, and re-passing conflicting switches is how you corrupt one.
- """
- _ensure_tools_injected() # before any ida_pro_mcp import
- import idapro
- idapro.enable_console_messages(False)
- args = load_args or None
- if args and _has_database(binpath):
- # The .i64 already records how this image was loaded. Passing the
- # switches again on reopen makes IDA fail outright (rc != 0) — the load
- # options belong to the FIRST open only.
- args = None
- if idapro.open_database(binpath, run_auto_analysis=True,
- args=args): # nonzero == failure
- if args:
- # With load switches in play they are the likeliest culprit by far:
- # IDA refuses an unknown -p name with no diagnostic of its own, so
- # saying "the database is locked" here sends people hunting a
- # problem they don't have.
- raise RuntimeError(
- f"failed to open {binpath} with load options {args!r}: IDA "
- f"rejected them \u2014 an unknown processor name is the usual "
- f"cause (see tools/verify_procs.py for the valid ones)")
- raise RuntimeError(
- f"failed to open {binpath}: the .i64 is likely held by a running "
- f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash "
- f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)")
- import ida_auto
- ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp)
-
- # importing the package registers all api_*/patched tools against MCP_SERVER
- from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433
-
- import ida_nalt
- module = os.path.basename(ida_nalt.get_root_filename() or binpath)
-
- def save():
- import idc
- try:
- idc.save_database(idc.get_idb_path(), 0)
- except Exception: # noqa: BLE001
- import ida_loader, ida_pro # noqa: WPS433
- ida_loader.save_database(idc.get_idb_path(), 0)
-
- return MCP_SERVER.tools.methods, module, save
-
-
-def serve(sockpath: str, binpath: str, load_args: str = "") -> None:
- tools, module, save = _open_and_register(binpath, load_args)
- sid = uuid.uuid4().hex[:8]
- deadline = _Deadline(TOOL_TIMEOUT_SEC) if TOOL_TIMEOUT_SEC > 0 else None
-
- def dispatch(name: str, args: dict):
- args = dict(args)
- args.pop("database", None) # single-DB worker: no session routing
- # session-management shims (were the supervisor's job):
- if name in ("idb_open",):
- return {"success": True,
- "session": {"session_id": sid, "module": module,
- "input_path": binpath}}
- if name in ("idb_save", "save"):
- save()
- return {"success": True}
- if name in ("server_health", "ping", "health", "state"):
- return {"module": module, "ok": True, "session_id": sid}
- if name in ("idb_list",):
- return {"sessions": [{"session_id": sid, "module": module,
- "input_path": binpath}]}
- fn = tools.get(name)
- if fn is None:
- raise KeyError(f"unknown tool: {name!r}")
- if deadline is None:
- result = fn(**args)
- else:
- deadline.arm()
- try:
- result = fn(**args)
- finally:
- deadline.disarm()
- # Match the MCP server's structuredContent: a dict passes through, any
- # other return (list/scalar) is wrapped as {"result": ...}. domain.py
- # parses that exact shape (e.g. lookup_funcs -> payload["result"]).
- return result if isinstance(result, dict) else {"result": result}
-
- try:
- os.unlink(sockpath)
- except OSError:
- pass
- srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
- srv.bind(sockpath)
- srv.listen(8)
- try:
- while True:
- conn, _ = srv.accept()
- try:
- while True:
- req = recv(conn)
- if req is None:
- break
- name, args = req
- if name == "__shutdown__":
- return
- try:
- send(conn, (True, dispatch(name, args)))
- except Exception as e: # noqa: BLE001 -- report, keep serving
- send(conn, (False, f"{type(e).__name__}: {e}"))
- except (ConnectionError, OSError):
- pass
- finally:
- conn.close()
- finally:
- try:
- import idapro
- idapro.close_database(save=False)
- except Exception: # noqa: BLE001
- pass
- try:
- os.unlink(sockpath)
- except OSError:
- pass
-
-
-def main(argv=None) -> None:
- argv = argv if argv is not None else sys.argv[1:]
- if len(argv) < 2:
- sys.stderr.write(
- "usage: python -m idatui.worker <sock> <binary> [ida-load-args]\n")
- raise SystemExit(2)
- try:
- serve(argv[0], argv[1], argv[2] if len(argv) > 2 else "")
- except SystemExit:
- raise
- except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1
- import traceback
- sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n")
- traceback.print_exc()
- sys.stderr.flush()
- raise SystemExit(1)
-
-
-if __name__ == "__main__":
- main()