nordisk/heimdall

crates/heimdall-plain/src/shim.rs

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Ignalina ApS and Vetra Automation AB

//! The shim: one short-lived process per job that drives youki's
//! `libcontainer` (create, start, wait, kill on the wall clock, delete).
//!
//! Heimdall is a defensive sandbox. The runner never forks inside its own
//! (multi-threaded) process: it starts this shim with a JSON request on stdin
//! and reads a JSON response from stdout. The shim never receives the
//! per-boot signing key, and the job's stdio is `/dev/null` (or the shim's
//! stderr with `HEIMDALL_PLAIN_DEBUG=1`), so a job cannot write into the
//! runner's streams.

use std::path::PathBuf;

use serde::{Deserialize, Serialize};

/// The argument that makes a host binary act as the shim
/// (see [`run_if_shim`]).
pub const SHIM_ARG: &str = "--heimdall-plain-shim";

/// What the runner asks the shim to do.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ShimRequest {
    /// Container id (unique per job).
    pub id: String,
    /// The bundle: `config.json` + `rootfs/`.
    pub bundle: PathBuf,
    /// libcontainer's state root.
    pub state_root: PathBuf,
    /// The wall-clock limit in milliseconds.
    pub wall_ms: u64,
    /// The job's cgroup v2 dir (to read `memory.events` / `pids.events`).
    pub cgroup_dir: PathBuf,
}

/// What happened.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ShimResponse {
    /// The container never ran: why.
    pub error: Option<String>,
    pub exit_code: Option<i32>,
    pub signal: Option<i32>,
    /// The shim killed the job at the wall-clock limit.
    pub wall_killed: bool,
    /// `oom_kill` count from the job's `memory.events`.
    pub oom_kills: u64,
    /// `max` count from the job's `pids.events`.
    pub pids_max: u64,
    /// Milliseconds the job ran.
    pub ran_ms: u64,
}

/// Read `key N` from a cgroup `*.events` file.
pub fn events_count(text: &str, key: &str) -> u64 {
    text.lines()
        .filter_map(|l| l.split_once(' '))
        .find(|(k, _)| *k == key)
        .and_then(|(_, v)| v.trim().parse().ok())
        .unwrap_or(0)
}

/// Call first thing in `main()` of a binary that should also serve as the
/// shim (`PlainConfig::shim_args = vec![SHIM_ARG.into()]`): when the first
/// argument is [`SHIM_ARG`], it runs the shim and exits.
#[cfg(feature = "youki")]
pub fn run_if_shim() {
    if std::env::args().nth(1).as_deref() == Some(SHIM_ARG) {
        std::process::exit(main());
    }
}

/// The shim's main: request on stdin, response on stdout. Returns the exit code.
#[cfg(feature = "youki")]
pub fn main() -> i32 {
    use std::io::Read;
    let mut input = Vec::new();
    if let Err(e) = std::io::stdin().take(1 << 20).read_to_end(&mut input) {
        return respond(&ShimResponse {
            error: Some(format!("reading the request: {e}")),
            ..Default::default()
        });
    }
    let resp = match serde_json::from_slice::<ShimRequest>(&input) {
        Ok(req) => youki::run(&req),
        Err(e) => ShimResponse {
            error: Some(format!("bad request: {e}")),
            ..Default::default()
        },
    };
    respond(&resp)
}

#[cfg(feature = "youki")]
fn respond(r: &ShimResponse) -> i32 {
    println!("{}", serde_json::to_string(r).expect("response serializes"));
    0
}

#[cfg(feature = "youki")]
mod youki {
    use std::fs::File;
    use std::time::{Duration, Instant};

    use libcontainer::container::builder::ContainerBuilder;
    use libcontainer::syscall::syscall::SyscallType;
    use libcontainer::workload::default::DefaultExecutor;
    use nix::sys::signal::Signal;
    use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus};

    use super::{events_count, ShimRequest, ShimResponse};

    fn stdio(debug: bool) -> std::io::Result<(File, File, File)> {
        use std::os::fd::AsFd;
        let null = || File::options().read(true).write(true).open("/dev/null");
        let err = if debug {
            File::from(std::io::stderr().as_fd().try_clone_to_owned()?)
        } else {
            null()?
        };
        Ok((null()?, null()?, err))
    }

    fn err(e: impl std::fmt::Display) -> ShimResponse {
        ShimResponse {
            error: Some(e.to_string()),
            ..Default::default()
        }
    }

    pub(super) fn run(req: &ShimRequest) -> ShimResponse {
        let debug = std::env::var_os("HEIMDALL_PLAIN_DEBUG").is_some();
        let (i, o, e) = match stdio(debug) {
            Ok(x) => x,
            Err(e) => return err(format!("stdio: {e}")),
        };
        if let Err(e) = std::fs::create_dir_all(&req.state_root) {
            return err(format!("state root: {e}"));
        }
        let builder = match ContainerBuilder::new(req.id.clone(), SyscallType::default())
            .with_executor(DefaultExecutor {})
            .with_root_path(req.state_root.clone())
        {
            Ok(b) => b,
            Err(e) => return err(format!("container root: {e}")),
        };
        let mut container = match builder
            .with_stdin(i)
            .with_stdout(o)
            .with_stderr(e)
            .as_init(&req.bundle)
            .with_systemd(false)
            .with_detach(true)
            .build()
        {
            Ok(c) => c,
            Err(e) => return err(format!("create: {e}")),
        };
        let Some(pid) = container.pid() else {
            let _ = container.delete(true);
            return err("create: no init pid");
        };
        // No swap beyond the memory limit. libcgroups 0.7.0 skips writing a
        // computed `memory.swap.max` of 0 (OCI swap == limit), which leaves
        // swap unlimited; write it here, before the job starts.
        let swap = req.cgroup_dir.join("memory.swap.max");
        if swap.exists() {
            if let Err(e) = std::fs::write(&swap, "0") {
                let _ = container.kill(Signal::SIGKILL, true);
                let _ = waitpid(pid, None);
                let _ = container.delete(true);
                return err(format!("memory.swap.max: {e}"));
            }
        }
        let t0 = Instant::now();
        if let Err(e) = container.start() {
            let _ = container.kill(Signal::SIGKILL, true);
            let _ = waitpid(pid, None);
            let _ = container.delete(true);
            return err(format!("start: {e}"));
        }
        let deadline = t0 + Duration::from_millis(req.wall_ms);
        let mut resp = ShimResponse::default();
        let status = loop {
            match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
                Ok(WaitStatus::StillAlive) => {
                    if Instant::now() >= deadline {
                        resp.wall_killed = true;
                        let _ = container.kill(Signal::SIGKILL, true);
                        break waitpid(pid, None);
                    }
                    std::thread::sleep(Duration::from_millis(5));
                }
                other => break other,
            }
        };
        resp.ran_ms = t0.elapsed().as_millis() as u64;
        match status {
            Ok(WaitStatus::Exited(_, c)) => resp.exit_code = Some(c),
            Ok(WaitStatus::Signaled(_, s, _)) => resp.signal = Some(s as i32),
            Ok(other) => resp.error = Some(format!("unexpected wait status {other:?}")),
            Err(e) => resp.error = Some(format!("wait: {e}")),
        }
        let read = |f: &str| std::fs::read_to_string(req.cgroup_dir.join(f)).unwrap_or_default();
        resp.oom_kills = events_count(&read("memory.events"), "oom_kill");
        resp.pids_max = events_count(&read("pids.events"), "max");
        if let Err(e) = container.delete(true) {
            if resp.error.is_none() && debug {
                eprintln!("heimdall-plain-shim: delete: {e}");
            }
        }
        resp
    }
}