From 93240e2f0c0d716758c43154dd37802f556ac6d0 Mon Sep 17 00:00:00 2001 From: user Date: Fri, 7 Aug 2026 01:31:57 +0200 Subject: Stop ida-pro-mcp installing a sys.setprofile hook around every tool call. Its deadline mechanism profiles every python call/return so a pure-python tool loop can be interrupted; our tools are call-heavy, so it taxed the whole backend 3.3x. Worker now sets IDA_MCP_TOOL_TIMEOUT_SEC=0 and arms the deadline itself with one polling watchdog thread + ida_kernwin.set_cancelled() (the half that actually frees the IDA main thread). Also rewrote _idatui_spans to jump between colour tags instead of walking characters (byte-identical over 258k real lines). Result: {"status":"keep","total_ms":26923.9,"lg_boot_ms":762.2,"lg_decomp_ms":2631.7,"lg_graph_ms":941.8,"lg_hex_ms":1052,"lg_index_ms":67.2,"lg_listing_cold_ms":440.6,"lg_listing_warm_ms":530.5,"lg_nav_ms":10598.3,"lg_palette_ms":4.6,"lg_render_ms":227.8,"lg_search_ms":5265.5,"pure_graph_ms":237.9,"sm_boot_ms":535.3,"sm_decomp_ms":631.8,"sm_graph_ms":702.1,"sm_hex_ms":841.1,"sm_index_ms":0,"sm_listing_cold_ms":268.2,"sm_listing_warm_ms":269.4,"sm_nav_ms":443.6,"sm_palette_ms":0.3,"sm_render_ms":277.8,"sm_search_ms":194.4,"fails":0} --- .auto/diff_spans.py | 106 +++++++++++++++++++++++++++++++++++++++++++++ .auto/log.jsonl | 3 ++ idatui/worker.py | 70 +++++++++++++++++++++++++++++- server/patch_server.py | 115 +++++++++++++++++++++++++++---------------------- 4 files changed, 241 insertions(+), 53 deletions(-) create mode 100644 .auto/diff_spans.py create mode 100644 .auto/log.jsonl diff --git a/.auto/diff_spans.py b/.auto/diff_spans.py new file mode 100644 index 0000000..95721a5 --- /dev/null +++ b/.auto/diff_spans.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Differential check: the current `_idatui_spans` vs the one at a git ref. + +The span walker turns IDA's colour-tagged disassembly line into (spans, ops). +It is on the hot path of every listing row, so it is worth optimising -- but its +output drives highlighting, operand marking and the cursor's column arithmetic, +so "faster" is only acceptable if it is byte-identical. + +This pulls both implementations out of `server/patch_server.py` (the current +working tree, and whatever `--ref` names), runs them over every tagged line of +a real binary, and reports the first disagreement. + + /usr/bin/python3 .auto/diff_spans.py [--ref HEAD] [--target targets/bash] + [--limit 60000] +""" +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def load_impl(path: str, name: str): + """Exec just the span-walker section of a patch_server.py's injected BODY. + + BODY is a normal (non-raw) triple-quoted string, so the escapes in it are + resolved once by importing the module -- slicing the file text instead would + compile the *undecoded* source and silently test a different program (an + escaped backslash became a literal one, and every 's' in a comment turned + into a space). + """ + import importlib.util + spec = importlib.util.spec_from_file_location(f"_ps_{name}", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) # IDA-free at import time + body = mod.BODY + a = body.index("#: IDA colour tag -> the semantic kind") + b = body.index("def _idatui_unknown_row") + g = {"__name__": name} + exec(compile(body[a:b], name, "exec"), g) # noqa: S102 + return g["_idatui_spans"] + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--ref", default="HEAD") + ap.add_argument("--target", default="targets/bash") + ap.add_argument("--limit", type=int, default=60000) + a = ap.parse_args() + + d = tempfile.mkdtemp(prefix="diffspans-") + old_path = os.path.join(d, "patch_server_old.py") + with open(old_path, "w") as fh: + fh.write(subprocess.run( + ["git", "-C", ROOT, "show", f"{a.ref}:server/patch_server.py"], + capture_output=True, text=True, check=True).stdout) + new = load_impl(os.path.join(ROOT, "server", "patch_server.py"), "new") + old = load_impl(old_path, "old") + + binary = os.path.join(ROOT, a.target) + tgt = os.path.join(d, os.path.basename(binary)) + shutil.copy2(binary, tgt) + shutil.copy2(binary + ".i64", tgt + ".i64") + try: + import idapro + if idapro.open_database(tgt, run_auto_analysis=False) != 0: + raise SystemExit("could not open database") + import ida_bytes + import ida_lines + import ida_segment + import idaapi + + checked = bad = 0 + for si in range(ida_segment.get_segm_qty()): + seg = ida_segment.getnseg(si) + if seg is None: + continue + ea = seg.start_ea + while ea < seg.end_ea and checked < a.limit: + line = ida_lines.generate_disasm_line(ea, 0) + if line: + checked += 1 + ra, rb = old(line), new(line) + if ra != rb: + bad += 1 + if bad <= 3: + print(f"MISMATCH @ {ea:#x}\n line={line!r}\n" + f" old={ra!r}\n new={rb!r}") + nxt = ida_bytes.get_item_end(ea) + ea = nxt if nxt > ea else ea + 1 + if checked >= a.limit: + break + print(f"checked {checked} lines, {bad} mismatches") + idapro.close_database(False) + return 1 if bad else 0 + finally: + shutil.rmtree(d, ignore_errors=True) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.auto/log.jsonl b/.auto/log.jsonl new file mode 100644 index 0000000..ebef486 --- /dev/null +++ b/.auto/log.jsonl @@ -0,0 +1,3 @@ +{"type":"config","name":"ida-tui performance: cut the latency of navigation, listing, decomp, graph and search","metricName":"total_ms","metricUnit":"ms","bestDirection":"lower"} +{"run":1,"commit":"2910c93","metric":46572.1,"metrics":{"lg_boot_ms":863.5,"lg_decomp_ms":2881.3,"lg_graph_ms":951.7,"lg_hex_ms":900.9,"lg_index_ms":219.3,"lg_listing_cold_ms":434.4,"lg_listing_warm_ms":552.3,"lg_nav_ms":29106.9,"lg_palette_ms":4.8,"lg_render_ms":231.2,"lg_search_ms":5590.1,"pure_graph_ms":236.2,"sm_boot_ms":539.4,"sm_decomp_ms":621.4,"sm_graph_ms":702,"sm_hex_ms":850.5,"sm_index_ms":0,"sm_listing_cold_ms":270.8,"sm_listing_warm_ms":270.6,"sm_nav_ms":875.1,"sm_palette_ms":0.3,"sm_render_ms":272.3,"sm_search_ms":197.1,"fails":0},"status":"checks_failed","description":"Baseline run of the new bench harness. Benchmark clean (fails=0) but the pilot scenario suite reported 300 passed / 1 failed with ZERO code changes -> flaky, and checks.sh printed the tail instead of the FAIL line so the name is unknown.","timestamp":1786058161497,"segment":0,"confidence":null,"asi":{"hypothesis":"establish a baseline for total_ms","bottleneck":"lg_nav_ms=29107 is 62% of total_ms; lg_nav_worst_ms=28274 is ONE cold jump to a high address in bash. ListingModel.ensure_ea walks the segment forward in 500-head pages from seg_start, so landing near the end of a 224k-row listing costs ~450 sequential worker round trips.","second_bottleneck":"lg_search_ms=5590 (91783 hits over the whole segment)","cheap_phases":"palette/index/render/pure_graph are all <600ms; not where the time is","rollback_reason":"checks.sh flagged 1 scenario failure with no code change (flake)","next_action_hint":"make checks.sh print the FAIL line name on non-zero exit, re-run baseline, then attack ListingModel address->row lookup (needs a backend primitive in server/patch_server.py: heads walking anchored at an address, or a segment head-index built in one call)"}} +{"run":2,"commit":"a3f3400","metric":46685.5,"metrics":{"lg_boot_ms":885.7,"lg_decomp_ms":2871.8,"lg_graph_ms":930.9,"lg_hex_ms":947.3,"lg_index_ms":207.8,"lg_listing_cold_ms":564.3,"lg_listing_warm_ms":454.7,"lg_nav_ms":29018.5,"lg_palette_ms":4.6,"lg_render_ms":223.6,"lg_search_ms":5448.7,"pure_graph_ms":535.4,"sm_boot_ms":537.9,"sm_decomp_ms":643.8,"sm_graph_ms":680.4,"sm_hex_ms":854.9,"sm_index_ms":0,"sm_listing_cold_ms":265.5,"sm_listing_warm_ms":263.7,"sm_nav_ms":887.2,"sm_palette_ms":0.3,"sm_render_ms":264.1,"sm_search_ms":194.5,"fails":0},"status":"discard","description":"Baseline re-run with the fixed checks gate. Checks pass; total_ms reproduces to within 0.24% of run #1 (46572 -> 46686), so the noise floor is ~115ms on a 46.6s metric.","timestamp":1786058370297,"segment":0,"confidence":null,"asi":{"hypothesis":"confirm the baseline is reproducible and the checks gate is green","noise_floor_ms":115,"reproducibility":"run1 46572 / run2 46686 -> 0.24% spread; pure_graph_ms is the jumpiest single phase (236 -> 535, it is CPU-only and gets descheduled)","checks":"flaky-scenario retry logic works; suite green on a clean tree","next_action_hint":"attack ListingModel.ensure_ea / the heads tool: 29s of 46.6s is one cold address->row walk over bash"}} diff --git a/idatui/worker.py b/idatui/worker.py index a4e3509..556e69a 100644 --- a/idatui/worker.py +++ b/idatui/worker.py @@ -26,8 +26,68 @@ 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 @@ -147,6 +207,7 @@ def _open_and_register(binpath: str, load_args: str = ""): 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) @@ -167,7 +228,14 @@ def serve(sockpath: str, binpath: str, load_args: str = "") -> None: fn = tools.get(name) if fn is None: raise KeyError(f"unknown tool: {name!r}") - result = fn(**args) + 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"]). diff --git a/server/patch_server.py b/server/patch_server.py index 73147d0..b806be2 100644 --- a/server/patch_server.py +++ b/server/patch_server.py @@ -364,6 +364,8 @@ def _idatui_tag_map(): _IDATUI_TAGS = None _IDATUI_OPND_TAGS = None +_IDATUI_CTL = None # re: the three control characters a tagged line can hold +_IDATUI_WS = None # re: a run of whitespace, exactly what str.isspace() calls one def _idatui_opnd_tag_map(): @@ -393,75 +395,84 @@ def _idatui_spans(line): Unknown tags become 'text' rather than being dropped: a processor module can emit a colour we don't classify, and losing the characters would corrupt the line.""" - global _IDATUI_TAGS, _IDATUI_OPND_TAGS + global _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL, _IDATUI_WS import ida_lines if _IDATUI_TAGS is None: _IDATUI_TAGS = _idatui_tag_map() if _IDATUI_OPND_TAGS is None: _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() + if _IDATUI_CTL is None: + import re as _re + _IDATUI_CTL = _re.compile("[\\x01\\x02\\x03]") + # str.isspace() is true for \\x1c-\\x1f and \\x85 as well as the \\s + # class, so spell those out: this substitution has to agree with the + # plain-text collapse character for character (checked over every + # codepoint) or the row silently loses its highlighting. + _IDATUI_WS = _re.compile("[\\\\s\\x1c\\x1d\\x1e\\x1f\\x85]+") + tags, opnds, ctl = _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL on, off, esc = "\x01", "\x02", "\x03" addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) + # Jump between control characters and take the text in between as one slice. + # A per-character loop here was 68% of the whole `heads` tool: a disasm line + # is ~50 characters but only ~15 tags, and everything between two tags is + # already exactly one span's worth of text. spans, stack, buf = [], [], [] # stack entries: (kind, operand index|None) + kind, opnd = "text", None # state the current run of text belongs to i, n = 0, len(line) - - def _opnd(): - for _k, o in reversed(stack): - if o is not None: - return o - return None - - def flush(): - if buf: - spans.append([stack[-1][0] if stack else "text", "".join(buf), - _opnd()]) - del buf[:] - - while i < n: - ch = line[i] - if ch == on and i + 1 < n: - tag = line[i + 1] - if tag == addr_tag: - # An embedded target address, not display text: 16 hex digits - # that must not reach the screen. - i += 2 + addr_len - continue - flush() - stack.append((_IDATUI_TAGS.get(tag, "text"), - _IDATUI_OPND_TAGS.get(tag))) - i += 2 + for m in ctl.finditer(line): + j = m.start() + if j < i: # inside an address payload / after an esc continue - if ch == off and i + 1 < n: - flush() - if stack: - stack.pop() - i += 2 + if j + 1 >= n: # a trailing control char is literal text + break + ch = line[j] + if ch == esc: # escaped literal: keep the char it guards + buf.append(line[i:j]) + buf.append(line[j + 1]) + i = j + 2 continue - if ch == esc and i + 1 < n: # escaped literal - buf.append(line[i + 1]) - i += 2 + tag = line[j + 1] + if ch == on and tag == addr_tag: + # An embedded target address, not display text: 16 hex digits that + # must not reach the screen. Deliberately NOT a span boundary. + buf.append(line[i:j]) + i = j + 2 + addr_len continue - buf.append(ch) - i += 1 - flush() + buf.append(line[i:j]) + i = j + 2 + txt = "".join(buf) + if txt: + spans.append([kind, txt, opnd]) + del buf[:] + if ch == on: + stack.append((kind, opnd)) + kind = tags.get(tag, "text") + o = opnds.get(tag) + if o is not None: + opnd = o # operands nest: an inner colour keeps the operand + elif stack: + kind, opnd = stack.pop() + else: + kind, opnd = "text", None + if i < n: + buf.append(line[i:]) + txt = "".join(buf) + if txt: + spans.append([kind, txt, opnd]) # Collapse IDA's column padding EXACTLY as the plain text does. A run of - # spaces can straddle two spans, so this walks characters rather than - # collapsing each span on its own — otherwise the spans and `text` disagree - # about the line and the row silently loses its highlighting. + # spaces can straddle two spans, so the leading space of a span is dropped + # when the previous one ended in space — otherwise the spans and `text` + # disagree about the line and the row silently loses its highlighting. out, prev_space = [], False + ws = _IDATUI_WS for kind, txt, opnd in spans: - acc = [] - for ch in txt: - if ch.isspace(): - if prev_space: - continue - acc.append(" ") - prev_space = True - else: - acc.append(ch) - prev_space = False + acc = ws.sub(" ", txt) + if prev_space and acc[:1] == " ": + acc = acc[1:] if acc: - out.append([kind, "".join(acc), opnd]) + prev_space = acc[-1] == " " + out.append([kind, acc, opnd]) while out and out[0][1] == " ": out.pop(0) while out and out[-1][1] == " ": -- cgit v1.3.1-sl0p