nordisk/heimdall

crates/heimdall-cli/src/main.rs

// SPDX-License-Identifier: AGPL-3.0-only
//! heimdall-cli (spec §7): the `heimdall` command line.
//!
//! **Heimdall is a defensive sandbox.** Milestone 1 runs wasm jobs only:
//!
//! - `heimdall run --job wasm --module <guest> --trust-root <der> --in <dir> --out <dir> [--policy <file>]`
//! - `heimdall verify-result <file> --key <pinned key>`
//!
//! Native jobs and the `profile` verbs print "not in milestone 1" and exit 2.
//! Exit codes are documented in `docs/CLI.md`.

use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::path::{Path, PathBuf};
use std::process::ExitCode;

use clap::{Parser, Subcommand, ValueEnum};
use heimdall_core::{BootKey, JobId, JobTicket, Nonce, Params, Verdict};
use heimdall_run::verify::{default_files_dir, load_pinned_key, read_result_file, verify_result};
use heimdall_run::{run_job, GuestRunner, JobSpec, PolicyFile};

/// Exit code: the job ran and was Accepted, or the result verified and says Accepted.
pub const EXIT_OK: u8 = 0;
/// Exit code: the result did not verify (signature, key, ticket or outputs).
pub const EXIT_REJECTED: u8 = 1;
/// Exit code: a usage error, or a verb that is not in milestone 1.
pub const EXIT_USAGE: u8 = 2;
/// Exit code: the job could not start, or a file could not be read.
pub const EXIT_HOST: u8 = 3;
/// Exit code: the job ran and was Refused (the signed result was written), or
/// the result verified and says Refused.
pub const EXIT_REFUSED: u8 = 4;

#[derive(Parser, Debug)]
#[command(
    name = "heimdall",
    version,
    about = "Heimdall is a defensive sandbox: run untrusted work, trust only its signed verdict."
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand, Debug)]
enum Command {
    /// Run one job on an input dir; outputs and the signed result go into a fresh out dir.
    Run(RunArgs),
    /// Verify a signed result file against a pinned public key.
    VerifyResult(VerifyArgs),
    /// Syscall profiles: probe, gen, check (not in milestone 1).
    #[command(disable_help_flag = true)]
    Profile {
        #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
        rest: Vec<String>,
    },
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)]
enum JobKindArg {
    Wasm,
    Native,
}

#[derive(clap::Args, Debug)]
struct RunArgs {
    /// The job kind: wasm (milestone 1) or native (later).
    #[arg(long = "job", value_enum)]
    job: JobKindArg,
    /// The signed wasm guest (`<guest>.wasm` beside `.wasm.p7s` and `.wasm.blake3`).
    #[arg(long = "module")]
    module: Option<PathBuf>,
    /// A DER root certificate the guest's signature must chain to; may repeat.
    #[arg(long = "trust-root")]
    trust_roots: Vec<PathBuf>,
    /// Also demand that the guest's BLAKE3-256 equals this hex digest.
    #[arg(long = "pin-blake3")]
    pin_blake3: Option<String>,
    /// The input dir (exactly one regular file in milestone 1).
    #[arg(long = "in")]
    input: PathBuf,
    /// The out dir; it must not exist yet.
    #[arg(long = "out")]
    out: PathBuf,
    /// The job policy file (TOML, docs/POLICY.md).
    #[arg(long = "policy")]
    policy: Option<PathBuf>,
    /// Write this run's pinned public key record here (a new file).
    #[arg(long = "pin-out")]
    pin_out: Option<PathBuf>,
    /// The caller's job id (32 hex digits); random when absent.
    #[arg(long = "job-id", value_parser = parse_job_id)]
    job_id: Option<JobId>,
    /// The caller's nonce (64 hex digits); random when absent.
    #[arg(long = "nonce", value_parser = parse_nonce)]
    nonce: Option<Nonce>,
    /// A job parameter, `key=value`; may repeat (a key at most once).
    #[arg(long = "param", value_parser = parse_param)]
    params: Vec<(String, String)>,
}

#[derive(clap::Args, Debug)]
struct VerifyArgs {
    /// The signed result file (`out/result.signed`).
    file: PathBuf,
    /// The pinned public key: a file holding the record, or the record itself.
    #[arg(long = "key")]
    key: String,
    /// The dir whose files must match the result's outputs (default: `files/`
    /// beside the result file, when it exists).
    #[arg(long = "files")]
    files: Option<PathBuf>,
    /// The job id the result must be for (needs --nonce too).
    #[arg(long = "job-id", value_parser = parse_job_id, requires = "nonce")]
    job_id: Option<JobId>,
    /// The nonce the result must carry (needs --job-id too).
    #[arg(long = "nonce", value_parser = parse_nonce, requires = "job_id")]
    nonce: Option<Nonce>,
}

fn parse_param(s: &str) -> Result<(String, String), String> {
    let (k, v) = s
        .split_once('=')
        .ok_or_else(|| format!("param {s:?} is not key=value"))?;
    if k.is_empty() {
        return Err(format!("param {s:?} has an empty key"));
    }
    Ok((k.to_string(), v.to_string()))
}

fn parse_job_id(s: &str) -> Result<JobId, String> {
    JobId::from_hex(s).ok_or_else(|| format!("job id {s:?} is not {} hex digits", 2 * JobId::LEN))
}

fn parse_nonce(s: &str) -> Result<Nonce, String> {
    Nonce::from_hex(s).ok_or_else(|| format!("nonce {s:?} is not {} hex digits", 2 * Nonce::LEN))
}

const NOT_IN_M1: &str = "not in milestone 1";

fn main() -> ExitCode {
    let cli = match Cli::try_parse() {
        Ok(c) => c,
        Err(e) => {
            let _ = e.print();
            return ExitCode::from(match e.kind() {
                clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion => {
                    EXIT_OK
                }
                _ => EXIT_USAGE,
            });
        }
    };
    ExitCode::from(dispatch(cli))
}

fn dispatch(cli: Cli) -> u8 {
    match cli.command {
        Command::Profile { .. } => {
            eprintln!("heimdall profile: {NOT_IN_M1}");
            EXIT_USAGE
        }
        Command::Run(args) => run(args),
        Command::VerifyResult(args) => verify(args),
    }
}

/// Load the wasm guest. heimdall-wasm is wired in `heimdall_run::wasm`.
fn load_guest(args: &RunArgs, policy: &PolicyFile) -> Result<Box<dyn GuestRunner>, String> {
    let module = args.module.as_deref().ok_or("--module is required")?;
    heimdall_run::wasm::load_guest(
        module,
        &args.trust_roots,
        args.pin_blake3.as_deref(),
        &policy.limits,
    )
    .map(|g| Box::new(g) as Box<dyn GuestRunner>)
    .map_err(|e| e.to_string())
}

fn write_new_file(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .mode(0o644)
        .custom_flags(libc::O_NOFOLLOW)
        .open(path)?;
    f.write_all(bytes)?;
    f.sync_all()
}

fn run(args: RunArgs) -> u8 {
    if args.job == JobKindArg::Native {
        eprintln!("heimdall run --job native: {NOT_IN_M1}");
        return EXIT_USAGE;
    }
    if args.module.is_none() {
        eprintln!("heimdall run --job wasm needs --module <guest.wasm>");
        return EXIT_USAGE;
    }
    let mut params = Params::new();
    for (k, v) in &args.params {
        if params.insert(k.clone(), v.clone()).is_some() {
            eprintln!("heimdall run: --param {k} given twice");
            return EXIT_USAGE;
        }
    }
    let policy = match &args.policy {
        None => PolicyFile::default(),
        Some(p) => match PolicyFile::load(p) {
            Ok(p) => p,
            Err(e) => {
                eprintln!("heimdall run: policy {}: {e}", p.display());
                return EXIT_HOST;
            }
        },
    };
    let guest = match load_guest(&args, &policy) {
        Ok(g) => g,
        Err(e) => {
            eprintln!("heimdall run: {e}");
            return EXIT_HOST;
        }
    };
    // The per-boot key: one heimdall process is one boot. The secret half is
    // never written anywhere; the caller pins the public record.
    let key = BootKey::generate();
    let pin = key.pinned().to_text();
    if let Some(p) = &args.pin_out {
        if let Err(e) = write_new_file(p, pin.as_bytes()) {
            eprintln!("heimdall run: --pin-out {}: {e}", p.display());
            return EXIT_HOST;
        }
    }
    let mut spec = JobSpec::new(&args.input, &args.out, policy, params);
    if let Some(id) = args.job_id {
        spec.job_id = id;
    }
    if let Some(n) = args.nonce {
        spec.nonce = n;
    }
    let outcome = match run_job(&spec, guest.as_ref(), &key) {
        Ok(o) => o,
        Err(e) => {
            eprintln!("heimdall run: {e}");
            return EXIT_HOST;
        }
    };
    let r = &outcome.signed.result;
    print!("pinned-key: {pin}");
    println!("job-id: {}", r.job_id);
    println!("nonce: {}", r.nonce);
    println!("result: {}", outcome.result_path.display());
    match &r.verdict {
        Verdict::Accepted => {
            println!("verdict: accepted ({} outputs)", r.outputs.len());
            EXIT_OK
        }
        Verdict::Refused { code, reason } => {
            println!(
                "verdict: refused: {} [{}]",
                reason.escape_debug(),
                code.escape_debug()
            );
            EXIT_REFUSED
        }
    }
}

fn verify(args: VerifyArgs) -> u8 {
    let pinned = match load_pinned_key(&args.key) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("heimdall verify-result: {e}");
            return EXIT_HOST;
        }
    };
    let wire = match read_result_file(&args.file) {
        Ok(w) => w,
        Err(e) => {
            eprintln!("heimdall verify-result: {e}");
            return EXIT_HOST;
        }
    };
    let ticket = match (args.job_id, args.nonce) {
        (Some(job_id), Some(nonce)) => Some(JobTicket { job_id, nonce }),
        _ => None,
    };
    let files = args.files.clone().or_else(|| default_files_dir(&args.file));
    match verify_result(&wire, pinned, ticket.as_ref(), files.as_deref()) {
        Ok(v) => {
            let r = &v.signed.result;
            println!("verified: job {} key {}", r.job_id, v.signed.key_id);
            if let Some(d) = &v.files_dir {
                println!("outputs: {} checked in {}", v.files_checked, d.display());
            } else {
                println!("outputs: not checked (no files dir)");
            }
            match &r.verdict {
                Verdict::Accepted => {
                    println!("verdict: accepted");
                    EXIT_OK
                }
                Verdict::Refused { code, reason } => {
                    println!(
                        "verdict: refused: {} [{}]",
                        reason.escape_debug(),
                        code.escape_debug()
                    );
                    EXIT_REFUSED
                }
            }
        }
        Err(e) => {
            eprintln!("heimdall verify-result: {e}");
            if e.is_rejection() {
                EXIT_REJECTED
            } else {
                EXIT_HOST
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parse(args: &[&str]) -> Result<Cli, clap::Error> {
        Cli::try_parse_from(std::iter::once("heimdall").chain(args.iter().copied()))
    }

    #[test]
    fn run_wasm_with_all_flags_parses_into_run_args() {
        let cli = parse(&[
            "run", "--job", "wasm", "--module", "g.wasm", "--in", "i", "--out", "o", "--policy",
            "p.toml", "--param", "a=b", "--param", "c=",
        ])
        .unwrap();
        let Command::Run(r) = cli.command else {
            panic!("not run")
        };
        assert_eq!(r.job, JobKindArg::Wasm);
        assert_eq!(r.module, Some(PathBuf::from("g.wasm")));
        assert_eq!(r.input, PathBuf::from("i"));
        assert_eq!(r.out, PathBuf::from("o"));
        assert_eq!(r.policy, Some(PathBuf::from("p.toml")));
        assert_eq!(
            r.params,
            vec![("a".into(), "b".into()), ("c".into(), "".into())]
        );
    }

    #[test]
    fn run_without_in_is_a_usage_error() {
        assert!(parse(&["run", "--job", "wasm", "--module", "g", "--out", "o"]).is_err());
    }

    #[test]
    fn run_without_out_is_a_usage_error() {
        assert!(parse(&["run", "--job", "wasm", "--module", "g", "--in", "i"]).is_err());
    }

    #[test]
    fn run_with_unknown_job_kind_is_a_usage_error() {
        assert!(parse(&["run", "--job", "docker", "--in", "i", "--out", "o"]).is_err());
    }

    #[test]
    fn param_without_equals_sign_is_a_usage_error() {
        assert!(parse(&[
            "run", "--job", "wasm", "--module", "g", "--in", "i", "--out", "o", "--param", "noeq"
        ])
        .is_err());
    }

    #[test]
    fn param_with_empty_key_is_a_usage_error() {
        assert!(parse_param("=v").is_err());
    }

    #[test]
    fn param_value_may_contain_equals_signs() {
        assert_eq!(parse_param("k=a=b").unwrap(), ("k".into(), "a=b".into()));
    }

    #[test]
    fn verify_result_needs_a_pinned_key() {
        assert!(parse(&["verify-result", "r.signed"]).is_err());
        assert!(parse(&["verify-result", "r.signed", "--key", "k.pub"]).is_ok());
    }

    #[test]
    fn profile_accepts_any_trailing_words_so_it_can_say_not_in_milestone_1() {
        assert!(parse(&["profile", "probe", "holger-scan", "--anything"]).is_ok());
        assert!(parse(&["profile"]).is_ok());
    }

    #[test]
    fn unknown_subcommand_is_a_usage_error() {
        assert!(parse(&["exec", "x"]).is_err());
    }

    #[test]
    fn native_job_dispatch_exits_with_usage_code() {
        let cli = parse(&["run", "--job", "native", "--in", "i", "--out", "o"]).unwrap();
        assert_eq!(dispatch(cli), EXIT_USAGE);
    }

    #[test]
    fn profile_dispatch_exits_with_usage_code() {
        let cli = parse(&["profile", "gen", "w"]).unwrap();
        assert_eq!(dispatch(cli), EXIT_USAGE);
    }

    #[test]
    fn wasm_job_without_module_exits_with_usage_code() {
        let cli = parse(&["run", "--job", "wasm", "--in", "i", "--out", "o"]).unwrap();
        assert_eq!(dispatch(cli), EXIT_USAGE);
    }
}