1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
"""Profile an operation INSIDE the database process.
`bench_ops.py` says how long an operation takes; this says where that time
goes. The snippet ships cProfile into the Code Mode sandbox, runs the real
remote-library function there in a loop, and returns the stats as text -- so
the split between IDA's own calls and OUR python in `remote_tools.py` is
visible, which no client-side timer can see.
PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py [BINARY]
PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py --op decompile
Read the `tottime` column: time in that function excluding subcalls. IDA
builtins (generate_disasm_line, get_flags, next_head...) are the floor; a
python frame from ida_tui_remote near the top is ours, and ours is fixable.
"""
from __future__ import annotations
import argparse
import os
import sys
from idatui.codemode_client import CodeModeClient, _REMOTE_MODULE, _script
# Runs in the database process. `a` is the bound argument dict.
PROFILE = '''
import cProfile, pstats, io, sys
_m = sys.modules.get(%(mod)r)
if _m is None:
result = {"error": "remote lib not installed yet"}
else:
call = a["call"]
reps = int(a["reps"])
ns = {"_m": _m, "a": a}
src = "for _ in range(%%d):\\n _m.%%s" %% (reps, call)
code = compile(src, "<profile>", "exec")
pr = cProfile.Profile()
pr.enable()
exec(code, ns)
pr.disable()
buf = io.StringIO()
st = pstats.Stats(pr, stream=buf).sort_stats("tottime")
st.print_stats(int(a["lines"]))
result = {"stats": buf.getvalue(), "total": st.total_tt, "reps": reps}
''' % {"mod": _REMOTE_MODULE}
CALLS = {
# One full listing page, exactly as the background grower asks for it.
"heads": 'heads(addr=a["addr"], count=500, annotate=True)',
"heads_plain": 'heads(addr=a["addr"], count=500, annotate=False)',
"decompile": 'decompile(a["addr"])',
"disasm": 'disasm(a["addr"], 500)',
}
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("binary", nargs="?", default="targets/bash")
ap.add_argument("--op", default="heads", choices=sorted(CALLS))
ap.add_argument("--reps", type=int, default=20)
ap.add_argument("--lines", type=int, default=18)
ap.add_argument("--addr", default=None, help="default: the .text start")
args = ap.parse_args()
client = CodeModeClient(os.path.abspath(args.binary))
client.connect()
addr = args.addr
if addr is None:
regions = client.invoke("file_regions")
rows = regions.get("regions") or regions.get("result") or []
text = next((r for r in rows if ".text" in str(r.get("name", ""))), None)
addr = (text or rows[0])["start"] if rows else "0x0"
print(f"# {os.path.basename(args.binary)} op={args.op} addr={addr} reps={args.reps}")
# Prime: the remote lib installs lazily, and its lru_caches must be warm or
# the profile measures cache misses that the real workload never pays.
client.invoke("heads", addr=addr, count=500, annotate=True)
# _script binds the args as JSON and adds the pack epilogue, exactly as a
# real operation is shipped -- so this measures the same path, not a
# special one.
out = client._unpack(client.execute_python(_script(
{"call": CALLS[args.op], "addr": addr,
"reps": args.reps, "lines": args.lines}, PROFILE), timeout=600))
if "error" in out:
print("FAILED:", out["error"])
return 1
per = out["total"] / out["reps"] * 1000
print(f"# {out['total']*1000:.0f}ms total, {per:.1f}ms per call\n")
print(out["stats"])
return 0
if __name__ == "__main__":
raise SystemExit(main())
|
