Download code/bugfind/bench/translate_items.py from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 10.1 kB
-
https://huggingface.co/viai957/align/resolve/main/code/bugfind/bench/translate_items.py
- Command line
-
hf download hf://viai957/align/code/bugfind/bench/translate_items.py
-
curl -L -o translate_items.py https://huggingface.co/viai957/align/resolve/main/code/bugfind/bench/translate_items.py
10.1 kB
| #!/usr/bin/env python3 -I | |
| """Translate the affected function(s) of candidate items with Charon then Aeneas (Lean backend), and check | |
| that the generated Lean compiles against the Aeneas library. | |
| Advisory items: Charon `--start-from-if-exists` on the functions whose bodies changed between the vulnerable and | |
| the patched crate version (difffuncs.py), so the unit of translation is "the changed function plus its in-crate | |
| callees", not the whole crate (whole-crate translation of real crates fails on Aeneas limitations; see TOOLCHAIN.md). | |
| Mutant items (--mutants): the one-function crate is translated whole. | |
| Per-item output fields: | |
| charon_ok exit 0 and an .llbc was written aeneas_ok exit 0, no `[Error]` lines, .lean written | |
| fn_present every start function appears as a `def` in the Lean output (not dropped / opaque) | |
| lean_ok the generated Lean file compiles (lean --root, Aeneas lib on LEAN_PATH); only attempted if aeneas_ok | |
| translated charon_ok and aeneas_ok and fn_present and lean_ok | |
| Run after `source env.sh`. Concurrency: --jobs workers x CARGO_BUILD_JOBS=2 (default 2x2 = 4 build jobs), nice -n 15. | |
| """ | |
| import argparse, json, os, re, shutil, subprocess, sys, time | |
| from concurrent.futures import ThreadPoolExecutor | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import difffuncs | |
| TC = os.environ.get("TC", "/apps/development/rlhf/toolchains") | |
| CHARON = os.environ.get("CHARON", f"{TC}/src/aeneas/charon/bin/charon") | |
| AENEAS = os.environ.get("AENEAS", f"{TC}/src/aeneas/bin/aeneas") | |
| LEAN_DIR = f"{TC}/src/aeneas/backends/lean" | |
| ANSI = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]") | |
| _LP = None | |
| def lean_path(): | |
| global _LP | |
| if _LP is None: | |
| o = subprocess.run(["lake", "env", "printenv", "LEAN_PATH"], cwd=LEAN_DIR, capture_output=True, text=True).stdout.strip() | |
| _LP = ":".join(dict.fromkeys(o.split(":"))) | |
| return _LP | |
| def start_patterns(fns): | |
| pats = [] | |
| for _rel, mod, name in fns: | |
| base = "crate::" + (mod + "::" if mod else "") | |
| pats += [base + name, base + "_::" + name] | |
| return list(dict.fromkeys(pats)) | |
| def lean_compile(out, leanfile, timeout): | |
| env = dict(os.environ, LEAN_PATH=lean_path() + ":" + out) | |
| mod = os.path.splitext(os.path.basename(leanfile))[0] | |
| p = subprocess.run(["nice", "-n", "15", "lean", f"--root={out}", "-o", os.path.join(out, mod + ".olean"), leanfile], | |
| cwd=LEAN_DIR, env=env, capture_output=True, text=True, timeout=timeout) | |
| return p.returncode == 0, (p.stdout + p.stderr)[-300:] | |
| def one(it, work, timeout): | |
| t0 = time.time() | |
| name = it.get("adv_id") or it["item_id"] | |
| d = it["pre_dir"] | |
| out = os.path.join(work, name); shutil.rmtree(out, ignore_errors=True); os.makedirs(out) | |
| res = dict(item=name, crate=it.get("crate"), charon_ok=False, aeneas_ok=False, fn_present=False, lean_ok=False, translated=False) | |
| cmd = ["nice", "-n", "15", CHARON, "cargo", "--preset=aeneas", "--dest-file", os.path.join(out, "crate.llbc")] | |
| names = [] | |
| if "post_dir" in it: | |
| fns = difffuncs.changed_functions(d, it["post_dir"]) | |
| res["changed_functions"] = [f"{m}::{n}" if m else n for _r, m, n in fns] | |
| if not fns: | |
| res["error"] = "no_changed_function_identified"; res["seconds"] = round(time.time() - t0, 1); return res | |
| for p in start_patterns(fns): cmd += ["--start-from-if-exists", p] | |
| names = sorted({n for _r, _m, n in fns}) | |
| env = dict(os.environ, CARGO_TARGET_DIR=os.path.join(out, "target"), CARGO_BUILD_JOBS="2") | |
| res["mode"] = "start_from_changed_fns" if names else "whole_crate" | |
| try: | |
| p = subprocess.run(cmd, cwd=d, env=env, capture_output=True, text=True, timeout=timeout) | |
| log = ANSI.sub("", p.stdout + p.stderr) | |
| res["charon_ok"] = p.returncode == 0 and os.path.exists(os.path.join(out, "crate.llbc")) | |
| if not res["charon_ok"]: | |
| errs = [l for l in log.splitlines() if re.match(r"\s*(error|Error)", l)] | |
| res["error"] = ("charon: " + (errs[0] if errs else log[-300:]))[:300] | |
| except subprocess.TimeoutExpired: | |
| res["error"] = "charon timeout" | |
| if res["charon_ok"]: | |
| try: | |
| p = subprocess.run(["nice", "-n", "15", AENEAS, "-backend", "lean", "-no-progress-bar", os.path.join(out, "crate.llbc"), "-dest", out], | |
| capture_output=True, text=True, timeout=timeout) | |
| log = ANSI.sub("", p.stdout + p.stderr) | |
| leans = [f for f in os.listdir(out) if f.endswith(".lean")] | |
| nerr = len(re.findall(r"\[\s*Error\s*\]", log)) | |
| res["n_aeneas_errors"] = nerr | |
| res["aeneas_ok"] = p.returncode == 0 and nerr == 0 and bool(leans) | |
| if not res["aeneas_ok"]: | |
| errs = [l for l in log.splitlines() if "Error" in l] | |
| res["error"] = ("aeneas: " + (errs[0] if errs else log[-200:]))[:300] | |
| elif leans: | |
| src = open(os.path.join(out, leans[0]), encoding="utf8").read() | |
| res["fn_present"] = all(re.search(rf"^(?:@\[[^\]]*\]\s*)*def\s+(?:[\w.]*\.)?{re.escape(n)}(?:_loop)?\b", src, re.M) for n in names) if names else True | |
| res["lean_file"] = leans[0]; res["lean_lines"] = src.count("\n") | |
| if names and not res["fn_present"] and res["mode"] == "start_from_changed_fns": | |
| # patterns matched nothing (trait-impl methods, cfg'd-out code): retry from the whole crate | |
| res["mode"] = "whole_crate_fallback" | |
| cmd2 = ["nice", "-n", "15", CHARON, "cargo", "--preset=aeneas", "--dest-file", os.path.join(out, "crate.llbc")] | |
| subprocess.run(cmd2, cwd=d, env=env, capture_output=True, text=True, timeout=timeout) | |
| p = subprocess.run(["nice", "-n", "15", AENEAS, "-backend", "lean", "-no-progress-bar", os.path.join(out, "crate.llbc"), "-dest", out], | |
| capture_output=True, text=True, timeout=timeout) | |
| log = ANSI.sub("", p.stdout + p.stderr) | |
| res["n_aeneas_errors"] = len(re.findall(r"\[\s*Error\s*\]", log)) | |
| leans = [f for f in os.listdir(out) if f.endswith(".lean")] | |
| if leans: | |
| src = open(os.path.join(out, leans[0]), encoding="utf8").read() | |
| res["fn_present"] = all(re.search(rf"^(?:@\[[^\]]*\]\s*)*def\s+(?:[\w.]*\.)?{re.escape(n)}(?:_loop)?\b", src, re.M) for n in names) | |
| res["lean_file"] = leans[0]; res["lean_lines"] = src.count("\n") | |
| # failures elsewhere in the crate are tolerated here: Aeneas drops failing decls, and what we need is the target fn | |
| res["aeneas_ok"] = True | |
| if res["fn_present"]: | |
| ok, msg = lean_compile(out, os.path.join(out, leans[0]), timeout) | |
| res["lean_ok"] = ok | |
| if not ok: res["error"] = "lean: " + msg | |
| else: | |
| res["error"] = "start function not in generated Lean" | |
| except subprocess.TimeoutExpired: | |
| res["error"] = "aeneas/lean timeout" | |
| res["translated"] = res["charon_ok"] and res["aeneas_ok"] and res["fn_present"] and res["lean_ok"] | |
| shutil.rmtree(os.path.join(out, "target"), ignore_errors=True) # keep llbc + lean, drop build products | |
| res["seconds"] = round(time.time() - t0, 1) | |
| return res | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("items"); ap.add_argument("--out", required=True) | |
| ap.add_argument("--work", default="/apps/development/rlhf/bugfind_data/work/translate") | |
| ap.add_argument("--jobs", type=int, default=2); ap.add_argument("--timeout", type=int, default=600) | |
| ap.add_argument("--mutants", action="store_true") | |
| ap.add_argument("--stages", default="diff_safe,strict_safe"); ap.add_argument("--limit", type=int, default=0) | |
| ap.add_argument("--max-changed-fns", type=int, default=3, help="advisory items touching more functions are not 'targeted fixes'") | |
| a = ap.parse_args() | |
| for p in (a.out, a.work): | |
| if not os.path.abspath(p).startswith("/apps/"): sys.exit("paths must be under /apps") | |
| items = [json.loads(l) for l in open(a.items)] | |
| if a.mutants: # materialise each kept mutant as a one-file crate (`lib.rs` = the mutated function) | |
| items = [i for i in items if i.get("status") == "kept_nonequivalent"] | |
| for i in items: | |
| d = os.path.join(a.work, "_crates", i["item_id"]); os.makedirs(d + "/src", exist_ok=True) | |
| open(d + "/Cargo.toml", "w").write(f'[package]\nname = "m_{i["function"]}"\nversion = "0.1.0"\nedition = "2021"\n\n[workspace]\n') | |
| open(d + "/src/lib.rs", "w").write(i["mutant_source"] + "\n") | |
| i["pre_dir"], i["stage"] = d, "diff_safe" | |
| items = [i for i in items if "pre_dir" in i and i.get("stage") in a.stages.split(",")] | |
| pre_skip = [] | |
| if not a.mutants: | |
| keep = [] | |
| for i in items: | |
| n = len(difffuncs.changed_functions(i["pre_dir"], i["post_dir"])) | |
| i["n_changed_fns"] = n | |
| (keep if 1 <= n <= a.max_changed_fns else pre_skip).append(i) | |
| items = keep | |
| json.dump({"skipped_not_targeted_or_no_fn": len(pre_skip), "to_attempt": len(items)}, open(a.out + ".prefilter.json", "w")) | |
| items.sort(key=lambda i: (i.get("stage") != "strict_safe", i.get("n_changed_fns", 0))) | |
| if a.limit: items = items[: a.limit] | |
| done = {} | |
| if os.path.exists(a.out): | |
| for l in open(a.out): r = json.loads(l); done[r["item"]] = r | |
| todo = [i for i in items if (i.get("adv_id") or i["item_id"]) not in done] | |
| os.makedirs(a.work, exist_ok=True) | |
| with open(a.out, "a") as f, ThreadPoolExecutor(a.jobs) as ex: | |
| for r in ex.map(lambda i: one(i, a.work, a.timeout), todo): | |
| f.write(json.dumps(r) + "\n"); f.flush() | |
| print(r["item"], "T" if r["translated"] else "-", r.get("error", "")[:80], r["seconds"], flush=True) | |
| if __name__ == "__main__": | |
| main() | |