mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
Marius CiepluchandClaude Opus 5.5
41f56170bb
The worker sent tuple params as JSON arrays; taint.sc expects comma-joined strings (ujson InvalidData, exit 1). Implemented with Codex. Co-Authored-By: Claude Opus 5.5 <[email protected]>
384 lines
18 KiB
Python
384 lines
18 KiB
Python
"""Container-only Joern trace worker; host prepares all corpus and question metadata.
|
|
|
|
Keep this module independent of learn/plane and their optional dependencies. The
|
|
frozen host invokes this file directly in the shipped, offline analyzer image.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import selectors
|
|
import signal
|
|
import subprocess
|
|
import time
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
from openultrasast.cpg.backend import JoernBackend, _render
|
|
from openultrasast.model.trace import parse_trace
|
|
|
|
|
|
def write_json(path, value):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_suffix(".tmp")
|
|
tmp.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
|
|
tmp.replace(path)
|
|
|
|
|
|
def safe_path(root, name):
|
|
path = (root / name).resolve()
|
|
if not path.is_relative_to(root.resolve()):
|
|
raise ValueError(f"path escapes input root: {name}")
|
|
return path
|
|
|
|
|
|
FRONTENDS = {"php2cpg", "pysrc2cpg", "jssrc2cpg", "javasrc2cpg", "c2cpg", "joern-parse"}
|
|
MIN_CPG_BYTES = 1024
|
|
HEAP_MB = 2560
|
|
EXCLUDED_DIRS = frozenset({"tests", "test", "docs", "node_modules", "vendor", "dist", "build", "templates", "migrations"})
|
|
|
|
|
|
def instrument_failure(proof):
|
|
if any(stage.get("oom") for stage in proof.get("jvm", [])):
|
|
return "JVM out of memory"
|
|
if not proof.get("files") or not proof.get("bytes"):
|
|
return "instrument read zero files or zero bytes"
|
|
if proof.get("cpg_bytes", 0) <= MIN_CPG_BYTES:
|
|
return "frontend wrote no CPG"
|
|
for stage in proof.get("jvm", []):
|
|
if stage["command"] in FRONTENDS:
|
|
if stage.get("cpg_bytes", 0) <= MIN_CPG_BYTES:
|
|
return "frontend wrote no CPG"
|
|
elif stage["command"] == "joern" and stage["seconds"] < 5:
|
|
return f"implausibly fast JVM step: {stage['command']} in {stage['seconds']:.3f}s (<5s)"
|
|
if not proof.get("jvm"):
|
|
return "no JVM step recorded"
|
|
return ""
|
|
|
|
|
|
class MeasuredBackend(JoernBackend):
|
|
"""Use the shipped backend with a process watchdog, no concurrent query threads."""
|
|
|
|
def __init__(self, deadline, question_deadline, checkpoint):
|
|
super().__init__(build_timeout=deadline, query_timeout=deadline, session_transport=False, heap_mb=HEAP_MB)
|
|
self.deadline = time.monotonic() + deadline
|
|
self.question_deadline = question_deadline
|
|
self.checkpoint = checkpoint
|
|
self.stages = []
|
|
self.asked = []
|
|
self.answers = {}
|
|
self.completion_counts = Counter()
|
|
self.timed_out = False
|
|
self.fatal = ""
|
|
|
|
def _run(self, command, *, timeout, cwd=None):
|
|
if self.fatal or self.timed_out:
|
|
return None
|
|
if time.monotonic() >= self.deadline:
|
|
self.timed_out = True
|
|
return None
|
|
started = time.monotonic()
|
|
end = min(self.deadline, started + timeout)
|
|
question_end = end
|
|
binary = Path(command[0]).name
|
|
frontend = binary in FRONTENDS
|
|
query = binary == "joern" and "--script" in command
|
|
output = bytearray()
|
|
pending = b""
|
|
process = subprocess.Popen(
|
|
command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd, env=self._jvm_env(), start_new_session=True
|
|
)
|
|
try:
|
|
with selectors.DefaultSelector() as selector:
|
|
selector.register(process.stdout, selectors.EVENT_READ)
|
|
while selector.get_map():
|
|
if time.monotonic() >= min(end, question_end):
|
|
self.timed_out = True
|
|
break
|
|
for key, _ in selector.select(min(0.2, max(0, min(end, question_end) - time.monotonic()))):
|
|
chunk = os.read(key.fileobj.fileno(), 65536)
|
|
if not chunk:
|
|
selector.unregister(key.fileobj)
|
|
continue
|
|
output.extend(chunk)
|
|
pending += chunk
|
|
while b"\n" in pending:
|
|
line, pending = pending.split(b"\n", 1)
|
|
try:
|
|
event = json.loads(line)
|
|
except ValueError:
|
|
continue
|
|
if not isinstance(event, dict):
|
|
continue
|
|
if "__question__" in event:
|
|
self.asked.append(event["__question__"])
|
|
question_end = min(end, time.monotonic() + self.question_deadline)
|
|
if "id" in event and "rows" in event:
|
|
self.answers[event["id"]] = event["rows"]
|
|
self.completion_counts[event["id"]] += 1
|
|
question_end = end
|
|
self.checkpoint()
|
|
if not self.timed_out:
|
|
process.wait(timeout=max(0.01, end - time.monotonic()))
|
|
finally:
|
|
with contextlib.suppress(ProcessLookupError):
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
process.wait(timeout=5)
|
|
process.stdout.close()
|
|
elapsed = time.monotonic() - started
|
|
if frontend or query:
|
|
stage = {
|
|
"command": binary,
|
|
"argv": command,
|
|
"heap_mb": self._heap_mb(),
|
|
"java_tool_options": self._jvm_env()["JAVA_TOOL_OPTIONS"],
|
|
"seconds": elapsed,
|
|
"exit": process.returncode,
|
|
"oom": b"OutOfMemoryError" in output,
|
|
}
|
|
self.stages.append(stage)
|
|
if query and process.returncode != 0:
|
|
# stderr is merged into stdout above; drop JVM frames before taking
|
|
# the tail so they cannot crowd out the exception's message.
|
|
lines = [
|
|
line.strip()
|
|
for line in output.decode(errors="replace").splitlines()
|
|
if line.strip() and not line.strip().startswith(("at ", "... "))
|
|
]
|
|
stage["output_tail"] = "\n".join(lines[-20:])
|
|
if not self.timed_out:
|
|
self.fatal = f"joern query exited {process.returncode}"
|
|
if frontend:
|
|
# Capture the exact frontend output before build failure cleanup or overlays.
|
|
output_arg = next((i + 1 for i, arg in enumerate(command[:-1]) if arg in {"-o", "--output"}), None)
|
|
cpg_path = (Path(cwd or Path.cwd()) / command[output_arg]).resolve() if output_arg is not None else None
|
|
stage["cpg_path"] = str(cpg_path) if cpg_path is not None else None
|
|
stage["cpg_bytes"] = cpg_path.stat().st_size if cpg_path is not None and cpg_path.is_file() else 0
|
|
if stage["cpg_bytes"] <= MIN_CPG_BYTES:
|
|
self.fatal = "frontend wrote no CPG"
|
|
elif elapsed < 5 and not self.fatal:
|
|
self.fatal = f"implausibly fast JVM step: {command[0]} in {elapsed:.3f}s (<5s)"
|
|
if b"OutOfMemoryError" in output:
|
|
self.fatal = "JVM out of memory"
|
|
self.checkpoint()
|
|
return subprocess.CompletedProcess(command, process.returncode if not self.fatal else 1, output.decode(errors="replace"), "")
|
|
|
|
|
|
def unit_record(unit, status, reason="", **extra):
|
|
return {
|
|
**unit,
|
|
"status": status,
|
|
"reason": reason,
|
|
"witness_rows": [],
|
|
"traces": [],
|
|
"instrument": {"files": [], "bytes": 0, "cpg_path": None, "cpg_bytes": 0, "jvm": []},
|
|
"questions": {},
|
|
"questions_asked": [],
|
|
"questions_completed": [],
|
|
**extra,
|
|
}
|
|
|
|
|
|
def worker_language(pin, root, out, deadline, question_deadline, persist=write_json):
|
|
proof = {
|
|
"files": [],
|
|
"bytes": 0,
|
|
"cpg_path": None,
|
|
"cpg_bytes": 0,
|
|
"jvm": [],
|
|
"heap_mb": HEAP_MB,
|
|
"excludes": sorted(EXCLUDED_DIRS),
|
|
"materialization": pin.get("materialization", {}),
|
|
}
|
|
record = {
|
|
**pin,
|
|
"done": False,
|
|
"instrument": proof,
|
|
"units": [unit_record(u, "unsupported" if not u["supported"] else "timeout", u["reason"]) for u in pin["units"]],
|
|
}
|
|
|
|
def checkpoint():
|
|
# Checkpoints retain raw completed rows even if the outer container deadline fires.
|
|
# They remain timeout until the finished JVM's instrument proof is validated.
|
|
for stage in backend.stages:
|
|
if "cpg_path" in stage:
|
|
proof["cpg_path"] = stage["cpg_path"]
|
|
proof["cpg_bytes"] = stage["cpg_bytes"]
|
|
for row in record["units"]:
|
|
rid = row["unit"]
|
|
row["instrument"] = proof
|
|
row["questions_asked"] = [rid] if rid in backend.asked else []
|
|
if rid in backend.answers:
|
|
row["questions_completed"] = [rid]
|
|
row["witness_rows"] = [r for r in backend.answers[rid] if isinstance(r, dict) and "kind" not in r and "sink" in r]
|
|
persist(out, record)
|
|
|
|
backend = MeasuredBackend(deadline, question_deadline, checkpoint)
|
|
proof["jvm"] = backend.stages
|
|
cpg = None
|
|
try:
|
|
active = [u for u in pin["units"] if u["supported"]]
|
|
for file in sorted({u["file"] for u in active}):
|
|
data = safe_path(root, file).read_bytes()
|
|
proof["files"].append({"file": file, "bytes": len(data), "sha256": hashlib.sha256(data).hexdigest()})
|
|
proof["bytes"] += len(data)
|
|
print(f"read {file}: {len(data)} bytes", flush=True)
|
|
if not proof["bytes"]:
|
|
raise ValueError("instrument read zero files or zero bytes")
|
|
languages = {"javascript" if u["language"] == "typescript" else u["language"] for u in active}
|
|
if len(languages) != 1:
|
|
raise ValueError(f"one pin has multiple frontend languages: {sorted(languages)}")
|
|
language = next(iter(languages))
|
|
requests = {rid: {key: _render(value) for key, value in params.items()} for rid, params in pin["questions"].items()}
|
|
if set(requests) != {u["unit"] for u in active}:
|
|
raise ValueError("prepared questions do not match supported units")
|
|
record["questions"] = requests
|
|
for unit_row in record["units"]:
|
|
unit_row["questions"] = {unit_row["unit"]: requests[unit_row["unit"]]} if unit_row["unit"] in requests else {}
|
|
persist(out, record)
|
|
cpg = backend.build(root, language=language)
|
|
if cpg is None:
|
|
raise ValueError(backend.fatal or backend.last_failure or ("pin deadline" if backend.timed_out else "CPG build failed"))
|
|
if proof["cpg_path"] is None:
|
|
proof["cpg_path"] = str(cpg.cpg_path)
|
|
proof["cpg_bytes"] = cpg.cpg_path.stat().st_size if cpg.cpg_path.is_file() else 0
|
|
failure = instrument_failure(proof)
|
|
if failure:
|
|
raise ValueError(failure)
|
|
if cpg.run_batch_once is None:
|
|
raise ValueError("backend offers no batch operation")
|
|
answers = cpg.run_batch_once("taint", requests) or {}
|
|
record["census"] = answers.get("__census__", [])
|
|
failure = backend.fatal or instrument_failure(proof)
|
|
if not backend.asked and not backend.timed_out:
|
|
failure = failure or "query produced no markers"
|
|
census = record["census"]
|
|
if census and int(census[0].get("files", 0)) <= 0:
|
|
failure = failure or "query census read zero files"
|
|
elif not census and not backend.timed_out:
|
|
failure = failure or "query census is missing"
|
|
for index, unit in enumerate(pin["units"]):
|
|
if not unit["supported"]:
|
|
continue
|
|
rid = unit["unit"]
|
|
extra = {
|
|
"questions": {rid: requests[rid]},
|
|
"questions_asked": [rid] if rid in backend.asked else [],
|
|
"questions_completed": [rid] if rid in answers else [],
|
|
"instrument": proof,
|
|
}
|
|
rows = answers.get(rid)
|
|
# A batch across backend shards is complete only when every shard answered.
|
|
# The legacy merger can return a union even when one shard was killed.
|
|
shards = int(census[0].get("shards", 1)) if census else 1
|
|
shard_incomplete = shards > 1 and getattr(backend, "completion_counts", {}).get(rid, 0) < shards
|
|
if rows is not None:
|
|
raw_flows = [r for r in rows if isinstance(r, dict) and "kind" not in r and "sink" in r]
|
|
extra.update(witness_rows=raw_flows, traces=[parse_trace(r) for r in raw_flows])
|
|
if shard_incomplete:
|
|
extra["questions_completed"] = []
|
|
dropped = any(p.endswith(unit["file"]) for p in cpg.unparsed)
|
|
names = census[0].get("file_names", []) if census else []
|
|
file_missing = not any(str(p).removeprefix("./") == unit["file"] or str(p).endswith("/" + unit["file"]) for p in names)
|
|
if file_missing and rows is not None:
|
|
dropped = True
|
|
if failure or dropped:
|
|
status, reason = "failed", failure or "labelled file was dropped or absent from query census"
|
|
elif rows is None or shard_incomplete:
|
|
status, reason = ("timeout" if backend.timed_out else "failed"), "question did not complete"
|
|
elif not rows and rid not in backend.asked:
|
|
status, reason = "failed", "query produced no markers"
|
|
else:
|
|
flows = [r for r in rows if isinstance(r, dict) and "kind" not in r and "sink" in r]
|
|
if any(not row.get("trace") for row in flows):
|
|
raise ValueError("taint witness returned without requested trace")
|
|
extra.update(witness_rows=flows, traces=[parse_trace(r) for r in flows])
|
|
status, reason = ("path" if flows else "asked-nothing"), ""
|
|
record["units"][index] = unit_record(unit, status, reason, **extra)
|
|
except (OSError, ValueError, RuntimeError) as exc:
|
|
for index, unit in enumerate(pin["units"]):
|
|
if unit["supported"]:
|
|
record["units"][index] = unit_record(
|
|
unit,
|
|
"timeout" if backend.timed_out and not backend.fatal else "failed",
|
|
str(exc),
|
|
instrument=proof,
|
|
questions={unit["unit"]: record["questions"][unit["unit"]]} if unit["unit"] in record.get("questions", {}) else {},
|
|
witness_rows=[r for r in backend.answers.get(unit["unit"], []) if isinstance(r, dict) and "sink" in r],
|
|
questions_completed=[unit["unit"]] if unit["unit"] in backend.answers else [],
|
|
questions_asked=[unit["unit"]] if unit["unit"] in backend.asked else [],
|
|
)
|
|
finally:
|
|
if cpg and cpg.cleanup:
|
|
cpg.cleanup()
|
|
for stage in backend.stages:
|
|
if stage["command"] == "joern" and stage.get("exit") and "output_tail" in stage:
|
|
detail = f"joern query exited {stage['exit']}: {stage['output_tail'] or 'no output'}"
|
|
for row in record["units"]:
|
|
if row["supported"]:
|
|
row["reason"] = f"{row['reason']}\n{detail}".strip()
|
|
record["questions_asked"] = list(dict.fromkeys(backend.asked))
|
|
record["questions_completed"] = [rid for u in record["units"] for rid in u["questions_completed"]]
|
|
record["done"] = True
|
|
persist(out, record)
|
|
return record
|
|
|
|
|
|
def worker(pin, root, out, deadline, question_deadline):
|
|
"""Run language partitions sequentially under one shared container deadline."""
|
|
languages = sorted({"javascript" if u["language"] == "typescript" else u["language"] for u in pin["units"] if u["supported"]})
|
|
if len(languages) <= 1:
|
|
return worker_language(pin, root, out, deadline, question_deadline)
|
|
end = time.monotonic() + deadline
|
|
record = {
|
|
**pin,
|
|
"done": False,
|
|
"units": [
|
|
unit_record(u, "timeout" if u["supported"] else "unsupported", "pin deadline" if u["supported"] else u["reason"])
|
|
for u in pin["units"]
|
|
],
|
|
"instruments": {},
|
|
}
|
|
|
|
def persist(_, partial):
|
|
updates = {u["unit"]: u for u in partial["units"]}
|
|
record["units"] = [updates.get(u["unit"], u) for u in record["units"]]
|
|
record["instruments"][language] = partial["instrument"]
|
|
record["questions_asked"] = [rid for u in record["units"] for rid in u["questions_asked"]]
|
|
record["questions_completed"] = [rid for u in record["units"] for rid in u["questions_completed"]]
|
|
write_json(out, record)
|
|
|
|
write_json(out, record)
|
|
for language in languages:
|
|
remaining = end - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
units = [
|
|
u for u in pin["units"] if u["supported"] and ("javascript" if u["language"] == "typescript" else u["language"]) == language
|
|
]
|
|
questions = {u["unit"]: pin["questions"][u["unit"]] for u in units}
|
|
worker_language({**pin, "units": units, "questions": questions}, root, out, remaining, question_deadline, persist)
|
|
record["done"] = True
|
|
write_json(out, record)
|
|
return record
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--worker", type=Path, required=True)
|
|
parser.add_argument("--deadline", type=int, default=900)
|
|
parser.add_argument("--question-deadline", type=int, default=120)
|
|
args = parser.parse_args()
|
|
if min(args.deadline, args.question_deadline) <= 0:
|
|
parser.error("deadlines must be positive")
|
|
worker(json.loads(args.worker.read_text()), Path("/case"), Path("/out/result.json"), args.deadline, args.question_deadline)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|