nordisk/heimdall

bench/znippy-unpack/enforce.py

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Heimdall is a defensive sandbox. Enforce the generated znippy-unpack profile on every run
# of its workload: each run in a rootless podman container (heim-m0, the image the binary
# was built in) whose seccomp filter is the checked-in .oci.json plus the two calls crun
# makes after loading a filter (close_range, execve; the probe starts the program itself
# and records from its first instruction). Every run must end with its expect_exit, never
# SIGSYS; then a call outside the profile (znippy compress reads a directory: getdents64)
# must be killed by SIGSYS (exit 159).
#   enforce.py            (from bench/znippy-unpack; TMPDIR for the outputs: ./tmp)
import json, os, shutil, subprocess, sys, tomllib

here = os.path.dirname(os.path.abspath(__file__))
prof = os.path.join(here, "profiles", "znippy-unpack")
w = tomllib.load(open(os.path.join(prof, "workload.toml"), "rb"))
oci = json.load(open(os.path.join(prof, "x86_64-glibc.oci.json")))
oci["syscalls"].append({"names": ["close_range", "execve"], "action": "SCMP_ACT_ALLOW"})
tmp = os.environ.get("ZNIPPY_ENFORCE_TMP", os.path.join(here, "tmp"))
os.makedirs(tmp, exist_ok=True)
filt = os.path.join(tmp, "enforce.oci.json")
json.dump(oci, open(filt, "w"))
corpus = os.path.join(prof, "corpus")
binary = os.path.join(here, "bin", "znippy")
img = os.environ.get("ZNIPPY_ENFORCE_IMAGE", "localhost/nordisk-cargo-toolchain:heim-m0")

def run(name, args):
    t = os.path.join(tmp, "run-" + name)
    shutil.rmtree(t, ignore_errors=True)
    os.makedirs(t)
    args = [a.replace("{corpus}", "/corpus").replace("{tmp}", "/work") for a in args]
    cmd = ["podman", "run", "--rm", "--network", "none", "--security-opt", "seccomp=" + filt,
           "--security-opt", "no-new-privileges", "--cap-drop", "all",
           "-v", corpus + ":/corpus:ro", "-v", t + ":/work", "-v", binary + ":/znippy:ro",
           "-w", "/work", "-e", "RUST_BACKTRACE=0", "--entrypoint", "/znippy", img] + args
    r = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    shutil.rmtree(t, ignore_errors=True)
    return r.returncode

failed = 0
for r in w["run"]:
    want = r.get("expect_exit", 0)
    got = run(r["name"], r["args"])
    ok = got == want
    failed += not ok
    print(f"test znippy-unpack x86_64 glibc {r['name']} ... {'ok' if ok else 'FAILED'} (exit {got}, want {want})", flush=True)
got = run("outside", ["compress", "-i", "/corpus", "-o", "/work/x.znippy", "--format", "generic"])
ok = got == 159
failed += not ok
print(f"test a_call_outside_the_profile_is_killed_by_sigsys ... {'ok' if ok else 'FAILED'} (exit {got}, want 159)")
print(f"\nenforce: {len(w['run']) + 1 - failed} passed; {failed} failed")
sys.exit(1 if failed else 0)