File size: 6,789 Bytes
646ccfc | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """Build compact tree-sitter def-use/reachability features for TRAIN + VALID."""
from __future__ import annotations
import argparse
import json
import signal
from collections import Counter, deque
from multiprocessing import Pool
from pathlib import Path
import pandas as pd
from src.cpg.defuse import build_defuse_graph
from src.cpg.trace_miner import _extract_param_names, _known_dangerous_call_names, is_source_candidate
TIMEOUT_S = 15
TYPE_NAMES = ("call_expression", "assignment_expression", "declaration", "init_declarator",
"return_statement", "if_statement", "for_statement", "while_statement",
"switch_statement", "expression_statement", "binary_expression")
DANGEROUS_CALL_NAMES = _known_dangerous_call_names(None)
class RowTimeout(Exception):
pass
def _alarm(signum, frame):
raise RowTimeout()
def feature_row(task):
rec_id, pair_id, target, code = task
signal.signal(signal.SIGALRM, _alarm)
signal.alarm(TIMEOUT_S)
try:
graph = build_defuse_graph(code)
params = _extract_param_names(code)
nodes = [s.node_id for s in graph.statements]
node_set = set(nodes)
out = {n: [] for n in nodes}
indegree = Counter()
for a, b in graph.edges:
if a in node_set and b in node_set:
out[a].append(b)
indegree[b] += 1
sources = {n for n in nodes if is_source_candidate(graph, n, params)}
distances = {n: 0 for n in sources}
queue = deque(sources)
while queue:
node = queue.popleft()
for nxt in out[node]:
if nxt not in distances:
distances[nxt] = distances[node] + 1
queue.append(nxt)
dangerous_sinks = set()
call_count = 0
for stmt in graph.statements:
if stmt.node_type == "call_expression":
call_count += 1
if graph.uses.get(stmt.node_id, set()) & DANGEROUS_CALL_NAMES:
dangerous_sinks.add(stmt.node_id)
reachable_dangerous = [distances[n] for n in dangerous_sinks if n in distances]
n_nodes, n_edges = len(nodes), len(graph.edges)
degrees_out = [len(out[n]) for n in nodes]
degrees_in = [indegree[n] for n in nodes]
n_def, n_use = sum(bool(v) for v in graph.defines.values()), sum(bool(v) for v in graph.uses.values())
features = {
"n_statements": n_nodes, "n_defuse_edges": n_edges,
"edges_per_statement": n_edges / max(1, n_nodes),
"n_sources": len(sources), "source_fraction": len(sources) / max(1, n_nodes),
"n_reachable_statements": len(distances),
"reachable_fraction": len(distances) / max(1, n_nodes),
"n_calls": call_count, "n_dangerous_call_candidates": len(dangerous_sinks),
"n_reachable_dangerous_calls": len(reachable_dangerous),
"dangerous_reachability_fraction": len(reachable_dangerous) / max(1, len(dangerous_sinks)),
"has_source_to_dangerous_call": int(bool(reachable_dangerous)),
"shortest_source_dangerous_distance": min(reachable_dangerous) if reachable_dangerous else 0,
"mean_source_dangerous_distance": sum(reachable_dangerous) / max(1, len(reachable_dangerous)),
"n_defining_statements": n_def, "n_using_statements": n_use,
"defined_fraction": n_def / max(1, n_nodes), "used_fraction": n_use / max(1, n_nodes),
"max_in_degree": max(degrees_in, default=0), "mean_in_degree": sum(degrees_in) / max(1, n_nodes),
"max_out_degree": max(degrees_out, default=0), "mean_out_degree": sum(degrees_out) / max(1, n_nodes),
"n_control_dependent": len(graph.control_parent),
}
type_counts = Counter(s.node_type for s in graph.statements)
for name in TYPE_NAMES:
features[f"type_{name}"] = type_counts[name]
return {"sample_id": rec_id, "pair_id": pair_id, "target": target,
"status": "ok", "features": features}
except RowTimeout:
return {"sample_id": rec_id, "pair_id": pair_id, "target": target, "status": "timeout"}
except Exception as exc:
return {"sample_id": rec_id, "pair_id": pair_id, "target": target,
"status": "error", "error": f"{type(exc).__name__}: {exc}"}
finally:
signal.alarm(0)
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--corpus", type=Path, default=Path("data/processed/corpus.parquet"))
p.add_argument("--train", type=Path, default=Path("data/splits/chronological/train.json"))
p.add_argument("--valid", type=Path, default=Path("data/splits/chronological/valid.json"))
p.add_argument("--output", type=Path, default=Path("results/phase7/graph_feature_table.jsonl"))
p.add_argument("--workers", type=int, default=8)
a = p.parse_args()
train_ids, valid_ids = set(json.loads(a.train.read_text())), set(json.loads(a.valid.read_text()))
if train_ids & valid_ids:
raise ValueError("TRAIN/VALID pair overlap")
df = pd.read_parquet(a.corpus, columns=["sample_id", "pair_id", "target", "func_vuln", "func_patch"])
df = df[df.pair_id.isin(train_ids | valid_ids)]
tasks = []
for row in df.to_dict("records"):
code = row["func_vuln"] if int(row["target"]) == 1 else row["func_patch"]
tasks.append((row["sample_id"], row["pair_id"], int(row["target"]), str(code)))
a.output.parent.mkdir(parents=True, exist_ok=True)
results = []
with Pool(max(1, a.workers)) as pool:
for i, rec in enumerate(pool.imap_unordered(feature_row, tasks, chunksize=5), 1):
results.append(rec)
if i % 500 == 0:
print(f"built {i}/{len(tasks)}", flush=True)
results.sort(key=lambda x: (x["pair_id"], x["target"]))
with a.output.open("w") as f:
for rec in results:
f.write(json.dumps(rec, sort_keys=True) + "\n")
counts = Counter(r["status"] for r in results)
report = {"n_rows": len(results), "n_pairs": len(train_ids | valid_ids),
"train_pairs": len(train_ids), "valid_pairs": len(valid_ids),
"statuses": dict(counts), "timeout_seconds": TIMEOUT_S,
"features_per_row": len(next((r["features"] for r in results if r["status"] == "ok"), {})),
"test_read": False}
report_path = a.output.with_suffix(".report.json")
report_path.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if counts["timeout"] + counts["error"] > 0.02 * len(results):
raise RuntimeError("graph feature failure rate exceeds 2%; do not train classifier")
if __name__ == "__main__":
main()
|