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 | |
| 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.
| -rw-r--r-- | experiments/worker_smoke.py | 58 | ||||
| -rw-r--r-- | idatui/worker.py | 160 | ||||
| -rw-r--r-- | idatui/worker_client.py | 165 |
3 files changed, 383 insertions, 0 deletions
diff --git a/experiments/worker_smoke.py b/experiments/worker_smoke.py new file mode 100644 index 0000000..b215a57 --- /dev/null +++ b/experiments/worker_smoke.py @@ -0,0 +1,58 @@ +"""Runnable read-path smoke: drives the REAL domain.Program through WorkerClient +(our idalib worker over a unix socket). Run when idalib can spawn: + ~/ida-venv/bin/python experiments/worker_smoke.py +""" +import os, sys, shutil, time +REPO=os.path.expanduser("~/dev/ida-tui-maybe"); sys.path.insert(0, REPO); os.chdir(REPO) +# fresh copy so the worker's idalib doesn't fight any running server +src=f"{REPO}/targets/echo"; tmp="/tmp/echo_worker" +shutil.copy(src, tmp) +for e in ".i64 .id0 .id1 .id2 .nam .til".split(): + try: os.remove(tmp+e) + except OSError: pass + +from idatui.worker_client import WorkerClient +from idatui.domain import Program + +print("spawning worker + opening echo…", flush=True) +t=time.time() +cl=WorkerClient(tmp) +cl.connect(progress=lambda m: None) +print(f" worker ready in {time.time()-t:.2f}s session={cl.resolve_db()}", flush=True) +prog=Program(cl) + +# --- drive the REAL domain layer through the worker (read path) --- +main=prog.resolve("main") +print("resolve('main') =", hex(main), flush=True) + +idx=prog.functions(); idx.load_all() +print("functions() ->", len(idx), "funcs", flush=True) + +fn=prog.function_of(main) +print("function_of(main) ->", fn.name, hex(fn.addr), "size", fn.size, flush=True) + +b=prog.read_bytes(main, 16) +print("read_bytes(main,16) ->", b.hex(), flush=True) + +lm=prog.listing(main) +for _ in range(3): lm.load_next_page() +rows=[lm.get(i) for i in range(min(6,len(lm)))] +print("listing() first rows:", flush=True) +for h in rows: + if h: print(" ", hex(h.ea), h.kind, repr(h.text[:44]), flush=True) + +d=prog.decompile(main) +print("decompile(main) -> failed?", d.failed, "lines:", len((d.code or '').splitlines()), flush=True) + +regs=prog.file_regions() +print("file_regions ->", len(regs), "segments", flush=True) + +# xrefs to a called function +callee=next((f.addr for f in idx.all_loaded() if f.name.startswith("sub_")), None) +if callee: + xr=prog.xrefs_to(callee) + print("xrefs_to(", hex(callee), ") ->", len(xr), "refs", flush=True) + +ok = (fn.name=="main" and len(idx)>100 and b and not d.failed and len(regs)>0) +print("VERDICT:", "OK — domain.Program runs unchanged on the worker" if ok else "FAIL", flush=True) +cl.close() 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() diff --git a/idatui/worker_client.py b/idatui/worker_client.py new file mode 100644 index 0000000..49da8fa --- /dev/null +++ b/idatui/worker_client.py @@ -0,0 +1,165 @@ +"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own +idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's +HTTP/JSON transport. + +It exposes exactly the 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 tool functions), so ``domain.py`` and the +app are unchanged — you just construct a WorkerClient instead of an IDAClient. + +Concurrency: the app fires calls from several worker threads over one client; +the worker is single-threaded, so calls are serialized under a lock (the worker +processes one tool at a time anyway — and at ~50us/call that's free). +""" +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import threading +import time +import uuid +from typing import Any + +from .client import IDAToolError, IDAConnectionError, Session +from .worker import recv as _recv +from .worker import send as _send + + +class _NoopKeepAlive: + """The worker is ours and never idles out, so keepalive is a no-op.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self): + return self + + def stop(self) -> None: + pass + + +class WorkerClient: + def __init__(self, binary_path: str, *, ttl: int = 0, + python: str | None = None) -> None: + self._bin = os.path.abspath(os.path.expanduser(binary_path)) + self._python = python or sys.executable + self._sock_path = f"/tmp/idatui-worker-{os.getpid()}-{uuid.uuid4().hex[:8]}.sock" + self._proc: subprocess.Popen | None = None + self._sock: socket.socket | None = None + self._sid = uuid.uuid4().hex[:8] + self._lock = threading.Lock() # serialize socket use + self._spawn_lock = threading.Lock() + + # -- lifecycle --------------------------------------------------------- # + def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": + """Spawn the worker (opens + analyzes the DB) and connect once ready.""" + with self._spawn_lock: + if self._sock is not None: + return self + if self._proc is None or self._proc.poll() is not None: + self._proc = subprocess.Popen( + [self._python, "-m", "idatui.worker", + self._sock_path, self._bin], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + deadline = time.time() + timeout + t0 = time.time() + while time.time() < deadline: + try: + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.connect(self._sock_path) + self._sock = s + return self + except OSError: + if self._proc.poll() is not None: + raise IDAConnectionError( + f"worker exited during startup (code " + f"{self._proc.returncode})") + if progress: + progress(f"auto-analyzing {os.path.basename(self._bin)}… " + f"({int(time.time() - t0)}s)") + time.sleep(0.2) + raise IDAConnectionError("worker did not become ready in time") + + def close(self) -> None: + with self._lock: + s = self._sock + self._sock = None + if s is not None: + try: + _send(s, ("__shutdown__", {})) + except Exception: # noqa: BLE001 + pass + try: + s.close() + except Exception: # noqa: BLE001 + pass + if self._proc is not None: + try: + self._proc.terminate() + self._proc.wait(timeout=5) + except Exception: # noqa: BLE001 + try: + self._proc.kill() + except Exception: # noqa: BLE001 + pass + + # -- the call surface -------------------------------------------------- # + def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: + if self._sock is None: + self.connect() + with self._lock: + s = self._sock + if s is None: + raise IDAConnectionError("worker connection is closed") + try: + _send(s, (tool, args)) + reply = _recv(s) + except (OSError, ConnectionError) as e: + self._sock = None + raise IDAConnectionError(f"worker transport failed: {e}") from e + if reply is None: + self._sock = None + raise IDAConnectionError("worker closed the connection") + ok, payload = reply + if not ok: + raise IDAToolError(tool, str(payload)) + return payload + + def call_envelope(self, tool: str, *, timeout: float | None = None, + **args) -> dict: + # domain.decompile() reads result.structuredContent — mirror that shape. + return {"result": {"structuredContent": self.call(tool, timeout=timeout, + **args)}} + + # -- session shims (single-DB worker) --------------------------------- # + def set_db(self, db: str | None) -> None: + if db: + self._sid = db + + def resolve_db(self) -> str: + return self._sid + + def list_sessions(self) -> list[Session]: + return [Session(session_id=self._sid, + filename=os.path.basename(self._bin), + input_path=self._bin, is_active=True)] + + def health(self) -> dict: + try: + return self.call("server_health") + except IDAToolError: + return {"module": os.path.basename(self._bin), "ok": True} + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + return _NoopKeepAlive() + + # context manager parity with IDAClient + def __enter__(self) -> "WorkerClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() |
