nordisk/heimdall

crates/heimdall-cli/tests/profile_corpus.rs

// SPDX-License-Identifier: AGPL-3.0-only
//! Writes the probe corpus of the `heimdall-wasm-host` profile
//! (profiles/heimdall-wasm-host/corpus): small WAT guests turned into
//! components, signed with a throwaway CA, their inputs and policies.
//! Heimdall is a defensive sandbox; the adversarial guests attack only
//! Heimdall's own limits.
//!
//! The real verify-bundle guest's part of the corpus (verify-bundle/ and its
//! runs in workload.toml) is guests/verify-bundle/tests/probe_corpus.rs: its
//! Holger bundles need that workspace's dev-dependencies.
//!
//! Run by hand when the corpus changes, then probe and regenerate:
//!   HEIMDALL_WRITE_CORPUS=profiles/heimdall-wasm-host/corpus \
//!     cargo test -p nordisk-heimdall-cli --test profile_corpus -- --ignored

#[path = "../../heimdall-wasm/tests/common/mod.rs"]
mod common;

use std::path::Path;

use common::{component, good_wat, sidecar, Kit, Wat};

const SPIN: &str = "(loop $l (br $l)) (i32.const 0)";

fn copy_guest(kit_module: &Path, to: &Path, name: &str) {
    for ext in ["", ".p7s", ".blake3"] {
        let from = if ext.is_empty() {
            kit_module.to_path_buf()
        } else {
            sidecar(kit_module, &ext[1..])
        };
        std::fs::copy(&from, to.join(format!("{name}.wasm{ext}"))).unwrap();
    }
}

fn input(to: &Path, dir: &str, bytes: &[u8]) {
    std::fs::create_dir_all(to.join(dir)).unwrap();
    std::fs::write(to.join(dir).join("input.bin"), bytes).unwrap();
}

#[test]
#[ignore = "writes the checked-in probe corpus; run by hand"]
fn writes_the_heimdall_wasm_host_probe_corpus() {
    let Some(out) = std::env::var_os("HEIMDALL_WRITE_CORPUS") else {
        panic!("set HEIMDALL_WRITE_CORPUS=<dir>");
    };
    let out = Path::new(&out);
    let _ = std::fs::remove_dir_all(out);
    std::fs::create_dir_all(out).unwrap();
    let kit = Kit::new("corpus");
    std::fs::copy(&kit.root, out.join("root.der")).unwrap();

    let guests: Vec<(&str, String)> = vec![
        ("good", good_wat()),
        ("spin", Wat::new(SPIN).text()),
        (
            "membomb",
            Wat::new("(loop $l (br_if $l (i32.ne (memory.grow (i32.const 16)) (i32.const -1)))) unreachable").text(),
        ),
        (
            "recurse",
            Wat::new("(drop (call $deep (i32.const 0))) (call $ok (i32.const 16) (i32.const 0))")
                .extra("(func $deep (param i32) (result i32) (call $deep (i32.add (local.get 0) (i32.const 1))))")
                .text(),
        ),
        (
            "manyfiles",
            Wat::new(
                r#"(local $i i32)
    (loop $l
      (call $hexname (local.get $i))
      (call $emit (i32.const 64) (i32.const 8) (i32.const 16) (i32.const 1))
      (local.set $i (i32.add (local.get $i) (i32.const 1)))
      (br_if $l (i32.lt_u (local.get $i) (i32.const 100000))))
    (call $ok (i32.const 16) (i32.const 0))"#,
            )
            .text(),
        ),
    ];
    for (name, wat) in &guests {
        let m = kit.place(name, &component(wat));
        copy_guest(&m, out, name);
    }
    // A guest whose signature is for another module: refused before load.
    let good = component(&good_wat());
    let other = kit.place("other", &component(&Wat::new(SPIN).text()));
    std::fs::write(out.join("badsig.wasm"), &good).unwrap();
    std::fs::copy(sidecar(&other, "p7s"), out.join("badsig.wasm.p7s")).unwrap();
    std::fs::copy(sidecar(&other, "blake3"), out.join("badsig.wasm.blake3")).unwrap();

    input(out, "in-good", b"a good input for the probe corpus\n");
    input(out, "in-tampered", b"T: tampered input\n");
    input(out, "in-malformed", b"M: malformed input\n");
    let policies = [
        (
            "policy-fuel.toml",
            "version = 1\n[limits]\nfuel = 1000000\nwall_time_ms = 60000\n",
        ),
        (
            "policy-time.toml",
            "version = 1\n[limits]\nfuel = 1000000000000\nwall_time_ms = 300\n",
        ),
        (
            "policy-mem.toml",
            "version = 1\n[limits]\nmemory_bytes = 8388608\n",
        ),
        (
            "policy-files.toml",
            "version = 1\n[limits]\nmax_output_files = 8\n",
        ),
    ];
    for (n, t) in policies {
        std::fs::write(
            out.join(n),
            format!("# Heimdall is a defensive sandbox: probe corpus policy.\n{t}"),
        )
        .unwrap();
    }
    std::fs::write(
        out.join("README.md"),
        "# heimdall-wasm-host probe corpus\n\nHeimdall is a defensive sandbox. Generated by\n`crates/heimdall-cli/tests/profile_corpus.rs` (guests signed with a throwaway\ntest CA, `root.der`); the runs are in `../workload.toml`.\n",
    )
    .unwrap();
}