learn: trace-enabled Joern run over labelled units (opt-in trace in taint.sc, trace parsing, resumable runner)

Implemented with Codex. Default scan queries do not request a trace.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

tests/test_engine_trace.py

@@ -0,0 +1,355 @@
+"""Offline controls for the trace experiment; no Docker, Joern or network."""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from openultrasast.cpg.backend import BEGIN, END, CpgResult, JoernBackend
+from openultrasast.model.specs import taint_specs
+from openultrasast.model.taint import request_params
+from openultrasast.model.trace import parse_trace
+
+SCRIPT = Path(__file__).resolve().parents[1] / "benchmarks/learn/engine_trace.py"
+_spec = importlib.util.spec_from_file_location("engine_trace", SCRIPT)
+assert _spec and _spec.loader
+trace = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(trace)
+
+
+def row(**extra):
+    return {
+        "sink": "os.system(cmd)",
+        "source": "request.args['cmd']",
+        "sinkFile": "app.py",
+        "sinkMethod": "run",
+        "sinkLine": 4,
+        "sourceKind": "framework",
+        "length": 3,
+        "paths": 2,
+        "sanitized": False,
+        "bounded": False,
+        "bound": "",
+        "trace": ["app.py:run:2 [CALL] request.args['cmd']", "app.py:run:3 [IDENTIFIER] cmd", "app.py:run:4 [CALL] os.system(cmd)"],
+        **extra,
+    }
+
+
+def unit(id="a", pin="pin", **extra):
+    return {
+        "unit": id,
+        "repo": "owner/repo",
+        "pin": pin,
+        "family": "injection",
+        "source": "vibe-py",
+        "excerpt": True,
+        "file": "app.py",
+        "function": "run",
+        "language": "python",
+        "pin_role": "vulnerable",
+        "pair": "pair",
+        "label": 1,
+        "fold": "outer-0",
+        "group": "owner/repo",
+        "supported": True,
+        "reason": "",
+        **extra,
+    }
+
+
+def test_trace_roles_and_path_level_sanitizer():
+    result = parse_trace(row(sanitized=True, bound="cmd", bounded=True))
+    assert [s["roles"] for s in result["steps"]] == [["source"], ["step", "guard"], ["sink"]]
+    assert result["steps"][0] == {
+        "file": "app.py",
+        "method": "run",
+        "line": 2,
+        "node_label": "CALL",
+        "code": "request.args['cmd']",
+        "roles": ["source"],
+        "raw": "app.py:run:2 [CALL] request.args['cmd']",
+    }
+    assert result["sanitizer"] == {"on_path": True, "step": None}
+    assert result["guard"] == {"text": "cmd", "bounded": True}
+    off_path = parse_trace(row(bound="len < 20", bounded=True))
+    assert all("guard" not in s["roles"] for s in off_path["steps"])
+    assert off_path["guard"]["text"] == "len < 20"
+
+
+def test_trace_unknown_location_and_method_colons():
+    parsed = parse_trace(row(trace=["?:-1 [UNKNOWN] x", "a.php:A::run:4 [CALL] sink(x)"]))
+    assert parsed["steps"][0]["line"] == -1
+    assert parsed["steps"][0]["method"] == "?"
+    assert parsed["steps"][1]["method"] == "A::run"
+    assert parse_trace(row(trace=["a.py:run:2 [CALL] x"]))["steps"][0]["roles"] == ["source", "sink"]
+    with pytest.raises(ValueError, match="invalid trace"):
+        parse_trace(row(trace=["unreadable"]))
+
+
+def test_rendered_default_scan_params_do_not_request_trace(tmp_path):
+    spec = taint_specs(language="python")["injection"]
+    rendered = []
+
+    def runner(command, **kwargs):
+        request_file = next(arg.split("=", 1)[1] for arg in command if arg.startswith("requestsFile="))
+        rendered.append(json.loads(Path(request_file).read_text()))
+        return subprocess.CompletedProcess(command, 0, BEGIN + '\n{"r": []}\n' + END, "")
+
+    backend = JoernBackend(runner=runner, session_transport=False)
+    cpg = tmp_path / "cpg.bin"
+    cpg.write_bytes(b"graph")
+    backend._batch_once(cpg, "taint", {"r": request_params(spec)})
+    backend._batch_once(cpg, "taint", {"r": request_params(spec, trace=True)})
+    assert "trace" not in rendered[0]["r"]
+    assert rendered[1]["r"]["trace"] == "true"
+    script = backend.queries_dir / "taint.sc"
+    assert "traceS = trace" in script.read_text()  # single-query forwarding too
+
+
+def test_plan_groups_skips_done_and_retains_other_families(tmp_path):
+    units = [unit(), unit("b", label=0), unit("c", "other")]
+    plan = trace.make_plan(units, tmp_path)
+    assert sorted(len(p["units"]) for p in plan) == [1, 2]
+    done = next(p for p in plan if p["pin"] == "pin")
+    trace.write_json(tmp_path / trace.pin_name(done), {**done, "done": True})
+    assert len(trace.make_plan(units, tmp_path)) == 1
+    assert trace.make_plan(units, tmp_path, only_family="path") == []
+    assert len(trace.make_plan(units, tmp_path, limit=1)) == 1
+    # An interrupted pin is rerun, even if a checkpoint file exists.
+    trace.write_json(tmp_path / trace.pin_name(done), {**done, "done": False})
+    assert len(trace.make_plan(units, tmp_path)) == 2
+
+
+def test_join_requires_exact_examples_and_ts_opt_in():
+    u = unit()
+    example = {**u, "id": "a", "candidate": "app.ts::run", "source": "pairs", "language": "typescript"}
+    with pytest.raises(ValueError, match="missing 1/1"):
+        trace.join_units([u], [])
+    assert not trace.join_units([u], [example])[0]["supported"]
+    assert trace.join_units([u], [example], include_typescript=True)[0]["supported"]
+
+
+class FakeBackend:
+    seconds = 20
+    census_files = 1
+    timed_out = False
+    fatal = ""
+    last_failure = ""
+    seen = []
+
+    def __init__(self, deadline, question_deadline, checkpoint):
+        self.stages = [{"command": "joern-parse", "seconds": self.seconds, "exit": 0}]
+        self.asked = []
+        self.answers = {}
+
+    def build(self, root, language):
+        graph = root / "graph.bin"
+        graph.write_bytes(b"CPG bytes")
+
+        def batch(query, requests):
+            self.seen.append(requests)
+            self.asked.extend(requests)
+            self.stages.append({"command": "joern", "seconds": self.seconds, "exit": 0})
+            return {
+                **{rid: [row()] if rid == "a" else [] for rid in requests},
+                "__census__": [{"files": str(self.census_files), "file_names": ["app.py"]}],
+            }
+
+        return CpgResult(graph, lambda q, p: None, run_batch_once=batch)
+
+
+def run_fake(tmp_path, monkeypatch, backend=FakeBackend, content="def run():\n    pass\n"):
+    monkeypatch.setattr(trace, "MeasuredBackend", backend)
+    (tmp_path / "app.py").write_text(content)
+    pin = {"repo": "owner/repo", "pin": "pin", "units": [unit(), unit("b")]}
+    out = tmp_path / "result.json"
+    trace.worker(pin, tmp_path, out, 900, 120)
+    return json.loads(out.read_text())
+
+
+def test_worker_one_batch_and_preserves_witnesses(tmp_path, monkeypatch):
+    FakeBackend.seen = []
+    result = run_fake(tmp_path, monkeypatch)
+    assert [u["status"] for u in result["units"]] == ["path", "asked-nothing"]
+    assert len(FakeBackend.seen) == 1
+    assert all(p["trace"] == "true" and p["function"] == "run" and p["file"] == "app.py" for p in FakeBackend.seen[0].values())
+    assert result["units"][0]["witness_rows"] == [row()]
+    assert result["units"][0]["traces"][0]["steps"][-1]["roles"] == ["sink"]
+    assert result["units"][1]["questions_completed"] == ["b"]
+    assert result["instrument"]["files"][0]["bytes"] > 0
+
+
+@pytest.mark.parametrize("failure", ["fast", "zero-input", "zero-census"])
+def test_invalid_instrument_is_persisted_failed(tmp_path, monkeypatch, failure):
+    class BadBackend(FakeBackend):
+        seconds = 0.5 if failure == "fast" else 20
+        census_files = 0 if failure == "zero-census" else 1
+
+    result = run_fake(tmp_path, monkeypatch, BadBackend, "" if failure == "zero-input" else "def run():\n    pass\n")
+    assert result["done"]
+    assert all(u["status"] == "failed" for u in result["units"])
+    assert all(u["reason"] for u in result["units"])
+    assert "<5s" in result["units"][0]["reason"] if failure == "fast" else "zero" in result["units"][0]["reason"]
+
+
+def test_missing_labelled_function_is_failure(tmp_path, monkeypatch):
+    result = run_fake(tmp_path, monkeypatch, content="def other():\n    pass\n")
+    assert result["units"][0]["status"] == "failed"
+    assert "not declared" in result["units"][0]["reason"]
+
+
+def test_docker_command_is_serial_offline_frozen():
+    cmd = trace.docker_command(Path("/export"), Path("/pin"), Path("/output"), "one", "openultrasast:dev", 900, 60)
+    assert cmd[cmd.index("--memory") + 1] == "3g"
+    assert cmd[cmd.index("--network") + 1] == "none"
+    assert cmd[cmd.index("--pull") + 1] == "never"
+    assert "/export:/frozen:ro" in cmd and "/pin:/case:ro" in cmd
+    assert cmd[-1] == "60"
+
+
+def test_summary_separates_unanswered_pairs(tmp_path):
+    rows = [
+        trace.unit_record(unit(), "path"),
+        trace.unit_record(unit("b", label=0), "asked-nothing"),
+        trace.unit_record(unit("c", pair="other"), "failed"),
+        trace.unit_record(unit("d", pair="other", label=0), "timeout"),
+    ]
+    trace.write_json(tmp_path / "one.json", {"done": True, "units": rows})
+    result = trace.summary(tmp_path)
+    assert result["pair_asymmetry"] == {"vulnerable-only": 1, "neither": 1}
+    assert result["pair_asymmetry_both_answered"] == {"vulnerable-only": 1}
+    assert result["by_family"]["injection"]["failed"] == 1
+
+
+def test_frozen_source_overlays_uncommitted_files_without_live_mount(tmp_path, monkeypatch):
+    root = tmp_path / "repo"
+    root.mkdir()
+    for name in trace.SNAPSHOT_FILES:
+        dest = root / name
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        dest.write_text("snapshot " + name)
+    monkeypatch.setattr(trace, "ROOT", root)
+    commands = []
+
+    def run(cmd, **kwargs):
+        commands.append(cmd)
+        return subprocess.CompletedProcess(cmd, 0, "head-sha\n", "")
+
+    monkeypatch.setattr(trace.subprocess, "run", run)
+    frozen = tmp_path / "frozen"
+    proof = trace.freeze_source(frozen)
+    assert commands[0][-3:] == ["archive", "HEAD", "src"]
+    for name in trace.SNAPSHOT_FILES:
+        assert (frozen / name).read_bytes() == (root / name).read_bytes()
+        assert len(proof["overlay_sha256"][name]) == 64
+
+
+def test_question_watchdog_retains_completed_rows_without_joern(tmp_path):
+    import sys
+
+    code = (
+        "import json,time\n"
+        'print(json.dumps({"__question__":"a"}),flush=True)\n'
+        'print(json.dumps({"id":"a","rows":[]}),flush=True)\n'
+        'print(json.dumps({"__question__":"b"}),flush=True)\n'
+        "time.sleep(60)\n"
+    )
+    backend = trace.MeasuredBackend(10, 0.05, lambda: None)
+    done = backend._run([sys.executable, "-c", code], timeout=10, cwd=tmp_path)
+    assert backend.timed_out
+    assert backend.asked == ["a", "b"]
+    assert backend.answers == {"a": []}
+    assert done.returncode != 0
+    assert backend.stages[0]["seconds"] < 5
+    assert "implausibly fast" in backend.fatal
+
+
+def test_pair_mapping_handles_harvest_normalized_crlf(tmp_path):
+    import hashlib
+    from types import SimpleNamespace
+
+    file = tmp_path / "excerpt.py"
+    file.write_bytes(b"def run():\r\n    pass\r\n")
+    normalized = file.read_text().encode()
+    pin = hashlib.sha1(b"blob " + str(len(normalized)).encode() + b"\0" + normalized).hexdigest()
+    case = SimpleNamespace(name="case", repo="owner/repo", relpath="app.py", vuln_file=file, fixed_file=file)
+    inputs = object.__new__(trace.Inputs)
+    inputs.pairs = [case]
+    inputs.hashes = {}
+    assert inputs.pair_for(unit(pin=pin)) == (case, "vuln")
+
+
+def test_missing_batch_answers_are_timeout_not_empty(tmp_path, monkeypatch):
+    class TimeoutBackend(FakeBackend):
+        def build(self, root, language):
+            graph = root / "graph.bin"
+            graph.write_bytes(b"CPG bytes")
+
+            def batch(query, requests):
+                self.timed_out = True
+                self.asked.append("a")
+                return None
+
+            return CpgResult(graph, lambda q, p: None, run_batch_once=batch)
+
+    result = run_fake(tmp_path, monkeypatch, TimeoutBackend)
+    assert [r["status"] for r in result["units"]] == ["timeout", "timeout"]
+    assert result["units"][0]["questions_asked"] == ["a"]
+    assert result["units"][1]["questions_asked"] == []
+
+
+def test_dry_run_makes_no_subprocess_calls(tmp_path, monkeypatch, capsys):
+    import sys
+
+    u = unit()
+    examples = [{**u, "id": "a", "candidate": "app.py::run", "source": "population-v1"}]
+    units = tmp_path / "units.jsonl"
+    units.write_text(json.dumps(u) + "\n")
+    example_file = tmp_path / "examples.jsonl"
+    example_file.write_text(json.dumps(examples[0]) + "\n")
+    monkeypatch.setattr(
+        sys, "argv", [str(SCRIPT), "--dry-run", "--units", str(units), "--examples", str(example_file), "--out", str(tmp_path / "results")]
+    )
+    monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("dry-run spawned a subprocess"))
+    assert trace.main() == 0
+    plan = json.loads(capsys.readouterr().out)
+    assert plan["runnable_pins"] == 1 and plan["units"] == 1
+    assert plan["estimated_seconds"] == 182
+    assert not (tmp_path / "results").exists()
+
+
+def test_family_filtered_results_merge_without_losing_prior_units(tmp_path):
+    first = trace.make_plan([unit(), unit("b", family="path")], tmp_path, only_family="injection")[0]
+    trace.save_pin(tmp_path / trace.pin_name(first), {**first, "done": True})
+    second = trace.make_plan([unit(), unit("b", family="path")], tmp_path, only_family="path")[0]
+    trace.save_pin(tmp_path / trace.pin_name(second), {**second, "done": True})
+    assert trace.make_plan([unit(), unit("b", family="path")], tmp_path) == []
+
+
+def test_failed_launch_persists_failed_pin_and_resume_skips_it(tmp_path, monkeypatch):
+    from types import SimpleNamespace
+
+    class Inputs:
+        def materialize(self, pin, target):
+            target.mkdir()
+            (target / "app.py").write_text("def run():\n    pass\n")
+
+    commands = []
+
+    def run(cmd, **kwargs):
+        commands.append(cmd)
+        return subprocess.CompletedProcess(cmd, 0, "", "")
+
+    monkeypatch.setattr(trace.subprocess, "run", run)
+    pin = trace.make_plan([unit()], tmp_path)[0]
+    args = SimpleNamespace(image="openultrasast:dev", deadline=900, question_deadline=120)
+    assert trace.run_pin(pin, tmp_path / "frozen", {}, Inputs(), tmp_path, args)
+    result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
+    assert result["units"][0]["status"] == "failed"
+    assert "produced no result" in result["units"][0]["reason"]
+    assert [c[1] for c in commands] == ["run", "rm"]
+    assert trace.make_plan([unit()], tmp_path) == []