| """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() |
|
|
