nordisk/heimdall

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

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

//! Test kit for heimdall-plain: OCI image layouts built on the fly (a static
//! busybox, or synthetic layers), job dirs, fake shims. Heimdall is a
//! defensive sandbox; these tests exercise only its own containment.
#![allow(dead_code)]

use std::io::Write;
use std::path::{Path, PathBuf};
use std::time::Duration;

use heimdall_core::*;
use sha2::{Digest, Sha256};

pub fn scratch(name: &str) -> PathBuf {
    let base = std::env::var_os("HEIMDALL_PLAIN_SCRATCH")
        .map(PathBuf::from)
        .unwrap_or_else(std::env::temp_dir);
    let d = base.join(format!("heimdall-plain-{name}-{}", JobId::random()));
    std::fs::create_dir_all(&d).unwrap();
    d
}

/// One layer entry.
pub enum Entry {
    Dir(&'static str),
    File(String, Vec<u8>, u32),
    Symlink(String, String),
    HardLink(String, String),
    /// A raw path written into the header (for traversal tests).
    RawFile(Vec<u8>, Vec<u8>),
}

fn sha(b: &[u8]) -> String {
    format!("sha256:{}", heimdall_core::hex::encode(&Sha256::digest(b)))
}

/// An OCI image layout under construction.
pub struct Layout {
    pub dir: PathBuf,
}

impl Layout {
    pub fn new(dir: &Path) -> Layout {
        std::fs::create_dir_all(dir.join("blobs/sha256")).unwrap();
        std::fs::write(dir.join("oci-layout"), br#"{"imageLayoutVersion":"1.0.0"}"#).unwrap();
        std::fs::write(
            dir.join("index.json"),
            br#"{"schemaVersion":2,"manifests":[]}"#,
        )
        .unwrap();
        Layout {
            dir: dir.to_path_buf(),
        }
    }

    pub fn blob(&self, bytes: &[u8]) -> (String, u64) {
        let d = sha(bytes);
        std::fs::write(self.blob_path(&d), bytes).unwrap();
        (d, bytes.len() as u64)
    }

    pub fn blob_path(&self, digest: &str) -> PathBuf {
        self.dir
            .join("blobs/sha256")
            .join(digest.strip_prefix("sha256:").unwrap())
    }

    pub fn tar(entries: &[Entry]) -> Vec<u8> {
        let mut b = tar::Builder::new(Vec::new());
        for e in entries {
            let mut h = tar::Header::new_gnu();
            match e {
                Entry::Dir(p) => {
                    h.set_entry_type(tar::EntryType::Directory);
                    h.set_mode(0o755);
                    h.set_size(0);
                    b.append_data(&mut h, p, &[][..]).unwrap();
                }
                Entry::File(p, data, mode) => {
                    h.set_entry_type(tar::EntryType::Regular);
                    h.set_mode(*mode);
                    h.set_size(data.len() as u64);
                    b.append_data(&mut h, p, &data[..]).unwrap();
                }
                Entry::Symlink(p, target) => {
                    h.set_entry_type(tar::EntryType::Symlink);
                    h.set_size(0);
                    h.set_mode(0o777);
                    b.append_link(&mut h, p, target).unwrap();
                }
                Entry::HardLink(p, target) => {
                    h.set_entry_type(tar::EntryType::Link);
                    h.set_size(0);
                    h.set_mode(0o644);
                    b.append_link(&mut h, p, target).unwrap();
                }
                Entry::RawFile(p, data) => {
                    h.set_entry_type(tar::EntryType::Regular);
                    h.set_mode(0o644);
                    h.set_size(data.len() as u64);
                    let name = &mut h.as_old_mut().name;
                    name[..p.len()].copy_from_slice(p);
                    h.set_cksum();
                    b.append(&h, &data[..]).unwrap();
                }
            }
        }
        b.into_inner().unwrap()
    }

    pub fn gzip(raw: &[u8]) -> Vec<u8> {
        let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
        e.write_all(raw).unwrap();
        e.finish().unwrap()
    }

    /// Store layers (gzip or plain) + config + manifest; return the image ref.
    pub fn image(&self, layers: &[(Vec<u8>, bool)], env: &[&str]) -> ImageRef {
        let mut descs = Vec::new();
        for (raw, gz) in layers {
            let bytes = if *gz { Self::gzip(raw) } else { raw.clone() };
            let (d, n) = self.blob(&bytes);
            descs.push(serde_json::json!({
                "mediaType": if *gz { "application/vnd.oci.image.layer.v1.tar+gzip" } else { "application/vnd.oci.image.layer.v1.tar" },
                "digest": d, "size": n
            }));
        }
        let cfg = serde_json::json!({
            "architecture": "amd64", "os": "linux",
            "config": {"Env": env},
            "rootfs": {"type": "layers", "diff_ids": []}
        });
        let (cd, cn) = self.blob(&serde_json::to_vec(&cfg).unwrap());
        let m = serde_json::json!({
            "schemaVersion": 2,
            "mediaType": "application/vnd.oci.image.manifest.v1+json",
            "config": {"mediaType": "application/vnd.oci.image.config.v1+json", "digest": cd, "size": cn},
            "layers": descs
        });
        let (md, _) = self.blob(&serde_json::to_vec(&m).unwrap());
        ImageRef {
            name: "test/image".into(),
            digest: md,
        }
    }

    /// The manifest JSON of `img`.
    pub fn manifest(&self, img: &ImageRef) -> serde_json::Value {
        serde_json::from_slice(&std::fs::read(self.blob_path(&img.digest)).unwrap()).unwrap()
    }
}

/// The busybox applets the container tests use.
pub const APPLETS: &[&str] = &[
    "sh",
    "cp",
    "cat",
    "echo",
    "nc",
    "wget",
    "dd",
    "mkdir",
    "ln",
    "mkfifo",
    "touch",
    "sleep",
    "seq",
    "grep",
    "id",
    "hostname",
    "unshare",
    "ls",
    "true",
    "false",
    "head",
    "kill",
    "mount",
    "rm",
    "sha256sum",
    "yes",
    "printf",
    "test",
    "chmod",
    "ps",
    "readlink",
];

/// Where a static busybox is found (HEIMDALL_PLAIN_BUSYBOX or /bin/busybox).
pub fn busybox() -> Option<PathBuf> {
    let p = std::env::var_os("HEIMDALL_PLAIN_BUSYBOX")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/bin/busybox"));
    p.exists().then_some(p)
}

/// A one-layer busybox image (static busybox at /bin/busybox + applet links).
pub fn busybox_image(layout: &Layout) -> ImageRef {
    let bb = std::fs::read(busybox().expect("a static busybox")).unwrap();
    let mut entries = vec![
        Entry::Dir("bin"),
        Entry::Dir("tmp"),
        Entry::Dir("in"),
        Entry::Dir("out"),
        Entry::Dir("proc"),
        Entry::Dir("dev"),
        Entry::Dir("etc"),
        Entry::File("bin/busybox".into(), bb, 0o755),
        Entry::File("etc/marker".into(), b"from-the-image".to_vec(), 0o644),
    ];
    for a in APPLETS {
        entries.push(Entry::Symlink(format!("bin/{a}"), "busybox".into()));
    }
    layout.image(
        &[(Layout::tar(&entries), true)],
        &["PATH=/bin", "HEIMDALL_TEST_ENV=from-image"],
    )
}

/// A job directory: in/, out/.
pub struct JobDirs {
    pub base: PathBuf,
    pub input: PathBuf,
    pub output: PathBuf,
}

pub fn job_dirs(base: &Path) -> JobDirs {
    let input = base.join(format!("in-{}", JobId::random()));
    let output = base.join(format!("out-{}", JobId::random()));
    std::fs::create_dir_all(&input).unwrap();
    std::fs::create_dir_all(&output).unwrap();
    JobDirs {
        base: base.to_path_buf(),
        input,
        output,
    }
}

pub fn native_job(image: &ImageRef, argv: &[&str], d: &JobDirs) -> Job {
    let mut policy = Policy::default();
    policy.limits.wall = Duration::from_secs(20);
    policy.limits.memory_bytes = 128 << 20;
    policy.limits.pids = 64;
    Job::new(
        JobKind::Native {
            image: image.clone(),
            argv: argv.iter().map(|s| s.to_string()).collect(),
        },
        d.input.clone(),
        d.output.clone(),
        serde_json::Value::Null,
        policy,
    )
}

pub fn sh(script: &str) -> Vec<&str> {
    vec!["/bin/sh", "-c", script]
}

/// The verdict code, or "accepted".
pub fn code(r: &JobResult) -> String {
    r.verdict.code().unwrap_or("accepted").to_string()
}