"""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 .errors import IDAToolError, IDAConnectionError, Session from .worker import recv as _recv from .worker import send as _send _WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") _worker_python_cache: str | None = None def _find_worker_python() -> str: """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily the TUI's python. On a typical box the TUI runs under a venv that has textual + idalib but not ida_pro_mcp, while the system python has idalib + ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" global _worker_python_cache if _worker_python_cache: return _worker_python_cache override = os.environ.get("IDATUI_WORKER_PYTHON") candidates = [override] if override else [] candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] for py in candidates: if not py or not os.path.exists(py): continue try: r = subprocess.run([py, "-c", "import ida_pro_mcp"], capture_output=True, timeout=30) if r.returncode == 0: _worker_python_cache = py return py except Exception: # noqa: BLE001 continue return sys.executable # last resort; the worker will report the real error 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, load_args: str = "") -> None: self._bin = os.path.abspath(os.path.expanduser(binary_path)) self._load_args = load_args or "" # IDA switches for a headerless blob self._python = python or _find_worker_python() tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" self._sock_path = f"/tmp/idatui-worker-{tag}.sock" self._log_path = f"/tmp/idatui-worker-{tag}.log" 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() #: Set by close(). A dropped socket is respawned on the next call (the #: worker segfaulted and we want it back); a CLOSED one must not be. #: Teardown and binary-switch both close while @work threads are still #: in flight, so without this, quitting during a decompile spawned a #: fresh idalib worker that re-opened the database nobody was looking #: at any more -- a stray process holding the .i64 we just released. self._closed = False # -- 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: self._closed = False # an explicit reconnect revives this client if self._sock is not None: return self if self._proc is None or self._proc.poll() is not None: # run worker.py as a SCRIPT (not -m idatui.worker) so we don't # import the textual-dependent idatui package __init__ under the # IDA python, which usually has no textual. argv = [self._python, _WORKER_PY, self._sock_path, self._bin] if self._load_args: argv.append(self._load_args) self._proc = subprocess.Popen( argv, stdout=open(self._log_path, "wb"), stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, ) deadline = time.time() + timeout t0 = time.time() # Poll fast at first, then back off. A flat 0.2s cost every caller a # fifth of a second even when the worker was ready in milliseconds # (a small binary, or a seeded .i64), which is most of the time in # the tests and noticeable on a re-open. # # Backing off all the way to 0.2s was too eager: a seeded database # is ready at ~250ms, by which point the delay has grown to 134ms, so # every open waited ~350ms whatever the binary -- the same number for # a 47KB `echo` and a 1.2MB `bash`, which is what gives a polling # artefact away. Cap the backoff at 25ms instead: the overshoot on a # fast open is bounded by that, and 40 probes a second is nothing # next to an auto-analysis that runs for minutes. # # Do NOT be tempted to hold the 5ms rate instead. This poll runs on a # background thread while the UI thread is drawing, and 200 wakeups a # second through a cold analysis cost enough GIL time to delay the # app's own startup -- it left the loading overlay up long enough for # project mode's first keypress to land on it. delay = 0.005 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 (code {self._proc.returncode}): " f"{self._log_tail()} [full log: {self._log_path}]") if progress: progress(f"auto-analyzing {os.path.basename(self._bin)}… " f"({int(time.time() - t0)}s)") time.sleep(delay) delay = min(delay * 1.6, 0.025) raise IDAConnectionError("worker did not become ready in time") @property def pid(self) -> int | None: """The worker process id (for memory accounting), or None if not spawned.""" return self._proc.pid if self._proc is not None else None def close(self, grace: float = 20.0) -> None: """Shut the worker down cleanly. After ``__shutdown__`` the worker still has to ``close_database()``, which re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. Signalling it before that finishes is what leaves databases wedged, so wait out the grace period first and only escalate if it really is stuck. """ with self._lock: s = self._sock self._sock = None self._closed = True 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.wait(timeout=grace) # let it close the DB properly except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck 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._closed: raise IDAConnectionError( f"{tool}: this worker was closed (call connect() to revive it)") 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() def _log_tail(self, n: int = 400) -> str: """Last meaningful line(s) of the worker log (skip IDA's licence banner), so a startup crash surfaces the real cause instead of just 'code 1'.""" try: with open(self._log_path, encoding="utf-8", errors="replace") as f: lines = [ln.strip() for ln in f if ln.strip()] except OSError: return "(no worker log)" # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash for ln in reversed(lines): if ln.startswith("WORKER-FATAL:"): return ln[len("WORKER-FATAL:"):].strip()[-n:] skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") meaningful = [ln for ln in lines if not any(s in ln.lower() for s in skip)] return " | ".join((meaningful or lines)[-3:])[-n:] # context manager parity with IDAClient def __enter__(self) -> "WorkerClient": return self.connect() def __exit__(self, *exc) -> None: self.close()