File size: 10,071 Bytes
4951e76
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
#!/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()