nordisk/heimdall

crates/heimdall-run/src/runner.rs

// SPDX-License-Identifier: AGPL-3.0-only
//! [`WasmJobRunner`]: the milestone 1 wasm path behind heimdall-core's
//! [`JobRunner`] trait.
//!
//! **Heimdall is a defensive sandbox.** For every job the runner verifies and
//! loads the signed guest (signature chain to a trust root plus the job's
//! blake3 pin), runs it with no WASI under the job's limits through
//! [`run_job`] into a private staging dir, then moves the emitted files into
//! the job's `output_dir` and returns the signed result. The per-boot key
//! stays in the runner; the guest never sees it.

use std::path::{Path, PathBuf};

use heimdall_core::{
    BootKey, Features, Job, JobKind, JobRunner, Limits, Params, RunError, SignedResult,
    VerifyingKey,
};

use crate::job::{run_job, wasm_features, JobError, JobSpec};
use crate::outdir::FILES_DIR;
use crate::policy_file::{PolicyFile, RunLimits, DEFAULT_MAX_INPUT_BYTES, LIMIT_KEYS};
use crate::wasm::load_guest;

/// Runs wasm jobs; refuses native jobs with [`RunError::Unsupported`].
pub struct WasmJobRunner {
    key: BootKey,
    trust_roots: Vec<PathBuf>,
    work_dir: PathBuf,
}

impl WasmJobRunner {
    /// A runner with a fresh per-boot key. `work_dir` holds per-job staging
    /// dirs (best on the same file system as the jobs' output dirs).
    pub fn new(trust_roots: Vec<PathBuf>, work_dir: PathBuf) -> Self {
        Self::with_key(BootKey::generate(), trust_roots, work_dir)
    }

    /// A runner with a given key (tests and fixtures).
    pub fn with_key(key: BootKey, trust_roots: Vec<PathBuf>, work_dir: PathBuf) -> Self {
        Self {
            key,
            trust_roots,
            work_dir,
        }
    }

    /// The pinned-key record of this runner's per-boot key.
    pub fn pinned(&self) -> heimdall_core::PinnedKey {
        self.key.pinned()
    }
}

/// heimdall-core limits as the wasm run's limits: every shared limit must be
/// between 1 and Heimdall's ceiling, and fuel is always metered.
pub fn run_limits_from_core(l: &Limits) -> Result<RunLimits, RunError> {
    let fuel = l.fuel.ok_or_else(|| RunError::InvalidJob {
        reason: "limits.fuel is None: a Heimdall wasm run always meters fuel".into(),
    })?;
    let r = RunLimits {
        fuel,
        wall_time_ms: l.wall_ms(),
        memory_bytes: l.memory_bytes,
        table_elements: l.table_elements,
        instances: l.instances,
        tables: l.tables,
        memories: l.memories,
        stack_bytes: l.stack_bytes,
        max_input_bytes: DEFAULT_MAX_INPUT_BYTES,
        max_output_files: l.max_output_files,
        max_output_file_bytes: l.max_output_file_bytes,
        max_output_total_bytes: l.max_output_total_bytes,
    };
    let c = RunLimits::ceilings();
    let pairs = [
        (r.fuel, c.fuel),
        (r.wall_time_ms, c.wall_time_ms),
        (r.memory_bytes, c.memory_bytes),
        (r.table_elements.into(), c.table_elements.into()),
        (r.instances.into(), c.instances.into()),
        (r.tables.into(), c.tables.into()),
        (r.memories.into(), c.memories.into()),
        (r.stack_bytes, c.stack_bytes),
        (r.max_input_bytes, c.max_input_bytes),
        (r.max_output_files.into(), c.max_output_files.into()),
        (r.max_output_file_bytes, c.max_output_file_bytes),
        (r.max_output_total_bytes, c.max_output_total_bytes),
    ];
    for ((v, ceiling), (key, _)) in pairs.into_iter().zip(LIMIT_KEYS) {
        if v == 0 || v > ceiling {
            return Err(RunError::InvalidJob {
                reason: format!("limit {key} = {v} is outside 1..={ceiling}"),
            });
        }
    }
    if r.max_output_file_bytes > r.max_output_total_bytes {
        return Err(RunError::InvalidJob {
            reason: "max_output_file_bytes is above max_output_total_bytes".into(),
        });
    }
    Ok(r)
}

/// `Job::params` as string pairs: `null` or a JSON object of strings.
pub fn params_from_json(v: &serde_json::Value) -> Result<Params, RunError> {
    let bad = || RunError::InvalidJob {
        reason: "params must be null or a JSON object of string values".into(),
    };
    match v {
        serde_json::Value::Null => Ok(Params::new()),
        serde_json::Value::Object(m) => m
            .iter()
            .map(|(k, v)| Ok((k.clone(), v.as_str().ok_or_else(bad)?.to_string())))
            .collect(),
        _ => Err(bad()),
    }
}

fn io(e: impl std::fmt::Display) -> RunError {
    RunError::Io {
        reason: e.to_string(),
    }
}

fn check_empty_dir(dir: &Path) -> Result<(), RunError> {
    let meta = std::fs::symlink_metadata(dir).map_err(|e| RunError::InvalidJob {
        reason: format!("output_dir {}: {e}", dir.display()),
    })?;
    if !meta.is_dir() {
        return Err(RunError::InvalidJob {
            reason: format!("output_dir {} is not a directory", dir.display()),
        });
    }
    if std::fs::read_dir(dir).map_err(io)?.next().is_some() {
        return Err(RunError::InvalidJob {
            reason: format!("output_dir {} is not empty", dir.display()),
        });
    }
    Ok(())
}

/// Move every file of `from` into `to` (rename, or copy when on another file system).
fn move_files(from: &Path, to: &Path) -> Result<(), RunError> {
    if !from.exists() {
        return Ok(());
    }
    for e in std::fs::read_dir(from).map_err(io)? {
        let e = e.map_err(io)?;
        let dst = to.join(e.file_name());
        if std::fs::rename(e.path(), &dst).is_err() {
            std::fs::copy(e.path(), &dst).map_err(io)?;
        }
    }
    Ok(())
}

impl JobRunner for WasmJobRunner {
    fn run(&self, job: Job) -> Result<SignedResult, RunError> {
        let JobKind::Wasm { module } = &job.kind else {
            return Err(RunError::Unsupported {
                what: "native jobs: the heimdall-wasm runner runs only wasm guests".into(),
            });
        };
        if job.policy.net.is_some() {
            return Err(RunError::Unsupported {
                what: "network access: a wasm guest has no network".into(),
            });
        }
        if !job.policy.read.is_empty() || !job.policy.write.is_empty() {
            return Err(RunError::Unsupported {
                what: "extra read/write paths: a wasm guest reads only its input".into(),
            });
        }
        let limits = run_limits_from_core(&job.policy.limits)?;
        let params = params_from_json(&job.params)?;
        check_empty_dir(&job.output_dir)?;
        std::fs::create_dir_all(&self.work_dir).map_err(io)?;
        let guest = load_guest(
            &module.path,
            &self.trust_roots,
            Some(&module.blake3),
            &limits,
        )
        .map_err(|e| RunError::Image {
            reason: e.to_string(),
        })?;
        let staging = self.work_dir.join(format!("job-{}", job.id));
        let spec = JobSpec {
            input_dir: job.input_dir.clone(),
            out_dir: staging.clone(),
            policy: PolicyFile {
                job: Some("wasm".into()),
                limits,
            },
            params,
            job_id: job.id,
            nonce: job.nonce,
        };
        let outcome = run_job(&spec, &guest, &self.key);
        let moved = match &outcome {
            Ok(_) => move_files(&staging.join(FILES_DIR), &job.output_dir),
            Err(_) => Ok(()),
        };
        let _ = std::fs::remove_dir_all(&staging);
        moved?;
        match outcome {
            Ok(o) => Ok(o.signed),
            Err(JobError::Input(e)) => Err(RunError::InvalidJob {
                reason: format!("input: {e}"),
            }),
            Err(JobError::OutDir(e)) => Err(io(e)),
            Err(JobError::Sign(e)) => Err(RunError::Sign(e)),
        }
    }

    fn verifying_key(&self) -> VerifyingKey {
        self.key.verifying_key()
    }

    fn features(&self) -> Features {
        wasm_features()
    }
}