File size: 4,438 Bytes
5bfc214 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Phase 2 per-sample annotation (spec §4.2-§4.4): ties statement_mapper,
defuse, and trace_miner together to produce, for one corpus row, the L1/L2/L3
localization targets and trace ground truth.
Uses the tree-sitter fallback (defuse.py) by default — not real Joern — since
running Joern per-sample across the full ~13.6k-pair corpus is a real,
multi-hour, CPU-parallelized batch job (spec §4: "budget accordingly...
target 16 workers... resumable") that hasn't been executed yet (see
DECISIONS.md). `annotate_sample(..., use_joern=True)` is wired for when that
batch job runs; this module's job is to define *what* an annotation record
looks like and compute it correctly, independent of which backend produced
the underlying graph.
"""
from __future__ import annotations
from dataclasses import asdict, dataclass, field
from src.cpg.defuse import backward_slice, build_defuse_graph
from src.cpg.statement_mapper import map_patch_lines_to_statements_detailed, statement_lines
from src.cpg.trace_miner import mine_trace
@dataclass
class SampleAnnotation:
sample_id: str
l1_lines: list[int]
l2_lines: list[int]
l3_lines: list[int] # whole function (degenerate upper bound, spec §4.2)
trace_gt: list[list[str]] # stringified node ids (JSON-safe), spec §4.3
trace_unverifiable: bool
trace_reason: str
annotation_backend: str # "tree_sitter" | "joern"
trace_lengths: list[int] = field(default_factory=list) # for the trace-length distribution report
used_fallback_sink_search: bool = False # True if L1 was empty and the old dangerous-call search ran
n_dangerous_call_sinks: int = 0 # enrichment count, spec amendment §1.3-a
cwe: str | None = None # carried through for the per-CWE trace-coverage table
l1_mapping_reasons: list[str] = field(default_factory=list) # spec amendment §2: per-line reason log
def annotate_sample(
sample_id: str,
func_vuln: str,
patch_lines_vuln: list[int],
cwe: str | None,
func_len_lines: int | None = None,
) -> SampleAnnotation:
"""spec §4.2-§4.3 for a single sample. Only meaningful for target=1 rows
(patch_lines_vuln is empty for target=0 benign rows by construction —
spec §3.2 — so L1/L2/trace are empty/unverifiable for those, correctly)."""
if not patch_lines_vuln:
n_lines = func_len_lines or len(func_vuln.splitlines())
return SampleAnnotation(
sample_id=sample_id, l1_lines=[], l2_lines=[], l3_lines=list(range(1, n_lines + 1)),
trace_gt=[], trace_unverifiable=True, trace_reason="target_0_no_patch_lines",
annotation_backend="tree_sitter", cwe=cwe,
)
detailed = map_patch_lines_to_statements_detailed(func_vuln, patch_lines_vuln)
l1_mapping_reasons = [reason for _stmt, reason in detailed]
l1_statements = [stmt for stmt, _reason in detailed if stmt is not None]
l1_lines = sorted(statement_lines(l1_statements))
graph = build_defuse_graph(func_vuln)
l1_ids = {s.node_id for s in l1_statements}
l2_ids = backward_slice(graph, l1_ids, max_depth=3)
l2_statement_objs = [s for s in graph.statements if s.node_id in l2_ids]
l2_lines = sorted(statement_lines(l2_statement_objs))
n_lines = func_len_lines or len(func_vuln.splitlines())
l3_lines = list(range(1, n_lines + 1))
# Trace amendment: sinks = L1 statements themselves, not a dangerous-call
# table lookup (see DECISIONS.md "Phase 2 trace-mining redefinition").
trace_result = mine_trace(func_vuln, cwe, l1_stmt_ids=sorted(l1_ids), graph=graph)
# stmt ids are (start_byte, end_byte) tuples — stringify for JSON safety
trace_gt = [[f"{a}-{b}" for (a, b) in path] for path in trace_result.trace_gt]
return SampleAnnotation(
sample_id=sample_id,
l1_lines=l1_lines,
l2_lines=sorted(set(l1_lines) | set(l2_lines)),
l3_lines=l3_lines,
trace_gt=trace_gt,
trace_unverifiable=trace_result.trace_unverifiable,
trace_reason=trace_result.reason,
annotation_backend="tree_sitter",
trace_lengths=[len(p) for p in trace_result.trace_gt],
used_fallback_sink_search=trace_result.used_fallback_sink_search,
n_dangerous_call_sinks=len(trace_result.dangerous_call_sinks),
cwe=cwe,
l1_mapping_reasons=l1_mapping_reasons,
)
def annotation_to_dict(ann: SampleAnnotation) -> dict:
return asdict(ann)
|