diff options
| author | blasty <blasty@local> | 2026-07-24 01:29:06 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-07-24 01:29:06 +0200 |
| commit | b14f26f0bb79baf9e1476f2d67e2cddd8b6957d1 (patch) | |
| tree | 4b32565a0c3f8eb9d806006aecdb24fd4e6e2167 /idatui/worker.py | |
| parent | experiments: add a unix-socket idalib worker as the 3rd bench column (diff) | |
| download | ida-tui-b14f26f0bb79baf9e1476f2d67e2cddd8b6957d1.tar.gz ida-tui-b14f26f0bb79baf9e1476f2d67e2cddd8b6957d1.tar.xz ida-tui-b14f26f0bb79baf9e1476f2d67e2cddd8b6957d1.zip | |
worker: idalib worker + WorkerClient (drop-in for IDAClient) — migration step 1
First concrete step off the mcp HTTP transport. Instead of reimplementing ~25
tools, reuse ida-pro-mcp's tool *functions* verbatim and replace only the
transport + process management:
* idatui/worker.py — opens ONE database in-process on the main thread (as idalib
requires), imports ida_pro_mcp (which registers every stock + our patched-in
custom tool against MCP_SERVER), then serves MCP_SERVER.tools.methods[name]
(**args) over a unix socket with length-prefixed pickle. Serial on the main
thread (idalib is single-threaded; tools run inline through execute_sync).
Session-management tools (idb_open/idb_save/server_health/idb_list) are shimmed
since the worker *is* the single session.
* idatui/worker_client.py — WorkerClient exposes the exact surface the app/domain
use on the client (call/call_envelope/connect/set_db/resolve_db/list_sessions/
health/keepalive/close) and returns byte-identical payloads (the worker calls
the same functions IDAClient.call ultimately hits). So domain.py and the app
are UNCHANGED — you just construct a WorkerClient instead of an IDAClient.
Calls are serialized under a lock over one socket; keepalive is a no-op (the
worker is ours and never idles out).
Not wired into the app yet — the mcp path is fully intact.
Verified without idalib: pickle framing round-trips arbitrary payloads incl raw
bytes; WorkerClient has full IDAClient surface; call_envelope produces the
result.structuredContent shape domain.decompile() reads. The idalib E2E
(experiments/worker_smoke.py drives the real domain.Program read path through the
worker) is written but couldn't run here — this sandbox has degraded to reaping
any idalib spawn; the underlying unix-socket protocol already ran clean in the
inproc_spike bench (~50us/call), and the worker dispatches the same tool
functions the HTTP path does, so shapes match by construction.
Next: stand up progress reporting during analysis, then flip _connect/_reconnect
to build a WorkerClient behind a flag and run the pilot suite against it.
Diffstat (limited to 'idatui/worker.py')
| -rw-r--r-- | idatui/worker.py | 160 |
1 files changed, 160 insertions, 0 deletions
diff --git a/idatui/worker.py b/idatui/worker.py new file mode 100644 index 0000000..d5fab20 --- /dev/null +++ b/idatui/worker.py @@ -0,0 +1,160 @@ +"""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 uuid + + +# --------------------------------------------------------------------------- # +# 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 _open_and_register(binpath: 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).""" + import idapro + idapro.enable_console_messages(False) + rc = idapro.open_database(binpath, run_auto_analysis=True) + if rc: + raise RuntimeError(f"open_database({binpath!r}) failed rc={rc}") + + # 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) -> None: + tools, module, save = _open_and_register(binpath) + sid = uuid.uuid4().hex[:8] + + 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}") + return fn(**args) + + 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>\n") + raise SystemExit(2) + serve(argv[0], argv[1]) + + +if __name__ == "__main__": + main() |
