aboutsummaryrefslogtreecommitdiffstats
path: root/idatui
diff options
context:
space:
mode:
authorblasty <blasty@local>2026-07-24 01:29:06 +0200
committerblasty <blasty@local>2026-07-24 01:29:06 +0200
commitb14f26f0bb79baf9e1476f2d67e2cddd8b6957d1 (patch)
tree4b32565a0c3f8eb9d806006aecdb24fd4e6e2167 /idatui
parentexperiments: add a unix-socket idalib worker as the 3rd bench column (diff)
downloadida-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')
-rw-r--r--idatui/worker.py160
-rw-r--r--idatui/worker_client.py165
2 files changed, 325 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()
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()