aboutsummaryrefslogtreecommitdiffstats
path: root/tests/test_domain.py
blob: 24cb781fcec282e7ce7b54ed8329dcd0c83b81ca (plain) (blame)
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
#!/usr/bin/env python3
"""Validate the domain/paging layer against a real, large binary.

Run with a session open on a big binary (e.g. libcrypto.so.3, 10k funcs):

    python3 tests/test_domain.py --db <session_id>

Checks: correct full pagination (advance by len, not next_offset), viewport
slicing across block boundaries, window caching (revisit is instant), prefetch
warming, cached instruction totals, decompile success + hard-failure handling,
and address resolution.
"""
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from idatui.client import IDAClient  # noqa: E402
from idatui.domain import DISASM_BLOCK, LIST_PAGE, Program  # noqa: E402

URL = "http://127.0.0.1:8745/mcp"
PASS = FAIL = 0


def check(name, cond, detail=""):
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"  ok   {name}")
    else:
        FAIL += 1
        print(f"  FAIL {name}   {detail}")


def ms(fn):
    t = time.time()
    r = fn()
    return (time.time() - t) * 1e3, r


def main(argv):
    db = None
    it = iter(argv)
    for a in it:
        if a == "--db":
            db = next(it)
    c = IDAClient(URL, db=db)
    c.connect()
    if db is None:
        c.set_db(c.resolve_db())
    prog = Program(c)
    module = c.health().get("module")
    total_funcs = c.call("survey_binary").get("statistics", {}).get("total_functions")
    print(f"db={c.db} module={module} survey_total_functions={total_funcs}")

    # ---- function index: full enumeration correctness ------------------ #
    print("\n[function index]")
    idx = prog.functions()
    dt, _ = ms(lambda: idx.load_all())
    n = len(idx)
    check("enumerated all functions (matches survey)", n == total_funcs,
          f"got {n} vs survey {total_funcs}")
    print(f"    loaded {n} funcs in {dt:.0f}ms ({dt / max(n,1):.3f}ms/func), "
          f"page={LIST_PAGE}")
    # no duplicates, monotonic-ish uniqueness by addr
    addrs = [idx.get(i).addr for i in range(min(n, 3000))]
    check("no duplicate addrs in first 3000", len(addrs) == len(set(addrs)))
    # viewport slice (adapt to binary size)
    wstart = min(1000, max(n - 50, 0))
    wlen = min(50, n - wstart)
    w = idx.window(wstart, wlen)
    check(f"window({wstart},{wlen}) returns {wlen}", len(w) == wlen, str(len(w)))

    # ---- filtered index uses server-side glob -------------------------- #
    print("\n[filtered index]")
    sub = prog.functions(filter="sub_*")
    sub.ensure(10)
    check("filter sub_* yields sub_ names", all(f.name.startswith("sub_") for f in sub.window(0, 10)),
          str([f.name for f in sub.window(0, 5)]))

    # ---- pick the fattest function for disasm stress ------------------- #
    print("\n[disasm windowing]")
    fattest = max((idx.get(i) for i in range(n)), key=lambda f: f.size)
    dm = prog.disasm(fattest.addr, fattest.name)
    dt, total = ms(dm.total)
    check("total() returns positive count", total > 0, str(total))
    dt2, total2 = ms(dm.total)
    check("total() cached (2nd call ~instant)", dt2 < dt / 2 + 1, f"{dt:.0f}ms -> {dt2:.1f}ms")
    print(f"    fattest {fattest.name}: {total} insns, total() {dt:.0f}ms then {dt2:.1f}ms")

    # viewport across a block boundary
    start = DISASM_BLOCK - 5
    win = dm.lines(start, 60, prefetch=False)
    check("viewport spans block boundary, right length",
          len(win) == min(60, max(total - start, 0)), f"got {len(win)}")
    # addresses strictly increasing and contiguous slice
    eas = [ln.ea for ln in win]
    check("viewport addrs strictly increasing", all(b > a for a, b in zip(eas, eas[1:])),
          str(eas[:4]))

    # deep window: first slow (O(offset)), revisit instant (cached)
    if total > DISASM_BLOCK * 4:
        deep = (total // DISASM_BLOCK - 1) * DISASM_BLOCK
        dt_cold, a = ms(lambda: dm.lines(deep, 60, prefetch=False))
        dt_warm, b = ms(lambda: dm.lines(deep, 60, prefetch=False))
        check("deep window revisit is cached/instant", dt_warm < dt_cold / 2 + 1,
              f"cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms")
        check("cached window identical", [l.ea for l in a] == [l.ea for l in b])
        print(f"    deep@{deep}: cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms")

    # prefetch warms the next block
    print("\n[prefetch]")
    dm2 = prog.disasm(fattest.addr + 0)  # same model (cached by ea)
    fresh = prog.disasm(idx.get(0).addr, idx.get(0).name)
    fresh.lines(0, 60, prefetch=True)  # should prefetch block 1
    time.sleep(0.3)
    check("prefetch warmed a neighbor block", fresh.cached_blocks() >= 2,
          f"cached_blocks={fresh.cached_blocks()}")

    # ---- decompile: success + hard-failure ----------------------------- #
    print("\n[decompile]")
    # a small function likely decompiles
    small = min((idx.get(i) for i in range(n)), key=lambda f: f.size if f.size > 4 else 1 << 30)
    d_small = prog.decompile(small.addr)
    check("small func decompiles or fails cleanly", isinstance(d_small.failed, bool))
    dt_c, _ = ms(lambda: prog.decompile(small.addr))
    check("decompile cached (2nd ~instant)", dt_c < 5, f"{dt_c:.1f}ms")
    # the monster should hard-fail as a soft error (not raise)
    d_big = prog.decompile(fattest.addr)
    check("monster decompile handled (failed flag, no raise)",
          d_big.failed or d_big.code is not None,
          f"failed={d_big.failed} err={d_big.error}")
    print(f"    small {small.name}: failed={d_small.failed} "
          f"trunc={d_small.truncated} chars={d_small.total_chars}")
    print(f"    monster {fattest.name}: failed={d_big.failed} err={d_big.error}")

    # ---- resolve ------------------------------------------------------- #
    print("\n[resolve]")
    check("resolve hex", prog.resolve(hex(fattest.addr)) == fattest.addr)
    check("resolve int passthrough", prog.resolve(fattest.addr) == fattest.addr)
    named = next((idx.get(i) for i in range(n) if not idx.get(i).name.startswith("sub_")), None)
    if named:
        try:
            r = prog.resolve(named.name)
            check("resolve symbol name", r == named.addr, f"{hex(r)} vs {hex(named.addr)} ({named.name})")
        except KeyError as e:
            check("resolve symbol name", False, str(e))

    # ---- flat listing (code + data + undefined heads) ------------------ #
    print("\n[listing]")
    seg = prog.segment_bounds(fattest.addr)
    check("segment_bounds finds the .text segment", seg is not None and seg[0] <= fattest.addr < seg[1],
          str(seg))
    lm = prog.listing(fattest.addr)
    check("listing() returns a model for a mapped address", lm is not None)
    if lm is not None:
        lm.ensure(20)
        w = lm.window(0, 20)
        check("listing window returns heads", len(w) == 20, str(len(w)))
        eas = [h.ea for h in w]
        check("listing head addrs strictly increasing", all(b > a for a, b in zip(eas, eas[1:])),
              str(eas[:4]))
        check("listing heads carry a kind", all(h.kind in ("code", "data", "unknown") for h in w),
              str({h.kind for h in w}))
        check("listing head sizes positive", all(h.size >= 1 for h in w),
              str([h.size for h in w[:6]]))
        # random access to a mid-segment address lands on the containing head
        mid = w[10].ea
        li = lm.ensure_ea(mid)
        check("ensure_ea lands on the exact head for a head address",
              li >= 0 and lm.get(li).ea == mid, f"idx={li}")
        # a mid-item byte resolves to its containing head
        if w[4].size > 1:
            inside = w[4].ea + 1
            li2 = lm.ensure_ea(inside)
            check("ensure_ea resolves a mid-item byte to its head",
                  li2 >= 0 and lm.get(li2).ea == w[4].ea, f"idx={li2} ea={w[4].ea:#x}")
        dt_cold, _ = ms(lambda: lm.window(0, 20))
        check("listing window revisit is cached/instant", dt_cold < 5, f"{dt_cold:.1f}ms")
        dt_all, _ = ms(lambda: lm.load_all())
        check("listing load_all completes the segment", lm.complete and len(lm) > 20,
              f"n={len(lm)} complete={lm.complete}")
        print(f"    {seg[2]}: {len(lm)} heads walked in {dt_all:.0f}ms")

    prog.close()
    print(f"\n{PASS} passed, {FAIL} failed")
    return 1 if FAIL else 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))