diff --git a/benchmarks/independent/verify_sinks.py b/benchmarks/independent/verify_sinks.py index 13d90d4..31248ee 100644 --- a/benchmarks/independent/verify_sinks.py +++ b/benchmarks/independent/verify_sinks.py @@ -20,7 +20,6 @@ import concurrent.futures import json import shutil import sys -import threading import time from pathlib import Path @@ -28,7 +27,7 @@ HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) import evaluate # noqa: E402 -from model_sinks import OPERATIONS, Meter # noqa: E402 +from model_sinks import OPERATIONS # noqa: E402 from openultrasast.config import load_config # noqa: E402 from openultrasast.model.endpoint import DEFAULT_DETECTOR_MODEL # noqa: E402 @@ -45,15 +44,26 @@ SYSTEM = ( ) -def verify(root: Path, candidate: list, family: str, args: argparse.Namespace, meter: Meter, local: threading.local) -> list[dict]: +_CLIENT = None + + +def _client(): + global _CLIENT + if _CLIENT is None: + from openultrasast import tool_hunter + + load_config() + _CLIENT = tool_hunter.resolve_hunter_client() + return _CLIENT + + +def verify(root: str, candidate: list, family: str, model: str, max_steps: int) -> dict: + """One candidate, in a worker process: CPU-bound repository tools do not serialise behind one interpreter.""" from openultrasast import tool_hunter from openultrasast.complexity.map import Hotspot - if not hasattr(local, "client"): - local.client = tool_hunter.resolve_hunter_client() - meter.add(local.client) - if meter.spent() >= args.budget_usd: - return [{"unverified": "budget"}] + client = _client() + before = float(client.cost_usd()) path, function, line = candidate spot = Hotspot(path=path, function_name=function, score=1.0, band="candidate", signals={}, rationale="model-classified sink", test_hint=None, inventory_finding_ids=()) prompt = ( @@ -62,21 +72,16 @@ def verify(root: Path, candidate: list, family: str, args: argparse.Namespace, m ) try: found = tool_hunter.run_tool_hunter( - root, [spot], client=local.client, model=args.model, max_steps=args.max_steps, + Path(root), [spot], client=client, model=model, max_steps=max_steps, system_prompt=SYSTEM, user_prompt=prompt, context_files=[path], tags=("sink-verifier",), ) # fmt: skip except Exception as error: # noqa: BLE001 -- recorded; an unverified candidate is not a clean one - return [{"unverified": f"{type(error).__name__}: {str(error)[:120]}"}] - return [ - { - "site": f"{f.path}:{f.line or line}:{f.function_name or function}", - "family": family, - "title": f.title, - "witness": f.rationale[:600], - "candidate": f"{path}::{function}", - } + return {"candidate": candidate, "unverified": f"{type(error).__name__}: {str(error)[:120]}", "usd": float(client.cost_usd()) - before} + findings = [ + {"site": f"{f.path}:{f.line or line}:{f.function_name or function}", "family": family, "title": f.title, "witness": f.rationale[:600], "candidate": f"{path}::{function}"} for f in found ] + return {"candidate": candidate, "findings": findings, "usd": float(client.cost_usd()) - before} def main() -> int: @@ -85,7 +90,7 @@ def main() -> int: parser.add_argument("--only", default="") parser.add_argument("--model", default=DEFAULT_DETECTOR_MODEL) parser.add_argument("--max-steps", type=int, default=3) - parser.add_argument("--workers", type=int, default=8) + parser.add_argument("--workers", type=int, default=4) parser.add_argument("--budget-usd", type=float, default=40.0) parser.add_argument("--fixed", action="store_true", help="recheck detected sites at the fixed pin") args = parser.parse_args() @@ -95,7 +100,7 @@ def main() -> int: out_dir = evaluate.RESULTS.with_name(evaluate.RESULTS.name + "-verified") / "scans" out_dir.mkdir(parents=True, exist_ok=True) only = {c.strip() for c in args.only.split(",") if c.strip()} - meter, local = Meter(args.budget_usd), threading.local() + spent = 0.0 for case in evaluate.cases(): if only and case["id"] not in only: continue @@ -103,32 +108,52 @@ def main() -> int: out = out_dir / f"{case['id']}--{label}.json" if out.is_file(): continue + log = out_dir / f"{case['id']}--{label}.jsonl" + done = [json.loads(row) for row in log.read_text().splitlines()] if log.is_file() else [] + spent += sum(float(d.get("usd", 0.0)) for d in done) record = json.loads((candidates_dir / f"{case['id']}.json").read_text()) - todo = sorted({(p, fn, ln) for p, fn, ln in record["candidates"]}, key=lambda c: (c[0], c[1], c[2])) - todo = list({(p, fn): (p, fn, ln) for p, fn, ln in todo}.values()) # one question per function + todo = list({(p, fn): (p, fn, ln) for p, fn, ln in sorted(record["candidates"])}.values()) # one question per function if args.fixed: vulnerable = json.loads((out_dir / f"{case['id']}--vulnerable_a.json").read_text()) detected = {f["candidate"] for f in vulnerable["findings"]} todo = [c for c in todo if f"{c[0]}::{c[1]}" in detected] - started, before = time.monotonic(), meter.spent() + finished = {tuple(d["candidate"][:2]) for d in done} + pending = [c for c in todo if (c[0], c[1]) not in finished] + started = time.monotonic() checkout = evaluate.export(case, label, case["fixed" if args.fixed else "vulnerable"]) + stopped = False try: - with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool: - results = list(pool.map(lambda c: verify(checkout, list(c), case["family"], args, meter, local), todo)) + with concurrent.futures.ProcessPoolExecutor(max_workers=args.workers) as pool, log.open("a") as sink: + futures = [] + for candidate in pending: + futures.append(pool.submit(verify, str(checkout), list(candidate), case["family"], args.model, args.max_steps)) + for future in concurrent.futures.as_completed(futures): + result = future.result() + spent += float(result.get("usd", 0.0)) + sink.write(json.dumps(result) + "\n") + sink.flush() + done.append(result) + if spent >= args.budget_usd and not stopped: + stopped = True + for other in futures: + other.cancel() finally: shutil.rmtree(checkout, ignore_errors=True) - findings = [f for rows in results for f in rows if "site" in f] - unverified = [f["unverified"] for rows in results for f in rows if "unverified" in f] + verified = {tuple(d["candidate"][:2]) for d in done if "findings" in d} + findings = [f for d in done for f in d.get("findings", [])] + unverified = [d["unverified"] for d in done if "unverified" in d] + complete = len(verified) + len(unverified) >= len(todo) result = { - "root": "/case", "families": [case["family"]], "questions": len(todo), "completed": len(todo) - len(unverified), - "completed_regions": sorted({f"{p}:{fn}" for p, fn, _ in todo}), "unverified": unverified[:20], - "usd": round(meter.spent() - before, 4), "seconds": round(time.monotonic() - started, 1), + "root": "/case", "families": [case["family"]], "questions": len(todo), "completed": len(verified), + "completed_regions": sorted(f"{p}:{fn}" for p, fn in verified), "unverified": unverified[:20], + "usd": round(sum(float(d.get("usd", 0.0)) for d in done), 4), "seconds": round(time.monotonic() - started, 1), "findings": sorted(findings, key=lambda f: f["site"]), } # fmt: skip - out.write_text(json.dumps(result, indent=1) + "\n") - print(json.dumps({"case": case["id"], "pin": label, "candidates": len(todo), "findings": len(findings), "unverified": len(unverified), "usd": result["usd"], "total_usd": round(meter.spent(), 4)}), flush=True) - if meter.spent() >= args.budget_usd: - print(json.dumps({"stopped": "budget", "spent_usd": round(meter.spent(), 4)}), flush=True) + if complete: + out.write_text(json.dumps(result, indent=1) + "\n") + print(json.dumps({"case": case["id"], "pin": label, "candidates": len(todo), "verified": len(verified), "findings": len(findings), "unverified": len(unverified), "usd": result["usd"], "total_usd": round(spent, 4), "complete": complete}), flush=True) + if stopped or spent >= args.budget_usd: + print(json.dumps({"stopped": "budget", "spent_usd": round(spent, 4)}), flush=True) break return 0 diff --git a/src/openultrasast/hunter_tools.py b/src/openultrasast/hunter_tools.py index 4dd4ded..fd5152a 100644 --- a/src/openultrasast/hunter_tools.py +++ b/src/openultrasast/hunter_tools.py @@ -32,13 +32,45 @@ def read_file(root: Path, path: str, *, max_chars: int) -> str: return clamped.read_text(errors="ignore")[:max_chars] +_Lines = tuple[tuple[str, tuple[str, ...]], ...] +_SOURCE_MEMO: dict[tuple[str, int], tuple[_Lines, tuple[tuple[Path, int], ...]]] = {} + + +def _source_lines(root: Path) -> _Lines: + """Every product source file under `root` with its lines, read ONCE per root state. + + Each `grep_repo` and `find_refs` call re-enumerated the tree -- walk, ignore rules, vendored checks -- and + re-read every file. On a 1,800-file repository a model verifying candidates spent its hours there, CPU-bound, + not waiting on the model. Keyed by the root's path and modification time, and a hit is re-validated file by + file -- an edit inside a subdirectory does not touch the root's mtime. The value is immutable. + """ + resolved_root = root.resolve() + key = (str(resolved_root), resolved_root.stat().st_mtime_ns) + entry = _SOURCE_MEMO.get(key) + if entry is not None and all(_mtime(path) == stamp for path, stamp in entry[1]): + return entry[0] + files = list(_iter_clamped_source_files(resolved_root)) + lines = tuple((relative, tuple(path.read_text(errors="ignore").splitlines())) for relative, path in files) + if len(_SOURCE_MEMO) >= 4: + _SOURCE_MEMO.pop(next(iter(_SOURCE_MEMO))) + _SOURCE_MEMO[key] = (lines, tuple((path, _mtime(path)) for _, path in files)) + return lines + + +def _mtime(path: Path) -> int: + try: + return path.stat().st_mtime_ns + except OSError: + return -1 + + def grep_repo(root: Path, pattern: str, *, max_matches: int) -> list[dict[str, object]]: if max_matches <= 0: return [] compiled = re.compile(pattern) matches: list[dict[str, object]] = [] - for relative, path in _iter_clamped_source_files(root): - for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), start=1): + for relative, lines in _source_lines(root): + for line_no, line in enumerate(lines, start=1): if compiled.search(line): matches.append({"path": relative, "line": line_no, "text": line}) if len(matches) >= max_matches: @@ -67,8 +99,8 @@ def find_refs(root: Path, symbol: str, mapping_index: object) -> list[dict[str, ) if symbol: compiled = re.compile(rf"\b{re.escape(symbol)}\b") - for relative, path in _iter_clamped_source_files(resolved_root): - for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), start=1): + for relative, lines in _source_lines(resolved_root): + for line_no, line in enumerate(lines, start=1): if compiled.search(line): refs.append({"path": relative, "name": symbol, "line": line_no, "text": line, "source": "text"}) return refs diff --git a/tests/test_hunter_tools.py b/tests/test_hunter_tools.py index 72c9589..d446ba4 100644 --- a/tests/test_hunter_tools.py +++ b/tests/test_hunter_tools.py @@ -121,3 +121,25 @@ def test_find_refs_rejects_mapping_path_escape(tmp_path: Path) -> None: find_refs(root, "passwd", [{"path": "../etc/passwd", "name": "passwd"}]) with pytest.raises(PathEscapesRepo): find_refs(root, "passwd", [{"path": "/etc/passwd", "name": "passwd"}]) + + +def test_grep_reads_the_tree_once_and_sees_an_edit(tmp_path, monkeypatch) -> None: + """Every call re-enumerated and re-read the repository; now once per root state, re-validated per file.""" + import os + + from openultrasast import hunter_tools + + sub = tmp_path / "pkg" + sub.mkdir() + source = sub / "app.py" + source.write_text("def run(q):\n return query(q)\n") + calls = [] + real = hunter_tools._iter_clamped_source_files + monkeypatch.setattr(hunter_tools, "_iter_clamped_source_files", lambda root: calls.append(root) or real(root)) + assert [m["line"] for m in hunter_tools.grep_repo(tmp_path, r"query\(", max_matches=5)] == [2] + assert [m["line"] for m in hunter_tools.grep_repo(tmp_path, r"def run", max_matches=5)] == [1] + assert len(calls) == 1, "the tree was enumerated again for an unchanged root" + source.write_text("def run(q):\n x = 1\n return query(q)\n") + stamp = source.stat().st_mtime_ns + 10_000_000 + os.utime(source, ns=(stamp, stamp)) + assert [m["line"] for m in hunter_tools.grep_repo(tmp_path, r"query\(", max_matches=5)] == [3], "a stale read survived an edit"