mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
perf(hunter): read a repository's source once per state, not once per tool call
Each grep_repo and find_refs call re-enumerated the tree (walk, ignore rules, vendored checks) and re-read every file. The v2 sink verifier spent almost three hours on its first case (pgAdmin, about 1,800 files) at 108% CPU with one open connection. It was CPU-bound in these tools, not waiting on the model. Source lines are now memoised per root, keyed by path and modification time. A hit is re-validated file by file, because an edit inside a subdirectory does not touch the root's mtime; a test covers one enumeration and a detected edit. verify_sinks.py now runs candidates in worker processes, not threads, and appends each candidate's result to a per-case log, so a stopped run resumes where it left off. GoPay: 12 candidates in 14 s, the real site found. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
023c24535f
commit
267bca5232
@@ -20,7 +20,6 @@ import concurrent.futures
|
|||||||
import json
|
import json
|
||||||
import shutil
|
import shutil
|
||||||
import sys
|
import sys
|
||||||
import threading
|
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -28,7 +27,7 @@ HERE = Path(__file__).resolve().parent
|
|||||||
sys.path.insert(0, str(HERE))
|
sys.path.insert(0, str(HERE))
|
||||||
|
|
||||||
import evaluate # noqa: E402
|
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.config import load_config # noqa: E402
|
||||||
from openultrasast.model.endpoint import DEFAULT_DETECTOR_MODEL # 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 import tool_hunter
|
||||||
from openultrasast.complexity.map import Hotspot
|
from openultrasast.complexity.map import Hotspot
|
||||||
|
|
||||||
if not hasattr(local, "client"):
|
client = _client()
|
||||||
local.client = tool_hunter.resolve_hunter_client()
|
before = float(client.cost_usd())
|
||||||
meter.add(local.client)
|
|
||||||
if meter.spent() >= args.budget_usd:
|
|
||||||
return [{"unverified": "budget"}]
|
|
||||||
path, function, line = candidate
|
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=())
|
spot = Hotspot(path=path, function_name=function, score=1.0, band="candidate", signals={}, rationale="model-classified sink", test_hint=None, inventory_finding_ids=())
|
||||||
prompt = (
|
prompt = (
|
||||||
@@ -62,21 +72,16 @@ def verify(root: Path, candidate: list, family: str, args: argparse.Namespace, m
|
|||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
found = tool_hunter.run_tool_hunter(
|
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",),
|
system_prompt=SYSTEM, user_prompt=prompt, context_files=[path], tags=("sink-verifier",),
|
||||||
) # fmt: skip
|
) # fmt: skip
|
||||||
except Exception as error: # noqa: BLE001 -- recorded; an unverified candidate is not a clean one
|
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 {"candidate": candidate, "unverified": f"{type(error).__name__}: {str(error)[:120]}", "usd": float(client.cost_usd()) - before}
|
||||||
return [
|
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}"}
|
||||||
"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
|
for f in found
|
||||||
]
|
]
|
||||||
|
return {"candidate": candidate, "findings": findings, "usd": float(client.cost_usd()) - before}
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
@@ -85,7 +90,7 @@ def main() -> int:
|
|||||||
parser.add_argument("--only", default="")
|
parser.add_argument("--only", default="")
|
||||||
parser.add_argument("--model", default=DEFAULT_DETECTOR_MODEL)
|
parser.add_argument("--model", default=DEFAULT_DETECTOR_MODEL)
|
||||||
parser.add_argument("--max-steps", type=int, default=3)
|
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("--budget-usd", type=float, default=40.0)
|
||||||
parser.add_argument("--fixed", action="store_true", help="recheck detected sites at the fixed pin")
|
parser.add_argument("--fixed", action="store_true", help="recheck detected sites at the fixed pin")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
@@ -95,7 +100,7 @@ def main() -> int:
|
|||||||
out_dir = evaluate.RESULTS.with_name(evaluate.RESULTS.name + "-verified") / "scans"
|
out_dir = evaluate.RESULTS.with_name(evaluate.RESULTS.name + "-verified") / "scans"
|
||||||
out_dir.mkdir(parents=True, exist_ok=True)
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
only = {c.strip() for c in args.only.split(",") if c.strip()}
|
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():
|
for case in evaluate.cases():
|
||||||
if only and case["id"] not in only:
|
if only and case["id"] not in only:
|
||||||
continue
|
continue
|
||||||
@@ -103,32 +108,52 @@ def main() -> int:
|
|||||||
out = out_dir / f"{case['id']}--{label}.json"
|
out = out_dir / f"{case['id']}--{label}.json"
|
||||||
if out.is_file():
|
if out.is_file():
|
||||||
continue
|
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())
|
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 sorted(record["candidates"])}.values()) # one question per function
|
||||||
todo = list({(p, fn): (p, fn, ln) for p, fn, ln in todo}.values()) # one question per function
|
|
||||||
if args.fixed:
|
if args.fixed:
|
||||||
vulnerable = json.loads((out_dir / f"{case['id']}--vulnerable_a.json").read_text())
|
vulnerable = json.loads((out_dir / f"{case['id']}--vulnerable_a.json").read_text())
|
||||||
detected = {f["candidate"] for f in vulnerable["findings"]}
|
detected = {f["candidate"] for f in vulnerable["findings"]}
|
||||||
todo = [c for c in todo if f"{c[0]}::{c[1]}" in detected]
|
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"])
|
checkout = evaluate.export(case, label, case["fixed" if args.fixed else "vulnerable"])
|
||||||
|
stopped = False
|
||||||
try:
|
try:
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
|
with concurrent.futures.ProcessPoolExecutor(max_workers=args.workers) as pool, log.open("a") as sink:
|
||||||
results = list(pool.map(lambda c: verify(checkout, list(c), case["family"], args, meter, local), todo))
|
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:
|
finally:
|
||||||
shutil.rmtree(checkout, ignore_errors=True)
|
shutil.rmtree(checkout, ignore_errors=True)
|
||||||
findings = [f for rows in results for f in rows if "site" in f]
|
verified = {tuple(d["candidate"][:2]) for d in done if "findings" in d}
|
||||||
unverified = [f["unverified"] for rows in results for f in rows if "unverified" in f]
|
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 = {
|
result = {
|
||||||
"root": "/case", "families": [case["family"]], "questions": len(todo), "completed": len(todo) - len(unverified),
|
"root": "/case", "families": [case["family"]], "questions": len(todo), "completed": len(verified),
|
||||||
"completed_regions": sorted({f"{p}:{fn}" for p, fn, _ in todo}), "unverified": unverified[:20],
|
"completed_regions": sorted(f"{p}:{fn}" for p, fn in verified), "unverified": unverified[:20],
|
||||||
"usd": round(meter.spent() - before, 4), "seconds": round(time.monotonic() - started, 1),
|
"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"]),
|
"findings": sorted(findings, key=lambda f: f["site"]),
|
||||||
} # fmt: skip
|
} # fmt: skip
|
||||||
out.write_text(json.dumps(result, indent=1) + "\n")
|
if complete:
|
||||||
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)
|
out.write_text(json.dumps(result, indent=1) + "\n")
|
||||||
if meter.spent() >= args.budget_usd:
|
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)
|
||||||
print(json.dumps({"stopped": "budget", "spent_usd": round(meter.spent(), 4)}), flush=True)
|
if stopped or spent >= args.budget_usd:
|
||||||
|
print(json.dumps({"stopped": "budget", "spent_usd": round(spent, 4)}), flush=True)
|
||||||
break
|
break
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|||||||
@@ -32,13 +32,45 @@ def read_file(root: Path, path: str, *, max_chars: int) -> str:
|
|||||||
return clamped.read_text(errors="ignore")[:max_chars]
|
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]]:
|
def grep_repo(root: Path, pattern: str, *, max_matches: int) -> list[dict[str, object]]:
|
||||||
if max_matches <= 0:
|
if max_matches <= 0:
|
||||||
return []
|
return []
|
||||||
compiled = re.compile(pattern)
|
compiled = re.compile(pattern)
|
||||||
matches: list[dict[str, object]] = []
|
matches: list[dict[str, object]] = []
|
||||||
for relative, path in _iter_clamped_source_files(root):
|
for relative, lines in _source_lines(root):
|
||||||
for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), start=1):
|
for line_no, line in enumerate(lines, start=1):
|
||||||
if compiled.search(line):
|
if compiled.search(line):
|
||||||
matches.append({"path": relative, "line": line_no, "text": line})
|
matches.append({"path": relative, "line": line_no, "text": line})
|
||||||
if len(matches) >= max_matches:
|
if len(matches) >= max_matches:
|
||||||
@@ -67,8 +99,8 @@ def find_refs(root: Path, symbol: str, mapping_index: object) -> list[dict[str,
|
|||||||
)
|
)
|
||||||
if symbol:
|
if symbol:
|
||||||
compiled = re.compile(rf"\b{re.escape(symbol)}\b")
|
compiled = re.compile(rf"\b{re.escape(symbol)}\b")
|
||||||
for relative, path in _iter_clamped_source_files(resolved_root):
|
for relative, lines in _source_lines(resolved_root):
|
||||||
for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), start=1):
|
for line_no, line in enumerate(lines, start=1):
|
||||||
if compiled.search(line):
|
if compiled.search(line):
|
||||||
refs.append({"path": relative, "name": symbol, "line": line_no, "text": line, "source": "text"})
|
refs.append({"path": relative, "name": symbol, "line": line_no, "text": line, "source": "text"})
|
||||||
return refs
|
return refs
|
||||||
|
|||||||
@@ -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"}])
|
find_refs(root, "passwd", [{"path": "../etc/passwd", "name": "passwd"}])
|
||||||
with pytest.raises(PathEscapesRepo):
|
with pytest.raises(PathEscapesRepo):
|
||||||
find_refs(root, "passwd", [{"path": "/etc/passwd", "name": "passwd"}])
|
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"
|
||||||
|
|||||||
Reference in New Issue
Block a user