nordisk/heimdall

crates/heimdall-wasm/tests/common/mod.rs

// SPDX-License-Identifier: AGPL-3.0-only
//! Test kit: tiny guests written as core WAT, turned into components with the
//! real `wit/heimdall.wit`, signed with a throwaway CA exactly as a release
//! door signs holger's plugins (`.wasm.p7s` CMS over blake3 + `.wasm.blake3`).
//!
//! Heimdall is a defensive sandbox; these guests attack only its own limits.
#![allow(dead_code)]

use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};

use guest_host::verify::sign::{dev, sign_module, Ed25519Signer};
use heimdall_wasm::{Guest, Host, HostConfig, LoadRefusal, RunLimits, WIT, WORLD};
use wit_component::{ComponentEncoder, StringEncoding};
use wit_parser::Resolve;

static NEXT: AtomicUsize = AtomicUsize::new(0);

/// A fresh scratch directory under cargo's per-target tmp dir.
pub fn scratch(name: &str) -> PathBuf {
    let n = NEXT.fetch_add(1, Ordering::SeqCst);
    let dir = Path::new(env!("CARGO_TARGET_TMPDIR"))
        .join(format!("hw-{name}-{}-{n}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    dir
}

/// A throwaway CA, a release signer under it, and a host trusting the CA.
pub struct Kit {
    pub dir: PathBuf,
    pub root: PathBuf,
    pub signer: Ed25519Signer,
}

impl Kit {
    pub fn new(name: &str) -> Kit {
        let dir = scratch(name);
        let (ca_key, ca_der) = dev::mint_ca("heimdall-test-ca").unwrap();
        let root = dir.join("root.der");
        std::fs::write(&root, &ca_der).unwrap();
        let signer = dev::new_ed25519_signer(&ca_key, &ca_der, "heimdall-test-release").unwrap();
        Kit { dir, root, signer }
    }

    pub fn config(&self) -> HostConfig {
        HostConfig::new([self.root.clone()])
    }

    pub fn host(&self) -> Host {
        Host::new(self.config()).expect("host")
    }

    /// Write `bytes` as `<name>.wasm` with its `.p7s` and `.blake3` beside it.
    pub fn place(&self, name: &str, bytes: &[u8]) -> PathBuf {
        let module = self.dir.join(format!("{name}.wasm"));
        std::fs::write(&module, bytes).unwrap();
        let cms = sign_module(bytes, &self.signer).unwrap();
        std::fs::write(sidecar(&module, "p7s"), cms).unwrap();
        std::fs::write(
            sidecar(&module, "blake3"),
            format!("{}\n", blake3::hash(bytes).to_hex()),
        )
        .unwrap();
        module
    }

    /// Place a WAT guest and load it; panics on a load refusal.
    pub fn guest(&self, name: &str, core_wat: &str) -> Guest {
        let path = self.place(name, &component(core_wat));
        self.host()
            .load(&path)
            .unwrap_or_else(|e| panic!("{name}: {e}"))
    }

    /// Place raw bytes (signed) and try to load them.
    pub fn try_load(&self, name: &str, bytes: &[u8]) -> Result<Guest, LoadRefusal> {
        let path = self.place(name, bytes);
        self.host().load(&path)
    }
}

pub fn sidecar(module: &Path, ext: &str) -> PathBuf {
    let mut n = module.file_name().unwrap().to_os_string();
    n.push(".");
    n.push(ext);
    module.with_file_name(n)
}

/// Componentize a core module against heimdall's own WIT world.
pub fn component(core_wat: &str) -> Vec<u8> {
    component_with(core_wat, &[("heimdall.wit", WIT)], WORLD)
}

/// Componentize a core module against any WIT packages and world.
pub fn component_with(core_wat: &str, wit: &[(&str, &str)], world: &str) -> Vec<u8> {
    let mut resolve = Resolve::default();
    let mut main = None;
    for (path, src) in wit {
        main = Some(resolve.push_str(path, src).unwrap());
    }
    let world = resolve.select_world(&[main.unwrap()], Some(world)).unwrap();
    let mut core =
        wat::parse_str(core_wat).unwrap_or_else(|e| panic!("bad test WAT: {e}\n{core_wat}"));
    wit_component::embed_component_metadata(&mut core, &resolve, world, StringEncoding::UTF8)
        .unwrap();
    ComponentEncoder::default()
        .module(&core)
        .unwrap()
        .validate(true)
        .encode()
        .unwrap()
}

/// The data area guests put constant strings in starts here.
pub const DATA: u32 = 16;

/// A core module implementing `run` with the canonical ABI by hand.
///
/// Layout: constant data from 16, the result record at 1024 (tag u8 @1024,
/// string ptr @1028, len @1032), a bump heap from 4096 served by
/// `cabi_realloc` (which grows memory as needed). `$ok`/`$refuse` take a
/// (ptr, len) string and return the result pointer. `$emit` is the host's
/// emit-file (name ptr, name len, bytes ptr, bytes len). `$hexname` writes
/// eight hex digits of its argument at 64 (a unique file name per counter).
pub struct Wat {
    pub pages: u32,
    pub extra: String,
    pub body: String,
    pub realloc_prelude: String,
}

impl Wat {
    pub fn new(body: &str) -> Wat {
        Wat {
            pages: 1,
            extra: String::new(),
            body: body.to_string(),
            realloc_prelude: String::new(),
        }
    }
    pub fn pages(mut self, pages: u32) -> Wat {
        self.pages = pages;
        self
    }
    pub fn extra(mut self, extra: &str) -> Wat {
        self.extra.push_str(extra);
        self
    }
    /// Code run at the start of every `cabi_realloc` call.
    pub fn realloc_prelude(mut self, code: &str) -> Wat {
        self.realloc_prelude = code.to_string();
        self
    }
    pub fn text(&self) -> String {
        format!(
            r#"(module
  (import "$root" "emit-file" (func $emit (param i32 i32 i32 i32)))
  (memory (export "memory") {pages})
  (global $heap (mut i32) (i32.const 4096))
  (func $realloc (export "cabi_realloc") (param $old i32) (param $old_size i32) (param $align i32) (param $size i32) (result i32)
    (local $p i32)
    {prelude}
    (local.set $p (i32.and (i32.add (global.get $heap) (i32.sub (local.get $align) (i32.const 1)))
                           (i32.xor (i32.sub (local.get $align) (i32.const 1)) (i32.const -1))))
    (global.set $heap (i32.add (local.get $p) (local.get $size)))
    (block $done
      (loop $grow
        (br_if $done (i64.le_u (i64.extend_i32_u (global.get $heap)) (i64.mul (i64.extend_i32_u (memory.size)) (i64.const 65536))))
        (if (i32.eq (memory.grow (i32.const 1)) (i32.const -1)) (then unreachable))
        (br $grow)))
    (local.get $p))
  (func $ok (param $p i32) (param $l i32) (result i32)
    (i32.store8 (i32.const 1024) (i32.const 0))
    (i32.store (i32.const 1028) (local.get $p))
    (i32.store (i32.const 1032) (local.get $l))
    (i32.const 1024))
  (func $refuse (param $p i32) (param $l i32) (result i32)
    (i32.store8 (i32.const 1024) (i32.const 1))
    (i32.store (i32.const 1028) (local.get $p))
    (i32.store (i32.const 1032) (local.get $l))
    (i32.const 1024))
  (func $hexname (param $n i32)
    (local $i i32) (local $d i32)
    (local.set $i (i32.const 7))
    (loop $l
      (local.set $d (i32.and (local.get $n) (i32.const 15)))
      (i32.store8 (i32.add (i32.const 64) (local.get $i))
        (i32.add (local.get $d) (select (i32.const 87) (i32.const 48) (i32.ge_u (local.get $d) (i32.const 10)))))
      (local.set $n (i32.shr_u (local.get $n) (i32.const 4)))
      (local.set $i (i32.sub (local.get $i) (i32.const 1)))
      (br_if $l (i32.ge_s (local.get $i) (i32.const 0)))))
  {extra}
  (func (export "run") (param $in i32) (param $in_len i32) (param $params i32) (param $params_len i32) (result i32)
    {body})
)"#,
            pages = self.pages,
            prelude = self.realloc_prelude,
            extra = self.extra,
            body = self.body
        )
    }
}

/// The well-behaved guest. Input starting with 'T' is refused by the guest
/// ("input signature mismatch"), 'M' traps (malformed input), anything else
/// is emitted as `report.txt` and accepted with the note "accepted".
pub fn good_wat() -> String {
    Wat::new(
        r#"(if (i32.gt_u (local.get $in_len) (i32.const 0))
      (then
        (if (i32.eq (i32.load8_u (local.get $in)) (i32.const 84))
          (then (return (call $refuse (i32.const 48) (i32.const 24)))))
        (if (i32.eq (i32.load8_u (local.get $in)) (i32.const 77))
          (then unreachable))))
    (call $emit (i32.const 16) (i32.const 10) (local.get $in) (local.get $in_len))
    (call $ok (i32.const 32) (i32.const 8))"#,
    )
    .extra(r#"(data (i32.const 16) "report.txt") (data (i32.const 32) "accepted") (data (i32.const 48) "input signature mismatch")"#)
    .text()
}

/// Defaults, but small enough that a bomb is caught quickly.
pub fn tight() -> RunLimits {
    RunLimits {
        fuel: 50_000_000,
        timeout: std::time::Duration::from_secs(5),
        memory_bytes: 16 << 20,
        ..RunLimits::default()
    }
}

pub fn reason(o: &heimdall_wasm::GuestOutcome) -> String {
    o.verdict
        .reason()
        .unwrap_or_else(|| panic!("expected a refusal, got {}", o.verdict))
        .to_string()
}

pub fn assert_trapped(o: &heimdall_wasm::GuestOutcome, needle: &str) {
    let r = reason(o);
    assert!(r.starts_with("guest trapped ("), "not a trap refusal: {r}");
    assert!(
        r.contains(needle),
        "the refusal does not name {needle:?}: {r}"
    );
    assert!(
        o.files.is_empty(),
        "a trapped run handed on {} file(s)",
        o.files.len()
    );
}

/// No job parameters.
pub fn none() -> heimdall_core::Params {
    heimdall_core::Params::new()
}

/// Job parameters from pairs.
pub fn params(pairs: &[(&str, &str)]) -> heimdall_core::Params {
    pairs
        .iter()
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .collect()
}