mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
Implemented with Codex. Live smoke pending the operator egress watcher. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1599 lines
71 KiB
Python
1599 lines
71 KiB
Python
"""Offline controls for the trace experiment; no Docker, Joern or network."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from collections import Counter
|
|
from dataclasses import replace
|
|
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
|
|
|
|
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)
|
|
with pytest.MonkeyPatch.context() as patch:
|
|
patch.syspath_prepend(str(SCRIPT.parent))
|
|
from engine_trace_parse import parse_trace
|
|
|
|
_spec.loader.exec_module(trace)
|
|
worker = sys.modules["engine_trace_worker"]
|
|
|
|
|
|
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, "cpg_bytes": 2048}]
|
|
self.asked = []
|
|
self.answers = {}
|
|
|
|
def build(self, root, language):
|
|
graph = root / "graph.bin"
|
|
graph.write_bytes(b"x" * 2048)
|
|
|
|
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(worker, "MeasuredBackend", backend)
|
|
(tmp_path / "app.py").write_text(content)
|
|
pin = {"repo": "owner/repo", "pin": "pin", "units": [unit(), unit("b")]}
|
|
out = tmp_path / "result.json"
|
|
pin["questions"] = trace.prepare_questions(pin, tmp_path, 120) if content else {}
|
|
worker.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("json_roundtrip", [False, True])
|
|
def test_worker_writes_string_requests_for_php_and_python(tmp_path, monkeypatch, json_roundtrip):
|
|
written = {}
|
|
submitted = {}
|
|
|
|
class RequestBackend(worker.MeasuredBackend):
|
|
def build(self, root, language):
|
|
graph = root / f"{language}.bin"
|
|
graph.write_bytes(b"x" * 2048)
|
|
self.stages.append({"command": "joern-parse", "seconds": 20, "exit": 0, "cpg_bytes": 2048})
|
|
|
|
def batch(query, requests):
|
|
submitted.update(requests)
|
|
return self._batch_once(graph, query, requests)
|
|
|
|
return CpgResult(graph, lambda q, p: None, run_batch_once=batch)
|
|
|
|
def _run(self, command, **kwargs):
|
|
request_file = Path(next(arg.split("=", 1)[1] for arg in command if arg.startswith("requestsFile=")))
|
|
data = request_file.read_bytes()
|
|
print(f"read {request_file.name}: {len(data)} bytes")
|
|
requests = json.loads(data)
|
|
written.update(requests)
|
|
self.asked.extend(requests)
|
|
self.stages.append({"command": "joern", "seconds": 20, "exit": 0})
|
|
answers = {**dict.fromkeys(requests, []), "__census__": [{"files": "1", "file_names": ["app.php", "app.py"]}]}
|
|
return subprocess.CompletedProcess(command, 0, BEGIN + "\n" + json.dumps(answers) + "\n" + END, "")
|
|
|
|
monkeypatch.setattr(worker, "MeasuredBackend", RequestBackend)
|
|
(tmp_path / "app.php").write_text("<?php function run() {}\n")
|
|
(tmp_path / "app.py").write_text("def run(): pass\n")
|
|
pin = {"units": [unit("php", language="php", file="app.php"), unit("python")]}
|
|
pin["questions"] = trace.prepare_questions(pin, tmp_path, 120)
|
|
assert all(isinstance(params["sources"], tuple) and params["sources"] for params in pin["questions"].values())
|
|
expected = {rid: ",".join(sorted(params["sources"])) for rid, params in pin["questions"].items()}
|
|
# Cover both direct tuple callers and arrays decoded from the host's pin.json.
|
|
if json_roundtrip:
|
|
pin = json.loads(json.dumps(pin))
|
|
out = tmp_path / "result.json"
|
|
worker.worker(pin, tmp_path, out, 900, 120)
|
|
result = json.loads(out.read_text())
|
|
assert all(row["status"] == "asked-nothing" for row in result["units"])
|
|
assert set(written) == set(submitted) == {"php", "python"}
|
|
for requests in (written, submitted, {rid: p for row in result["units"] for rid, p in row["questions"].items()}):
|
|
assert all(isinstance(value, str) for params in requests.values() for value in params.values())
|
|
assert {rid: params["sources"] for rid, params in requests.items()} == expected
|
|
|
|
|
|
def test_failed_query_records_non_stack_stderr_tail_in_unit_reason(tmp_path, monkeypatch):
|
|
error = 'ujson.Value$InvalidData: Expected ujson.Str (data: ["$_COOKIE","$_GET"])'
|
|
joern = tmp_path / "joern"
|
|
joern.write_text(
|
|
f"#!{sys.executable}\nimport sys\n"
|
|
"print('read bytes:', len(open(sys.argv[-1], 'rb').read()), flush=True)\n"
|
|
"for i in range(25): print(f'diagnostic {i}', file=sys.stderr)\n"
|
|
f"print({error!r}, file=sys.stderr)\n"
|
|
"for i in range(30): print(f'\\tat example.Frame.method(Frame.java:{i})', file=sys.stderr)\n"
|
|
"print('\\t... 30 more', file=sys.stderr)\n"
|
|
"sys.exit(1)\n"
|
|
)
|
|
joern.chmod(0o755)
|
|
|
|
class FailedQueryBackend(worker.MeasuredBackend):
|
|
def build(self, root, language):
|
|
graph = root / "graph.bin"
|
|
graph.write_bytes(b"x" * 2048)
|
|
self.stages.append({"command": "joern-parse", "seconds": 20, "exit": 0, "cpg_bytes": 2048})
|
|
|
|
def batch(query, requests):
|
|
done = self._run([str(joern), "--script", str(root / "app.py")], timeout=10, cwd=root)
|
|
assert done.returncode == 1
|
|
assert "read bytes: 20" in done.stdout
|
|
return None
|
|
|
|
return CpgResult(graph, lambda q, p: None, run_batch_once=batch)
|
|
|
|
result = run_fake(tmp_path, monkeypatch, FailedQueryBackend)
|
|
stage = result["instrument"]["jvm"][-1]
|
|
assert stage["exit"] == 1
|
|
assert stage["output_tail"].splitlines() == [*[f"diagnostic {i}" for i in range(6, 25)], error]
|
|
for row in result["units"]:
|
|
assert row["status"] == "failed"
|
|
assert "joern query exited 1" in row["reason"]
|
|
assert error in row["reason"]
|
|
assert "Frame.java" not in row["reason"]
|
|
assert "implausibly fast" not in row["reason"]
|
|
|
|
|
|
@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):
|
|
with pytest.raises(ValueError, match="not declared"):
|
|
run_fake(tmp_path, monkeypatch, content="def other():\n pass\n")
|
|
|
|
|
|
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"
|
|
assert "/frozen/benchmarks/learn/engine_trace_worker.py" in cmd
|
|
assert "benchmarks/learn/engine_trace_worker.py" in trace.SNAPSHOT_FILES
|
|
assert "benchmarks/learn/engine_trace_parse.py" in trace.SNAPSHOT_FILES
|
|
|
|
|
|
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 = worker.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 == []
|
|
assert not 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"x" * 2048)
|
|
|
|
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", "excerpt": False}]
|
|
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_replans_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")
|
|
return {"mode": "source"}
|
|
|
|
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 if c[0] == "docker"] == ["run", "rm"]
|
|
assert len(trace.make_plan([unit()], tmp_path)) == 1
|
|
assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
|
|
|
|
|
|
def import_container_worker():
|
|
"""Import in a fresh interpreter so host imports cannot mask the boundary."""
|
|
script = SCRIPT.with_name("engine_trace_worker.py")
|
|
code = """
|
|
import json, pathlib, runpy, sys
|
|
sys.modules['yaml'] = None
|
|
sys.modules['boto3'] = None
|
|
script = pathlib.Path(sys.argv[1])
|
|
size = len(script.read_bytes())
|
|
# Direct script execution in the image puts the sibling modules on sys.path.
|
|
sys.path.insert(0, str(script.parent))
|
|
runpy.run_path(str(script), run_name='trace_import_check')
|
|
print(json.dumps({'bytes': size, 'modules': sorted(
|
|
name for name in sys.modules if name == 'openultrasast' or name.startswith('openultrasast.')
|
|
)}))
|
|
"""
|
|
result = subprocess.run([sys.executable, "-c", code, str(script)], capture_output=True, text=True, timeout=30)
|
|
assert result.returncode == 0, result.stderr
|
|
return json.loads(result.stdout)
|
|
|
|
|
|
def test_trace_container_import_without_yaml_or_boto3():
|
|
assert import_container_worker()["bytes"] > 0
|
|
|
|
|
|
def test_trace_container_transitive_import_boundary():
|
|
# Explicit inventory: additions need review, particularly learn.examples/plane.*.
|
|
assert import_container_worker()["modules"] == [
|
|
"openultrasast",
|
|
"openultrasast.contracts",
|
|
"openultrasast.cpg",
|
|
"openultrasast.cpg.artifact",
|
|
"openultrasast.cpg.backend",
|
|
"openultrasast.cpg.session",
|
|
"openultrasast.model",
|
|
"openultrasast.model.contracts",
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("existing", [False, True])
|
|
def test_trace_dry_run_unsupported_units_write_nothing(tmp_path, monkeypatch, capsys, existing):
|
|
out = tmp_path / "results"
|
|
if existing:
|
|
out.mkdir()
|
|
(out / ".lock").write_bytes(b"existing lock")
|
|
(out / "checkpoint.json").write_text('{"done": false}')
|
|
(out / "previous.json").write_text('{"done": true, "units": []}')
|
|
before = {p.name: (p.read_bytes(), p.stat().st_mtime_ns) for p in out.glob("*")}
|
|
units = [unit("ts", file="app.ts"), unit("py", pin="python")]
|
|
examples = [{**u, "id": u["unit"], "candidate": u["file"] + "::run", "source": "population-v1", "excerpt": False} for u in units]
|
|
units_file, examples_file = tmp_path / "units.jsonl", tmp_path / "examples.jsonl"
|
|
units_file.write_text("\n".join(map(json.dumps, units)))
|
|
examples_file.write_text("\n".join(map(json.dumps, examples)))
|
|
monkeypatch.setattr(
|
|
sys, "argv", [str(SCRIPT), "--dry-run", "--units", str(units_file), "--examples", str(examples_file), "--out", str(out)]
|
|
)
|
|
monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("dry run spawned subprocess"))
|
|
monkeypatch.setattr(trace, "write_json", lambda *a: pytest.fail("dry run wrote JSON"))
|
|
assert trace.main() == 0
|
|
plan = json.loads(capsys.readouterr().out)
|
|
assert plan["pins"] == 2 and plan["runnable_pins"] == 1
|
|
assert out.exists() == existing
|
|
assert {p.name: (p.read_bytes(), p.stat().st_mtime_ns) for p in out.glob("*")} == before
|
|
|
|
|
|
def test_trace_real_unsupported_run_then_typescript_opt_in_replans(tmp_path, monkeypatch, capsys):
|
|
u = unit(file="app.ts")
|
|
example = {**u, "id": "a", "candidate": "app.ts::run", "source": "population-v1", "excerpt": False}
|
|
units_file, examples_file, out = tmp_path / "units.jsonl", tmp_path / "examples.jsonl", tmp_path / "results"
|
|
units_file.write_text(json.dumps(u) + "\n")
|
|
examples_file.write_text(json.dumps(example) + "\n")
|
|
argv = [str(SCRIPT), "--units", str(units_file), "--examples", str(examples_file), "--out", str(out)]
|
|
monkeypatch.setattr(sys, "argv", argv)
|
|
monkeypatch.setattr(trace, "freeze_source", lambda target: {})
|
|
monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("unsupported run invoked Docker"))
|
|
assert trace.main() == 0
|
|
result_file = out / trace.pin_name(u)
|
|
first = json.loads(result_file.read_text())
|
|
assert first["done"] and first["units"][0]["status"] == "unsupported"
|
|
progress = json.loads((out / "progress.json").read_text())
|
|
assert progress == {
|
|
"pins_done": 1,
|
|
"pins_total": 1,
|
|
"last_pin": {"repo": u["repo"], "pin": u["pin"]},
|
|
"statuses": {"unsupported": 1},
|
|
"lanes": {"docker": 1},
|
|
"pins_in_flight": 0,
|
|
"active_lanes": {},
|
|
}
|
|
assert trace.make_plan(trace.join_units([u], [example]), out) == []
|
|
|
|
called = []
|
|
|
|
def run_pin(pin, source, provenance, inputs, results, args):
|
|
called.append(pin)
|
|
trace.save_pin(
|
|
results / trace.pin_name(pin), {**pin, "done": True, "units": [trace.unit_record(r, "asked-nothing") for r in pin["units"]]}
|
|
)
|
|
return False
|
|
|
|
def inspect_image(command, **kwargs):
|
|
assert command[:3] == ["docker", "image", "inspect"]
|
|
return subprocess.CompletedProcess(command, 0, "sha256:fake\n", "")
|
|
|
|
monkeypatch.setattr(trace, "run_pin", run_pin)
|
|
monkeypatch.setattr(trace.subprocess, "run", inspect_image)
|
|
monkeypatch.setattr(sys, "argv", [*argv, "--include-typescript"])
|
|
assert trace.main() == 0
|
|
assert len(called) == 1 and called[0]["units"][0]["supported"]
|
|
assert called[0]["selection"] != first["selection"]
|
|
assert json.loads(result_file.read_text())["units"][0]["status"] == "asked-nothing"
|
|
assert trace.make_plan(trace.join_units([u], [example], include_typescript=True), out, include_typescript=True) == []
|
|
|
|
|
|
def test_trace_resume_identity_includes_selection_even_if_support_unchanged(tmp_path):
|
|
first = trace.make_plan([unit()], tmp_path)[0]
|
|
trace.save_pin(tmp_path / trace.pin_name(first), {**first, "done": True})
|
|
assert trace.make_plan([unit()], tmp_path) == []
|
|
assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
|
|
assert len(trace.make_plan([unit()], tmp_path, include_typescript=True)) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("binary", sorted(worker.FRONTENDS))
|
|
def test_fast_frontend_validates_exact_output(tmp_path, binary):
|
|
frontend = tmp_path / binary
|
|
frontend.write_text(
|
|
f"#!{sys.executable}\n"
|
|
"import pathlib, sys, time\n"
|
|
"source = pathlib.Path(sys.argv[1])\n"
|
|
"print('read bytes:', len(source.read_bytes()), flush=True)\n"
|
|
"time.sleep(2 if pathlib.Path(sys.argv[0]).name == 'php2cpg' else 0)\n"
|
|
"pathlib.Path(sys.argv[3]).write_bytes(b'x' * 34000)\n"
|
|
)
|
|
frontend.chmod(0o755)
|
|
(tmp_path / "source.php").write_text("<?php echo 'hello';")
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
flag = "--output" if binary == "joern-parse" else "-o"
|
|
done = backend._run([str(frontend), "source.php", flag, "out.graph"], timeout=10, cwd=tmp_path)
|
|
assert done.returncode == 0
|
|
assert "read bytes: 19" in done.stdout
|
|
assert not backend.fatal
|
|
stage = backend.stages[0]
|
|
assert stage["seconds"] < 5
|
|
if binary == "php2cpg":
|
|
assert stage["seconds"] >= 2
|
|
assert stage["cpg_path"] == str(tmp_path / "out.graph")
|
|
assert stage["cpg_bytes"] == 34000
|
|
proof = {"files": ["source.php"], "bytes": 19, **stage, "jvm": backend.stages}
|
|
assert worker.instrument_failure(proof) == ""
|
|
|
|
|
|
@pytest.mark.parametrize("size", [None, 0, 1024])
|
|
def test_frontend_without_sufficient_output_fails(tmp_path, size):
|
|
frontend = tmp_path / "php2cpg"
|
|
frontend.write_text(
|
|
f"#!{sys.executable}\nimport pathlib, sys\n"
|
|
+ (f"pathlib.Path(sys.argv[2]).write_bytes(b'x' * {size})\n" if size is not None else "")
|
|
)
|
|
frontend.chmod(0o755)
|
|
# A different non-empty CPG must not mask the missing requested output.
|
|
(tmp_path / "unrelated.bin").write_bytes(b"x" * 34000)
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
done = backend._run([str(frontend), "-o", "out.bin"], timeout=10, cwd=tmp_path)
|
|
assert done.returncode == 1
|
|
assert backend.stages[0]["exit"] == 0
|
|
assert backend.stages[0]["cpg_bytes"] == (size or 0)
|
|
assert backend.fatal == "frontend wrote no CPG"
|
|
|
|
|
|
def test_fast_joern_script_is_failed_launch(tmp_path):
|
|
joern = tmp_path / "joern"
|
|
joern.write_text(f"#!{sys.executable}\n")
|
|
joern.chmod(0o755)
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
done = backend._run([str(joern), "--script", "query.sc"], timeout=10, cwd=tmp_path)
|
|
assert done.returncode == 1
|
|
assert backend.stages[0]["exit"] == 0
|
|
assert "implausibly fast JVM step" in backend.fatal
|
|
|
|
|
|
@pytest.mark.parametrize("markers", [[], ["a"], ["a", "b"]])
|
|
def test_empty_answers_require_each_question_marker(tmp_path, monkeypatch, markers):
|
|
class EmptyBackend(FakeBackend):
|
|
def build(self, root, language):
|
|
cpg = super().build(root, language)
|
|
batch = cpg.run_batch_once
|
|
|
|
def empty_batch(query, requests):
|
|
answers = batch(query, requests)
|
|
self.asked = markers
|
|
return {**answers, **dict.fromkeys(requests, [])}
|
|
|
|
return replace(cpg, run_batch_once=empty_batch)
|
|
|
|
result = run_fake(tmp_path, monkeypatch, EmptyBackend)
|
|
for row in result["units"]:
|
|
assert row["status"] == ("asked-nothing" if row["unit"] in markers else "failed")
|
|
assert row["reason"] == ("" if row["unit"] in markers else "query produced no markers")
|
|
|
|
|
|
def test_frontend_measurement_survives_failed_build_cleanup(tmp_path, monkeypatch):
|
|
class CleanupBackend(worker.MeasuredBackend):
|
|
def build(self, root, language):
|
|
frontend = root / "php2cpg"
|
|
frontend.write_text(
|
|
f"#!{sys.executable}\nimport pathlib, sys\npathlib.Path(sys.argv[2]).write_bytes(b'x' * 34000)\nsys.exit(1)\n"
|
|
)
|
|
frontend.chmod(0o755)
|
|
self._run([str(frontend), "-o", "actual.graph"], timeout=10, cwd=root)
|
|
(root / "actual.graph").unlink()
|
|
return None
|
|
|
|
result = run_fake(tmp_path, monkeypatch, CleanupBackend)
|
|
assert all(row["status"] == "failed" for row in result["units"])
|
|
assert result["instrument"]["cpg_path"] == str(tmp_path / "actual.graph")
|
|
assert result["instrument"]["cpg_bytes"] == 34000
|
|
assert not (tmp_path / "actual.graph").exists()
|
|
|
|
|
|
def pair_inputs(tmp_path):
|
|
from types import SimpleNamespace
|
|
|
|
inputs = trace.Inputs(trace.ROOT, tmp_path / "cache")
|
|
excerpt = tmp_path / "excerpt.py"
|
|
excerpt.write_text("# commit: " + "f" * 40 + "\n# parent: " + "f" * 40 + "\ndef run(): pass\n")
|
|
case = SimpleNamespace(name="fixture", repo="owner/repo", relpath="pkg/app.py", vuln_file=excerpt, fixed_file=excerpt, context_files=())
|
|
inputs.catalog_rows[case.name] = [
|
|
("advisory-fixes-2", {"repo": case.repo, "relpath": case.relpath, "parent": "a" * 40, "commit": "b" * 40})
|
|
]
|
|
return inputs, case
|
|
|
|
|
|
@pytest.mark.parametrize("side,sha", [("vuln", "a" * 40), ("fixed", "b" * 40)])
|
|
def test_pair_materializes_catalog_side_not_header(tmp_path, monkeypatch, side, sha):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
case.vuln_file.write_text("# Provenance: example\n def run(self): pass\n")
|
|
fetched = tmp_path / "fetched"
|
|
(fetched / "pkg").mkdir(parents=True)
|
|
(fetched / case.relpath).write_text("class Handler:\n def run(self): pass\n")
|
|
calls = []
|
|
|
|
def fetch(repo, commit, roots):
|
|
calls.append((repo, commit, roots))
|
|
return fetched
|
|
|
|
monkeypatch.setattr(inputs, "fetch_source", fetch)
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, side))
|
|
target = tmp_path / "export"
|
|
proof = inputs.materialize({"units": [unit(file=case.relpath)]}, target)
|
|
assert calls == [("owner/repo", sha, ["pkg"])]
|
|
assert proof["commit"] == sha and proof["mode"] == "source"
|
|
assert (target / case.relpath).read_text().startswith("class Handler:")
|
|
|
|
|
|
@pytest.mark.parametrize("change", [{"parent": ""}, {"parent": "b" * 40}, {"commit": "HEAD"}])
|
|
def test_catalog_ambiguous_commits_refused(tmp_path, change):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name][0][1].update(change)
|
|
with pytest.raises(ValueError, match="ambiguous catalog commits"):
|
|
inputs.source_identity(case, "vuln")
|
|
|
|
|
|
def test_ambiguous_catalog_and_side_refused(tmp_path):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name] *= 2
|
|
with pytest.raises(ValueError, match="ambiguous catalog source"):
|
|
inputs.source_identity(case, "fixed")
|
|
with pytest.raises(ValueError, match="ambiguous pair side"):
|
|
inputs.pair_for(unit(pin_role="unknown"))
|
|
|
|
|
|
def test_sparse_export_filters_languages_dirs_and_caps_context(tmp_path):
|
|
source, target = tmp_path / "source", tmp_path / "export"
|
|
contents = {
|
|
"pkg/app.py": "def run(): pass\n",
|
|
"pkg/context.py": "x" * 200,
|
|
"pkg/other.js": "function run() {}",
|
|
"other/app.py": "other",
|
|
}
|
|
contents.update({f"pkg/{directory}/a.py": "excluded" for directory in worker.EXCLUDED_DIRS})
|
|
for rel, data in contents.items():
|
|
path = source / rel
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(data)
|
|
proof = trace.export_sources(source, target, [unit(file="pkg/app.py")], ["pkg"], cap=50)
|
|
assert [p.relative_to(target).as_posix() for p in target.rglob("*") if p.is_file()] == ["pkg/app.py"]
|
|
assert proof["cap_hit"] and proof["bytes"] == len(contents["pkg/app.py"])
|
|
assert proof["excludes"] == sorted(worker.EXCLUDED_DIRS)
|
|
assert trace.source_root("module/src/main/java/org/Handler.java", "java") == "module/src/main/java"
|
|
assert trace.source_root("pkg/sub/app.ts", "typescript") == "pkg"
|
|
assert trace.source_root("app.js", "javascript") == "."
|
|
|
|
|
|
def test_wrong_fetched_function_fails_without_fallback(tmp_path, monkeypatch):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
source = tmp_path / "source"
|
|
(source / "pkg").mkdir(parents=True)
|
|
(source / case.relpath).write_text("def runner(): pass\n")
|
|
inputs.allow_excerpt_fallback = True
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, "vuln"))
|
|
monkeypatch.setattr(inputs, "fetch_source", lambda *args: source)
|
|
with pytest.raises(ValueError, match="labelled function not found in fetched source"):
|
|
inputs.materialize({"units": [unit(file=case.relpath)]}, tmp_path / "export")
|
|
|
|
|
|
def test_excerpt_fallback_requires_explicit_flag(tmp_path, monkeypatch):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, "vuln"))
|
|
|
|
def unavailable(*args):
|
|
raise ValueError("fetch failed")
|
|
|
|
monkeypatch.setattr(inputs, "fetch_source", unavailable)
|
|
pin = {"units": [unit(file=case.relpath)]}
|
|
with pytest.raises(ValueError, match="fetch failed"):
|
|
inputs.materialize(pin, tmp_path / "export")
|
|
inputs.allow_excerpt_fallback = True
|
|
proof = inputs.materialize(pin, tmp_path / "fallback")
|
|
assert proof["mode"] == "excerpt-fallback" and proof["reason"] == "fetch failed"
|
|
|
|
|
|
def test_fetch_renders_sparse_pinned_host_commands(tmp_path, monkeypatch):
|
|
inputs, _ = pair_inputs(tmp_path)
|
|
commands = []
|
|
|
|
def git(cmd, **kwargs):
|
|
commands.append(cmd[3:])
|
|
assert kwargs["env"]["GIT_TERMINAL_PROMPT"] == "0"
|
|
assert "GIT_NO_LAZY_FETCH" not in kwargs["env"]
|
|
return subprocess.CompletedProcess(cmd, 0, "a" * 40 + "\n", "")
|
|
|
|
monkeypatch.setattr(trace.subprocess, "run", git)
|
|
clone = inputs.fetch_source("owner/repo", "a" * 40, ["module/src/main/java"])
|
|
assert clone.is_relative_to(inputs.cache)
|
|
assert ["fetch", "--depth", "1", "--filter=blob:none", "https://github.com/owner/repo.git", "a" * 40] in commands
|
|
assert ["sparse-checkout", "set", "--", "module/src/main/java"] in commands
|
|
assert ["checkout", "--detach", "--force", "a" * 40] in commands
|
|
|
|
|
|
def test_explicit_heap_in_frontend_and_query_commands(tmp_path, monkeypatch):
|
|
backend = worker.MeasuredBackend(900, 120, lambda: None)
|
|
commands = []
|
|
|
|
def run(command, **kwargs):
|
|
commands.append(command)
|
|
return subprocess.CompletedProcess(command, 0, BEGIN + '\n{"r": []}\n' + END, "")
|
|
|
|
monkeypatch.setattr(backend, "_run", run)
|
|
monkeypatch.setattr("openultrasast.cpg.backend.shutil.which", lambda name: "/opt/joern-cli/" + name)
|
|
graph = tmp_path / "cpg.bin"
|
|
graph.write_bytes(b"x" * 2048)
|
|
backend._build_with_frontend(tmp_path, graph, tmp_path, "python")
|
|
backend._batch_once(graph, "taint", {"r": {}})
|
|
assert {Path(cmd[0]).name for cmd in commands} == {"pysrc2cpg", "joern"}
|
|
assert all("-J-Xmx2560m" in cmd for cmd in commands)
|
|
assert "-Xmx2560m" in backend._jvm_env()["JAVA_TOOL_OPTIONS"]
|
|
|
|
|
|
@pytest.mark.parametrize("binary", ["pysrc2cpg", "joern"])
|
|
def test_oom_output_overrides_missing_graph_or_markers(tmp_path, binary):
|
|
launcher = tmp_path / binary
|
|
launcher.write_text(f"#!{sys.executable}\nprint('java.lang.OutOfMemoryError: Java heap space')\n")
|
|
launcher.chmod(0o755)
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
done = backend._run([str(launcher), "--script", "query.sc"], timeout=10)
|
|
assert done.returncode == 1
|
|
assert backend.fatal == "JVM out of memory"
|
|
assert worker.instrument_failure({"jvm": backend.stages}) == "JVM out of memory"
|
|
|
|
|
|
def test_two_language_pin_builds_sequential_frontends_and_routes_questions(tmp_path, monkeypatch):
|
|
events = []
|
|
|
|
class LanguagesBackend(FakeBackend):
|
|
def build(self, root, language):
|
|
events.append(("build", language))
|
|
graph = root / (language + ".bin")
|
|
graph.write_bytes(b"x" * 2048)
|
|
self.stages = [
|
|
{"command": {"python": "pysrc2cpg", "javascript": "jssrc2cpg"}[language], "seconds": 20, "exit": 0, "cpg_bytes": 2048}
|
|
]
|
|
|
|
def batch(query, requests):
|
|
events.append(("query", language, sorted(requests)))
|
|
self.asked.extend(requests)
|
|
return {
|
|
**dict.fromkeys(requests, []),
|
|
"__census__": [{"files": "1", "file_names": ["app.py" if language == "python" else "app.js"]}],
|
|
}
|
|
|
|
return CpgResult(graph, lambda q, p: None, run_batch_once=batch, cleanup=lambda: events.append(("cleanup", language)))
|
|
|
|
monkeypatch.setattr(worker, "MeasuredBackend", LanguagesBackend)
|
|
(tmp_path / "app.py").write_text("def run(): pass\n")
|
|
(tmp_path / "app.js").write_text("function run() {}\n")
|
|
pin = {"units": [unit(), unit("js", language="javascript", file="app.js")]}
|
|
pin["questions"] = trace.prepare_questions(pin, tmp_path, 120)
|
|
assert all(params["callDepth"] == "3" for params in pin["questions"].values())
|
|
result = worker.worker(pin, tmp_path, tmp_path / "result.json", 900, 120)
|
|
assert events == [
|
|
("build", "javascript"),
|
|
("query", "javascript", ["js"]),
|
|
("cleanup", "javascript"),
|
|
("build", "python"),
|
|
("query", "python", ["a"]),
|
|
("cleanup", "python"),
|
|
]
|
|
assert all(u["status"] == "asked-nothing" for u in result["units"])
|
|
assert set(result["instruments"]) == {"python", "javascript"}
|
|
assert all(u["instrument"]["heap_mb"] == 2560 for u in result["units"])
|
|
|
|
|
|
def test_resume_materialization_version_and_status_selection(tmp_path):
|
|
units = [unit(source="advisory-fixes-2"), unit("other", pin="other", source="population-v1")]
|
|
pin = trace.make_plan(units, tmp_path, only_source={"advisory-fixes-2"})[0]
|
|
record = {**pin, "done": True, "units": [trace.unit_record(pin["units"][0], "asked-nothing")]}
|
|
path = tmp_path / trace.pin_name(pin)
|
|
trace.write_json(path, record)
|
|
assert trace.make_plan(units, tmp_path, only_source={"advisory-fixes-2"}) == []
|
|
assert len(trace.make_plan(units, tmp_path, only_source={"advisory-fixes-2"}, rerun_status={"failed", "asked-nothing"})) == 1
|
|
trace.write_json(path, {**record, "materialization_version": trace.MATERIALIZATION_VERSION - 1})
|
|
assert len(trace.make_plan(units, tmp_path, only_source={"advisory-fixes-2"})) == 1
|
|
|
|
|
|
def test_sparse_fetch_from_local_git_without_network(tmp_path, monkeypatch):
|
|
origin = tmp_path / "origin"
|
|
origin.mkdir()
|
|
real_run = subprocess.run
|
|
|
|
def git(*args):
|
|
return real_run(["git", "-C", str(origin), *args], capture_output=True, text=True, check=True).stdout.strip()
|
|
|
|
git("init")
|
|
git("config", "user.name", "Offline Test")
|
|
git("config", "user.email", "[email protected]")
|
|
for rel in ("module/src/main/java/Handler.java", "module/src/main/java/docs/Other.java", "elsewhere/Other.java"):
|
|
path = origin / rel
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text("class Handler { void run() {} }\n")
|
|
git("add", ".")
|
|
git("commit", "-m", "local fixture")
|
|
sha = git("rev-parse", "HEAD")
|
|
inputs, _ = pair_inputs(tmp_path)
|
|
|
|
def local_only(cmd, **kwargs):
|
|
# Substitute a local repository for the one URL; no network-capable command runs.
|
|
cmd = [str(origin) if arg == "https://github.com/owner/repo.git" else arg for arg in cmd]
|
|
return real_run(cmd, **kwargs)
|
|
|
|
monkeypatch.setattr(trace.subprocess, "run", local_only)
|
|
clone = inputs.fetch_source("owner/repo", sha, ["module/src/main/java"])
|
|
assert not (clone / "elsewhere/Other.java").exists()
|
|
target = tmp_path / "export"
|
|
labelled = unit(file="module/src/main/java/Handler.java", language="java")
|
|
proof = trace.export_sources(clone, target, [labelled], ["module/src/main/java"])
|
|
trace.verify_functions(target, [labelled])
|
|
assert proof["files"] == 1 and proof["bytes"] > 0
|
|
|
|
|
|
@pytest.mark.parametrize("side", ["vuln", "fixed"])
|
|
def test_complete_catalog_twins_and_context_need_no_commits(tmp_path, monkeypatch, side):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name][0][1].update(parent="", commit="")
|
|
fixed = tmp_path / "fixed.py"
|
|
fixed.write_text("def run(): return 'fixed'\n")
|
|
case.fixed_file = fixed
|
|
context = tmp_path / "context.py"
|
|
context.write_text("registered = True\n")
|
|
case.context_files = (("pkg/context.py", context, context),)
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, side))
|
|
monkeypatch.setattr(inputs, "fetch_source", lambda *args: pytest.fail("catalog files must not fetch"))
|
|
target = tmp_path / "export"
|
|
proof = inputs.materialize({"units": [unit(file=case.relpath)]}, target)
|
|
assert proof["mode"] == "catalog-file"
|
|
assert (target / case.relpath).read_bytes() == (fixed if side == "fixed" else case.vuln_file).read_bytes()
|
|
assert (target / "pkg/context.py").read_bytes() == context.read_bytes()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"body,reason",
|
|
[
|
|
("# Provenance: example upstream_start: 25\n def run():\n return (\n", "excerpt does not parse:.*never closed"),
|
|
("def run(:\n", "not complete Python"),
|
|
],
|
|
)
|
|
def test_incomplete_catalog_without_commits_fails_specifically(tmp_path, monkeypatch, body, reason):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name][0][1].update(parent="", commit="")
|
|
case.vuln_file.write_text(body)
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, "vuln"))
|
|
with pytest.raises(ValueError, match=reason):
|
|
inputs.materialize({"units": [unit(file=case.relpath)]}, tmp_path / "export")
|
|
|
|
|
|
@pytest.mark.parametrize("comment,newline", [("#", "\n"), ("//", "\n"), ("#", "\r\n")])
|
|
@pytest.mark.parametrize("upstream", [True, False])
|
|
def test_parseable_python_method_excerpt_is_in_function_only(tmp_path, monkeypatch, comment, newline, upstream):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name][0][1].update(parent="", commit="")
|
|
header = f"{comment} Provenance: example\n" + (f"{comment} upstream_start: 25\n" if upstream else "")
|
|
body = header + "\n def run(self, cmd):\n os.system(cmd)\n"
|
|
case.vuln_file.write_bytes(body.replace("\n", newline).encode())
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, "vuln"))
|
|
monkeypatch.setattr(inputs, "fetch_source", lambda *args: pytest.fail("excerpt attempted fetch"))
|
|
pin = {"units": [unit(file=case.relpath)]}
|
|
target = tmp_path / "export"
|
|
pin["materialization"] = inputs.materialize(pin, target)
|
|
assert pin["materialization"]["mode"] == "excerpt-parseable"
|
|
assert (target / case.relpath).read_text() == header.replace("//", "#") + "\ndef run(self, cmd):\n os.system(cmd)\n"
|
|
assert trace.unit_record(pin["units"][0], "path")["evidence_scope"] == "in-function-only"
|
|
params = trace.prepare_questions(pin, target, 60)["a"]
|
|
assert params["callDepth"] == "0" and params["parameterSources"] == "true"
|
|
|
|
|
|
def test_java_fragment_is_not_parseable_excerpt(tmp_path, monkeypatch):
|
|
inputs, case = pair_inputs(tmp_path)
|
|
inputs.catalog_rows[case.name][0][1].update(parent="", commit="")
|
|
case.relpath = "pkg/App.java"
|
|
case.vuln_file = case.fixed_file = tmp_path / "excerpt.java"
|
|
case.vuln_file.write_text("// Provenance: example\n// upstream_start: 25\n public void run() {\n execute(cmd);\n")
|
|
inputs.catalog_rows[case.name][0][1]["relpath"] = case.relpath
|
|
monkeypatch.setattr(inputs, "pair_for", lambda row: (case, "vuln"))
|
|
with pytest.raises(ValueError, match="excerpt does not parse:"):
|
|
inputs.materialize({"units": [unit(file=case.relpath, language="java")]}, tmp_path / "export")
|
|
|
|
|
|
def test_excerpt_without_cheap_parser_is_rejected(tmp_path, monkeypatch):
|
|
monkeypatch.setattr("openultrasast.semantic.extra.grammar_for", lambda language: None)
|
|
with pytest.raises(ValueError, match="excerpt does not parse: no cheap parser for javascript"):
|
|
trace.parseable_excerpt(tmp_path / "app.js", "// Provenance: example upstream_start: 25\nfunction run() {}\n")
|
|
|
|
|
|
def test_parseable_javascript_excerpt_uses_repository_parser(tmp_path):
|
|
from openultrasast.semantic.extra import grammar_for
|
|
|
|
if grammar_for("javascript") is None:
|
|
pytest.skip("optional JavaScript grammar unavailable")
|
|
assert (
|
|
trace.parseable_excerpt(tmp_path / "app.js", "// Provenance: example\n function run(cmd) { execute(cmd); }\n")
|
|
== "// Provenance: example\nfunction run(cmd) { execute(cmd); }\n"
|
|
)
|
|
|
|
|
|
def test_export_keeps_labels_over_cap_and_excluded_ancestors(tmp_path):
|
|
source, target = tmp_path / "source", tmp_path / "export"
|
|
for rel in ("templates/preferences.php", "templates/other.php", "other.php"):
|
|
path = source / rel
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text("<?php function run() {}\n")
|
|
proof = trace.export_sources(source, target, [unit(file="templates/preferences.php", language="php")], cap=1)
|
|
assert (target / "templates/preferences.php").read_bytes() == (source / "templates/preferences.php").read_bytes()
|
|
assert not (target / "templates/other.php").exists()
|
|
assert proof["cap_hit"] and proof["bytes"] > proof["cap_bytes"]
|
|
|
|
|
|
def test_oversized_package_narrows_before_truncation_and_keeps_import(tmp_path):
|
|
source, target = tmp_path / "source", tmp_path / "export"
|
|
for rel, body in {
|
|
"pkg/sub/app.py": "from helper import value\ndef run(): pass\n",
|
|
"pkg/sub/context.py": "x = 1\n",
|
|
"pkg/huge.py": "#" * 500,
|
|
"helper.py": "value = 1\n",
|
|
}.items():
|
|
path = source / rel
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(body)
|
|
proof = trace.export_sources(source, target, [unit(file="pkg/sub/app.py")], ["pkg"], cap=100)
|
|
assert proof["roots"] == ["pkg/sub"] and proof["requested_roots"] == ["pkg"]
|
|
assert proof["scope_narrowed"] and not proof["cap_hit"]
|
|
assert (target / "helper.py").exists() and (target / "pkg/sub/context.py").exists()
|
|
assert not (target / "pkg/huge.py").exists()
|
|
|
|
|
|
def test_local_repository_export_scopes_packages(tmp_path, monkeypatch):
|
|
inputs, _ = pair_inputs(tmp_path)
|
|
|
|
def export(clone, pin, dest):
|
|
for rel in ("packages/server/run.js", "packages/other/skip.js"):
|
|
path = dest / rel
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text("function run() {}\n")
|
|
|
|
monkeypatch.setattr(trace, "export_pin", export)
|
|
proof = inputs.export_local(
|
|
tmp_path, {"pin": "sha", "units": [unit(file="packages/server/run.js", language="javascript")]}, tmp_path / "out"
|
|
)
|
|
assert proof["roots"] == ["packages/server"]
|
|
assert not (tmp_path / "out/packages/other/skip.js").exists()
|
|
|
|
|
|
def test_census_membership_is_uncapped_and_warning_is_not_absence(tmp_path, monkeypatch):
|
|
class LargeBackend(FakeBackend):
|
|
def build(self, root, language):
|
|
cpg = super().build(root, language)
|
|
original = cpg.run_batch_once
|
|
|
|
def batch(query, requests):
|
|
result = original(query, requests)
|
|
result["__census__"] = [{"files": "3001", "file_names": [f"f{i}.py" for i in range(3000)] + ["/case/app.py"]}]
|
|
return result
|
|
|
|
return replace(cpg, run_batch_once=batch, unparsed=("app.py",))
|
|
|
|
result = run_fake(tmp_path, monkeypatch, LargeBackend)
|
|
assert [u["status"] for u in result["units"]] == ["path", "asked-nothing"]
|
|
for name in ("census.sc", "taint.sc"):
|
|
assert '"file_names" -> ujson.Arr.from(cpg.file.name.l)' in (trace.ROOT / "src/openultrasast/cpg/queries" / name).read_text()
|
|
|
|
|
|
def test_absent_label_includes_frontend_parse_log(tmp_path, monkeypatch):
|
|
class MissingBackend(FakeBackend):
|
|
def build(self, root, language):
|
|
cpg = super().build(root, language)
|
|
self.stages[0]["labelled_file_logs"] = {"app.py": ["Failed to parse app.py: invalid syntax"]}
|
|
|
|
def batch(query, requests):
|
|
self.asked.extend(requests)
|
|
return {**dict.fromkeys(requests, []), "__census__": [{"files": "1", "file_names": ["<unknown>"]}]}
|
|
|
|
return replace(cpg, run_batch_once=batch)
|
|
|
|
result = run_fake(tmp_path, monkeypatch, MissingBackend)
|
|
assert all(u["status"] == "failed" and "Failed to parse app.py: invalid syntax" in u["reason"] for u in result["units"])
|
|
|
|
|
|
def test_frontend_retains_matching_log_lines(tmp_path):
|
|
launcher = tmp_path / "pysrc2cpg"
|
|
launcher.write_text(f"#!{sys.executable}\nprint('Failed to parse /case/pkg/app.py: syntax')\n")
|
|
launcher.chmod(0o755)
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
backend.labelled_files = ["pkg/app.py"]
|
|
backend._run([str(launcher)], timeout=10)
|
|
assert backend.stages[0]["labelled_file_logs"]["pkg/app.py"] == ["Failed to parse /case/pkg/app.py: syntax"]
|
|
|
|
|
|
@pytest.mark.parametrize("success", [True, False])
|
|
def test_engine_bug_retries_taint_only_once(tmp_path, monkeypatch, success):
|
|
backend = worker.MeasuredBackend(900, 120, lambda: None)
|
|
calls = []
|
|
|
|
def batch(self, graph, query, requests):
|
|
calls.append((self.queries_dir / "taint.sc").read_text())
|
|
if len(calls) == 1:
|
|
self.fatal = worker.ENGINE_BUG
|
|
self.asked.append("stale")
|
|
self.answers["stale"] = []
|
|
return None
|
|
assert not self.asked and not self.answers
|
|
assert "enhance = false" in calls[-1]
|
|
assert "requires saved dataflowOss overlay" in calls[-1]
|
|
if success:
|
|
return {"a": [], "__census__": [{"files": "1", "file_names": ["app.py"]}]}
|
|
self.fatal = "joern query exited 1"
|
|
return None
|
|
|
|
monkeypatch.setattr(JoernBackend, "_batch_once", batch)
|
|
result = backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}})
|
|
assert len(calls) == 2 and backend.retries[0]["completed"] == success
|
|
assert (result is not None) == success
|
|
assert backend.fatal == ("" if success else worker.ENGINE_BUG)
|
|
if not success:
|
|
backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}})
|
|
assert len(backend.retries) == 1
|
|
|
|
|
|
def test_named_engine_bug_from_process_output(tmp_path):
|
|
launcher = tmp_path / "joern"
|
|
launcher.write_text(
|
|
f"#!{sys.executable}\nprint('Pass io.joern.x2cpg.frontendspecific.jssrc2cpg.ObjectPropertyCallLinker failed')\n"
|
|
"print('java.lang.RuntimeException: Assignment statement with 3 arguments')\nraise SystemExit(1)\n"
|
|
)
|
|
launcher.chmod(0o755)
|
|
backend = worker.MeasuredBackend(30, 10, lambda: None)
|
|
backend._run([str(launcher), "--script", "query.sc"], timeout=10)
|
|
assert backend.fatal == worker.ENGINE_BUG
|
|
|
|
|
|
def test_rerun_failed_and_timeout_only_preserves_completed_v2(tmp_path):
|
|
units = [unit(status, pin=status) for status in trace.STATUSES]
|
|
for pin in trace.make_plan(units, tmp_path):
|
|
trace.save_pin(tmp_path / trace.pin_name(pin), {**pin, "done": True, "units": [trace.unit_record(pin["units"][0], pin["pin"])]})
|
|
assert {p["pin"] for p in trace.make_plan(units, tmp_path, rerun_status={"failed", "timeout"})} == {"failed", "timeout"}
|
|
|
|
|
|
def test_successful_retry_keeps_result_and_records_attempt(tmp_path, monkeypatch):
|
|
class RetriedBackend(FakeBackend):
|
|
def __init__(self, *args):
|
|
super().__init__(*args)
|
|
self.retries = [{"reason": worker.ENGINE_BUG, "mode": "taint-only", "completed": True}]
|
|
self.stages.append({"command": "joern", "seconds": 20, "exit": 1, "output_tail": worker.ENGINE_BUG})
|
|
|
|
result = run_fake(tmp_path, monkeypatch, RetriedBackend)
|
|
assert [u["status"] for u in result["units"]] == ["path", "asked-nothing"]
|
|
assert all(u["reason"] == "" for u in result["units"])
|
|
assert result["instrument"]["retries"][0]["completed"]
|
|
|
|
|
|
def test_overlay_engine_bug_allows_one_taint_only_attempt(tmp_path, monkeypatch):
|
|
backend = worker.MeasuredBackend(900, 120, lambda: None)
|
|
|
|
def overlay(self, graph, scratch):
|
|
self.fatal = worker.ENGINE_BUG
|
|
return False
|
|
|
|
monkeypatch.setattr(JoernBackend, "_apply_overlays", overlay)
|
|
assert not backend._apply_overlays(tmp_path / "cpg.bin", tmp_path)
|
|
assert backend.retry_pending and not backend.fatal
|
|
|
|
def batch(self, graph, query, requests):
|
|
assert "enhance = false" in (self.queries_dir / "taint.sc").read_text()
|
|
return {"a": []}
|
|
|
|
monkeypatch.setattr(JoernBackend, "_batch_once", batch)
|
|
assert backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}}) == {"a": []}
|
|
assert len(backend.retries) == 1
|
|
|
|
|
|
def k8s_args(tmp_path, **overrides):
|
|
from types import SimpleNamespace
|
|
|
|
return SimpleNamespace(
|
|
**{
|
|
"executor": "k8s",
|
|
"parallel": 2,
|
|
"kube_context": "test-context",
|
|
"kube_namespace": "ousast-engine",
|
|
"image": "ghcr.io/norandom/openultrasast:2.0.1",
|
|
"deadline": 900,
|
|
"question_deadline": 120,
|
|
"keep_jobs": False,
|
|
"out": tmp_path,
|
|
**overrides,
|
|
}
|
|
)
|
|
|
|
|
|
def test_kubernetes_manifests_keep_urls_in_secret(tmp_path):
|
|
from engine_trace_k8s import BOOTSTRAP, manifests
|
|
|
|
args = k8s_args(tmp_path)
|
|
urls = {name: "https://store/" + name + "?signature=secret" for name in ("SOURCE_URL", "ANALYZER_URL", "RESULT_URL")}
|
|
job, secret = manifests("job", "run", args, urls, {})
|
|
spec = job["spec"]
|
|
pod = spec["template"]["spec"]
|
|
container = pod["containers"][0]
|
|
assert spec["activeDeadlineSeconds"] == 900
|
|
assert spec["backoffLimit"] == 0 and spec["ttlSecondsAfterFinished"] > 0
|
|
assert pod["restartPolicy"] == "Never" and not pod["automountServiceAccountToken"]
|
|
assert pod["affinity"]["podAntiAffinity"]["requiredDuringSchedulingIgnoredDuringExecution"]
|
|
assert container["image"] == args.image
|
|
assert pod["securityContext"] == {"runAsUser": 1000, "runAsGroup": 1000, "fsGroup": 1000}
|
|
assert {m["mountPath"] for m in container["volumeMounts"]} == {"/case", "/frozen", "/out"}
|
|
assert all(v["emptyDir"] == {} for v in pod["volumes"])
|
|
assert container["resources"] == {"requests": {"memory": "3Gi", "cpu": "1"}, "limits": {"memory": "3Gi"}}
|
|
assert secret["stringData"] == urls
|
|
assert all(url not in json.dumps(job) for url in urls.values())
|
|
assert "boto3" not in BOOTSTRAP and "yaml" not in BOOTSTRAP
|
|
compile(BOOTSTRAP, "<bootstrap>", "exec")
|
|
|
|
|
|
@pytest.mark.parametrize("failure", [False, True])
|
|
def test_kubernetes_host_pin_success_failure_and_cleanup(tmp_path, monkeypatch, failure):
|
|
import io
|
|
import tarfile
|
|
|
|
from engine_trace_k8s import Kubernetes
|
|
|
|
from openultrasast.plane.memory import S3Store
|
|
|
|
objects, calls, expiries = {}, [], []
|
|
|
|
class Client:
|
|
def put(self, key, data, labels):
|
|
objects[key] = data
|
|
|
|
def get(self, key):
|
|
return (objects[key], None) if key in objects else None
|
|
|
|
def delete(self, key):
|
|
objects.pop(key, None)
|
|
|
|
def presign(self, method, key, expires):
|
|
expiries.append(expires.total_seconds())
|
|
return "https://store/" + key + "?secret"
|
|
|
|
store = object.__new__(S3Store)
|
|
store.client, store.prefix = Client(), "prefix"
|
|
args = k8s_args(tmp_path)
|
|
transport = Kubernetes(args, store)
|
|
args.kubernetes = transport
|
|
pin = trace.make_plan([unit()], tmp_path)[0]
|
|
source = tmp_path / "frozen"
|
|
source.mkdir()
|
|
(source / "worker.py").write_text("print(42)")
|
|
|
|
class Inputs:
|
|
def materialize(self, pin, target):
|
|
target.mkdir()
|
|
(target / "app.py").write_text("def run(): pass")
|
|
return {"bytes": 15, "files": 1}
|
|
|
|
def kubectl(command, **kwargs):
|
|
assert command[:5] == ["kubectl", "--context", "test-context", "--namespace", "ousast-engine"]
|
|
calls.append(command[5:])
|
|
action = command[5:]
|
|
stdout = "{}"
|
|
if action[0] == "create":
|
|
manifest = json.loads(kwargs["input"])
|
|
if manifest["kind"] == "Secret":
|
|
assert manifest["metadata"]["ownerReferences"][0]["uid"] == "job-uid"
|
|
if manifest["kind"] == "Job":
|
|
stdout = json.dumps({"metadata": {"uid": "job-uid"}})
|
|
env = manifest["spec"]["template"]["spec"]["containers"][0]["env"]
|
|
prepared = json.loads(next(e["value"] for e in env if e["name"] == "PIN"))
|
|
result = {**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing", "") for u in prepared["units"]]}
|
|
result_bytes = json.dumps(result).encode()
|
|
stream = io.BytesIO()
|
|
with tarfile.open(fileobj=stream, mode="w") as archive:
|
|
entry = tarfile.TarInfo("result.json")
|
|
entry.size = len(result_bytes)
|
|
archive.addfile(entry, io.BytesIO(result_bytes))
|
|
source_key = next(key for key in objects if key.endswith("source.tar"))
|
|
with tarfile.open(fileobj=io.BytesIO(objects[source_key])) as archive:
|
|
assert archive.extractfile("app.py").read() == b"def run(): pass"
|
|
objects[source_key.replace("source.tar", "result.tar")] = stream.getvalue()
|
|
if action[:2] == ["get", "job"]:
|
|
stdout = json.dumps(
|
|
{"status": {"conditions": [{"type": "Failed", "status": "True", "reason": "DeadlineExceeded"}]}}
|
|
if failure
|
|
else {"status": {"succeeded": 1}}
|
|
)
|
|
if action[:2] == ["get", "pods"]:
|
|
stdout = json.dumps({"items": [{"spec": {"nodeName": "worker-1"}}]})
|
|
if action[0] == "logs":
|
|
assert "--tail=40" in action
|
|
stdout = "last worker line https://store/private?secret"
|
|
return subprocess.CompletedProcess(command, 0, stdout, "")
|
|
|
|
monkeypatch.setattr(subprocess, "run", kubectl)
|
|
monkeypatch.setattr(trace, "prepare_questions", lambda *a: {"a": {}})
|
|
assert trace.run_pin(pin, source, {"head": "frozen"}, Inputs(), tmp_path, args) is failure
|
|
result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
|
|
assert result["done"] and result["executor"] == "k8s" and result["node"] == "worker-1"
|
|
assert result["analyzer"] == {"head": "frozen"}
|
|
assert result["container_exit"] == int(failure)
|
|
assert result["units"][0]["status"] == ("failed" if failure else "asked-nothing")
|
|
if failure:
|
|
assert "DeadlineExceeded" in result["units"][0]["reason"]
|
|
assert "last worker line" in result["units"][0]["reason"]
|
|
assert "?secret" not in json.dumps(result)
|
|
else:
|
|
assert trace.make_plan([unit()], tmp_path) == []
|
|
assert not objects
|
|
assert any(c[:2] == ["delete", "secret"] for c in calls)
|
|
assert any(c[:2] == ["delete", "job"] for c in calls)
|
|
assert min(expiries) >= args.deadline + 900
|
|
|
|
|
|
@pytest.mark.parametrize("executor,parallel", [("docker", 1), ("k8s", 2), ("mixed", 3)])
|
|
def test_scheduler_bounds_lanes_and_progress(tmp_path, monkeypatch, executor, parallel):
|
|
import threading
|
|
import time
|
|
from collections import Counter
|
|
|
|
active, peak = Counter(), Counter()
|
|
lock = threading.Lock()
|
|
args = k8s_args(tmp_path, executor=executor, parallel=parallel)
|
|
plan = trace.make_plan([unit(str(i), str(i)) for i in range(9)], tmp_path)
|
|
|
|
def run(pin, source, provenance, inputs, out, options):
|
|
with lock:
|
|
active[options.executor] += 1
|
|
active["total"] += 1
|
|
for key, value in active.items():
|
|
peak[key] = max(peak[key], value)
|
|
time.sleep(0.02)
|
|
trace.save_pin(
|
|
out / trace.pin_name(pin), {**pin, "done": True, "units": [worker.unit_record(u, "asked-nothing", "") for u in pin["units"]]}
|
|
)
|
|
with lock:
|
|
active[options.executor] -= 1
|
|
active["total"] -= 1
|
|
return False
|
|
|
|
monkeypatch.setattr(trace, "run_pin", run)
|
|
assert not trace.run_plan(plan, None, {}, None, args)
|
|
assert peak["total"] == parallel
|
|
if executor == "mixed":
|
|
assert peak["docker"] == 1 and peak["k8s"] == parallel - 1
|
|
progress = json.loads((tmp_path / "progress.json").read_text())
|
|
assert progress["pins_done"] == 9 and progress["pins_in_flight"] == 0
|
|
assert sum(progress["lanes"].values()) == parallel
|
|
assert trace.make_plan([unit(str(i), str(i)) for i in range(9)], tmp_path) == []
|
|
|
|
|
|
@pytest.mark.parametrize("keep", [False, True])
|
|
def test_k8s_check_roundtrip_and_retention(tmp_path, monkeypatch, capsys, keep):
|
|
from engine_trace_k8s import Kubernetes
|
|
|
|
from openultrasast.plane.memory import S3Store
|
|
|
|
data, manifests_seen, commands = {}, [], []
|
|
|
|
class Client:
|
|
def put(self, key, body, labels):
|
|
data[key] = body
|
|
|
|
def get(self, key):
|
|
return (b"x" * 1024, None)
|
|
|
|
def delete(self, key):
|
|
data.pop(key, None)
|
|
|
|
def presign(self, method, key, expires):
|
|
return "https://private/" + key
|
|
|
|
store = object.__new__(S3Store)
|
|
store.client, store.prefix = Client(), ""
|
|
runner = Kubernetes(k8s_args(tmp_path, keep_jobs=keep), store)
|
|
|
|
def fake(command, **kwargs):
|
|
commands.append(command)
|
|
action = command[5:]
|
|
result = {}
|
|
if kwargs.get("input"):
|
|
manifest = json.loads(kwargs["input"])
|
|
manifests_seen.append(manifest)
|
|
result = {"metadata": {"uid": "check-uid"}}
|
|
if action[:2] == ["get", "nodes"]:
|
|
result = {"items": [{"metadata": {"name": "node"}, "status": {"allocatable": {"memory": "6000000Ki"}}}]}
|
|
if action[:2] == ["get", "pods"]:
|
|
result = {"items": [{"spec": {"nodeName": "node"}}]}
|
|
if action[:2] == ["get", "job"]:
|
|
result = {"status": {"succeeded": 1}}
|
|
return subprocess.CompletedProcess(command, 0, json.dumps(result), "")
|
|
|
|
monkeypatch.setattr(subprocess, "run", fake)
|
|
assert runner.check() == 0
|
|
text = capsys.readouterr().out
|
|
assert "6000000Ki" in text and "1024" in text and "passed" in text
|
|
assert "https://" not in text
|
|
assert any(m["kind"] == "Namespace" for m in manifests_seen)
|
|
job = next(m for m in manifests_seen if m["kind"] == "Job")
|
|
assert {"name": "CHECK", "value": "1"} in job["spec"]["template"]["spec"]["containers"][0]["env"]
|
|
assert bool(data) is keep
|
|
assert any("delete" in c for c in commands) is not keep
|
|
|
|
|
|
@pytest.mark.parametrize("outcome", ["success", "timeout", "failure"])
|
|
def test_engine_queue_pin_persistence_resume_and_progress(tmp_path, monkeypatch, outcome):
|
|
from engine_trace_queue import Dispatcher
|
|
from engine_worker_service import Queue
|
|
from test_engine_queue import FakeS3
|
|
|
|
store = FakeS3()
|
|
args = k8s_args(tmp_path, executor="queue", queue_timeout=2, image="image@sha256:abc")
|
|
plan = trace.make_plan([unit()], tmp_path)
|
|
|
|
class Inputs:
|
|
def materialize(self, pin, target):
|
|
target.mkdir()
|
|
(target / "app.py").write_text("def run(): pass")
|
|
return {"files": 1, "bytes": 15}
|
|
|
|
def execute(client, task):
|
|
return {
|
|
**task["pin"],
|
|
"done": True,
|
|
"container_exit": 0,
|
|
"units": [worker.unit_record(u, "asked-nothing") for u in task["pin"]["units"]],
|
|
}
|
|
|
|
queue = Queue(store, "pod-worker")
|
|
if outcome == "timeout":
|
|
ticks = iter([0, 3])
|
|
args.dispatcher = Dispatcher(args, store, clock=lambda: next(ticks))
|
|
else:
|
|
args.dispatcher = Dispatcher(args, store, pause=lambda _: queue.process_one(execute))
|
|
if outcome == "failure":
|
|
|
|
def fail(*_):
|
|
raise RuntimeError("credential-must-not-appear")
|
|
|
|
monkeypatch.setattr(store, "put", fail)
|
|
monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
|
|
monkeypatch.setattr(subprocess, "run", lambda *_a, **_k: pytest.fail("queue must not launch Docker/kubectl"))
|
|
assert trace.run_plan(plan, None, {"image": args.image}, Inputs(), args) is (outcome != "success")
|
|
result = json.loads((tmp_path / trace.pin_name(plan[0])).read_text())
|
|
assert result["executor"] == "queue"
|
|
assert result["units"][0]["status"] == {"success": "asked-nothing", "timeout": "timeout", "failure": "failed"}[outcome]
|
|
assert "credential-must-not-appear" not in json.dumps(result)
|
|
if outcome == "success":
|
|
assert result["worker"] == "pod-worker"
|
|
assert trace.make_plan([unit()], tmp_path) == []
|
|
else:
|
|
assert trace.make_plan([unit()], tmp_path, rerun_status=("failed", "timeout"))
|
|
progress = json.loads((tmp_path / "progress.json").read_text())
|
|
assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
|
|
|
|
|
|
@pytest.mark.parametrize("executor", ["ax", "mixed-ax"])
|
|
def test_ax_scheduler_requeues_oom_to_single_vm_lane(tmp_path, monkeypatch, executor):
|
|
import threading
|
|
import time
|
|
from types import SimpleNamespace
|
|
|
|
plan = trace.make_plan([unit(str(i), str(i)) for i in range(7)], tmp_path)
|
|
args = SimpleNamespace(executor=executor, parallel=2, out=tmp_path, image="cluster", docker_image="vm")
|
|
lock = threading.Lock()
|
|
active = Counter()
|
|
peak = Counter()
|
|
attempts = []
|
|
|
|
def run(pin, source, provenance, inputs, out, options):
|
|
with lock:
|
|
active[options.executor] += 1
|
|
peak[options.executor] = max(peak[options.executor], active[options.executor])
|
|
attempts.append((pin["pin"], options.executor, options.image))
|
|
time.sleep(0.01)
|
|
oom = options.executor == "ax"
|
|
record = {
|
|
**pin,
|
|
"done": True,
|
|
"units": [worker.unit_record(u, "failed" if oom else "path", "JVM out of memory" if oom else "") for u in pin["units"]],
|
|
}
|
|
if getattr(options, "fallback", None):
|
|
record["fallback"] = options.fallback
|
|
trace.write_json(out / trace.pin_name(pin), record)
|
|
with lock:
|
|
active[options.executor] -= 1
|
|
return oom
|
|
|
|
monkeypatch.setattr(trace, "run_pin", run)
|
|
assert not trace.run_plan(plan, None, {}, None, args)
|
|
assert peak["ax"] <= 2 and peak["docker"] == 1
|
|
for pin in plan:
|
|
record = json.loads((tmp_path / trace.pin_name(pin)).read_text())
|
|
if (pin["pin"], "ax", "cluster") in attempts:
|
|
assert record["fallback"] == "vm-oom"
|
|
assert (pin["pin"], "docker", "vm") in attempts
|
|
progress = json.loads((tmp_path / "progress.json").read_text())
|
|
assert progress["pins_done"] == 7
|
|
assert progress["statuses"] == {"path": 7}
|
|
assert progress["pins_in_flight"] == 0
|
|
|
|
|
|
@pytest.mark.parametrize("timeout", [False, True])
|
|
def test_ax_run_pin_persists_worker_ip_timeout_and_resume(tmp_path, monkeypatch, timeout):
|
|
pin = trace.make_plan([unit()], tmp_path)[0]
|
|
calls = []
|
|
|
|
class Inputs:
|
|
def materialize(self, prepared, target):
|
|
target.mkdir()
|
|
(target / "app.py").write_text("def run(): pass")
|
|
return {"files": 1, "bytes": 15}
|
|
|
|
class Dispatcher:
|
|
def execute(self, checkout, output, prepared, name):
|
|
assert (checkout / "app.py").read_bytes() == b"def run(): pass"
|
|
assert prepared["questions"] == {"a": {}}
|
|
assert prepared["materialization"]["bytes"] == 15
|
|
calls.append(name)
|
|
if timeout:
|
|
return None, "10.0.0.12"
|
|
trace.write_json(
|
|
output / "result.json",
|
|
{**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing") for u in prepared["units"]]},
|
|
)
|
|
return subprocess.CompletedProcess([], 0, "", ""), "10.0.0.12"
|
|
|
|
args = k8s_args(tmp_path, executor="ax", ax_dispatcher=Dispatcher(), image="registry/engine@sha256:" + "a" * 64)
|
|
monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
|
|
monkeypatch.setattr(subprocess, "run", lambda *_a, **_k: pytest.fail("AX run_pin must use dispatcher only"))
|
|
assert trace.run_pin(pin, None, {"host": "unused"}, Inputs(), tmp_path, args) is timeout
|
|
result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
|
|
assert calls == [trace.pin_name(pin)[:-5]]
|
|
assert result["done"] and result["executor"] == "ax"
|
|
assert result["worker_ip"] == result["worker"] == "10.0.0.12"
|
|
assert result["analyzer"] == {"image": args.image, "code": "baked-in-image"}
|
|
assert result["units"][0]["status"] == ("timeout" if timeout else "asked-nothing")
|
|
assert result["container_exit"] == (None if timeout else 0)
|
|
assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
|
|
assert bool(trace.make_plan([unit()], tmp_path, rerun_status=("timeout",))) is timeout
|
|
|
|
|
|
@pytest.mark.parametrize("executor", ["ax", "mixed-ax"])
|
|
def test_ax_pending_fallback_resumes_in_docker_and_retains_attempt(tmp_path, monkeypatch, executor):
|
|
pin = trace.make_plan([unit()], tmp_path)[0]
|
|
attempt = {
|
|
**pin,
|
|
"done": True,
|
|
"executor": "ax",
|
|
"worker_ip": "10.0.0.12",
|
|
"fallback": "vm-oom",
|
|
"fallback_pending": True,
|
|
"units": [worker.unit_record(u, "failed", "JVM out of memory") for u in pin["units"]],
|
|
}
|
|
trace.write_json(tmp_path / trace.pin_name(pin), attempt)
|
|
resumed = trace.make_plan([unit()], tmp_path, rerun_status=())
|
|
assert len(resumed) == 1
|
|
calls = []
|
|
|
|
class Inputs:
|
|
def materialize(self, prepared, target):
|
|
target.mkdir()
|
|
(target / "app.py").write_text("def run(): pass")
|
|
return {"files": 1, "bytes": 15}
|
|
|
|
class Dispatcher:
|
|
def execute(self, *_):
|
|
pytest.fail("pending OOM fallback must not dispatch AX again")
|
|
|
|
def docker(command, **kwargs):
|
|
calls.append(command)
|
|
assert command[:2] in (["docker", "run"], ["docker", "rm"])
|
|
if command[1] == "run":
|
|
assert command[command.index("--heap-profile") + 1] == "vm"
|
|
assert "vm-image" in command
|
|
output = Path(next(arg.removesuffix(":/out") for arg in command if arg.endswith(":/out")))
|
|
prepared = json.loads((output / "pin.json").read_text())
|
|
trace.write_json(
|
|
output / "result.json",
|
|
{**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing") for u in prepared["units"]]},
|
|
)
|
|
return subprocess.CompletedProcess(command, 0, "", "")
|
|
|
|
args = k8s_args(tmp_path, executor=executor, ax_dispatcher=Dispatcher(), docker_image="vm-image")
|
|
monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
|
|
monkeypatch.setattr(subprocess, "run", docker)
|
|
assert not trace.run_plan(resumed, tmp_path / "frozen", {"code": "frozen"}, Inputs(), args)
|
|
result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
|
|
assert [command[1] for command in calls] == ["run", "rm"]
|
|
assert result["executor"] == "docker" and result["image"] == "vm-image"
|
|
assert result["fallback"] == "vm-oom" and not result["fallback_pending"]
|
|
assert result["ax_attempt"] == attempt
|
|
assert result["units"][0]["status"] == "asked-nothing"
|
|
assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
|
|
progress = json.loads((tmp_path / "progress.json").read_text())
|
|
assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
|
|
assert progress["fallbacks_pending"] == 0 and progress["statuses"] == {"asked-nothing": 1}
|