nordisk/heimdall

crates/heimdall-trace/src/image.rs

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

//! A trace image: a worker image (python3, node, the package managers) plus
//! one layer holding the static tracer at [`crate::TRACER_PATH`], written
//! into the same local OCI image layout and pinned by its new manifest
//! digest. Heimdall is a defensive sandbox; the image is verified at every
//! job start like any other (heimdall-container), so the tracer that runs is
//! exactly the one this function added.

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

use heimdall_core::{ImageRef, JobId, Sha256Digest};
use serde_json::{json, Value};

const OCI_MANIFEST: &str = "application/vnd.oci.image.manifest.v1+json";
const DOCKER_MANIFEST: &str = "application/vnd.docker.distribution.manifest.v2+json";
const OCI_LAYER_TAR: &str = "application/vnd.oci.image.layer.v1.tar";
const DOCKER_LAYER_TAR: &str = "application/vnd.docker.image.rootfs.diff.tar";
const MAX_JSON_BLOB: u64 = 4 << 20;

fn digest_of(b: &[u8]) -> String {
    format!("sha256:{}", Sha256Digest::of(b).to_hex())
}

fn blob_path(layout: &Path, digest: &str) -> io::Result<PathBuf> {
    let hex = digest
        .strip_prefix("sha256:")
        .filter(|h| h.len() == 64 && h.bytes().all(|c| c.is_ascii_hexdigit()))
        .ok_or_else(|| io::Error::other(format!("not a sha256 digest: {digest:?}")))?;
    Ok(layout.join("blobs/sha256").join(hex))
}

fn read_json_blob(layout: &Path, digest: &str) -> io::Result<Value> {
    let p = blob_path(layout, digest)?;
    if std::fs::metadata(&p)?.len() > MAX_JSON_BLOB {
        return Err(io::Error::other(format!("{digest} is over 4 MiB")));
    }
    let b = std::fs::read(&p)?;
    if digest_of(&b) != digest {
        return Err(io::Error::other(format!(
            "{digest} does not match its bytes"
        )));
    }
    serde_json::from_slice(&b).map_err(|e| io::Error::other(format!("{digest}: {e}")))
}

/// Store `bytes` as a blob (written beside, then renamed: safe when several
/// callers add the same blob at once); its digest.
fn write_blob(layout: &Path, bytes: &[u8]) -> io::Result<String> {
    let d = digest_of(bytes);
    let p = blob_path(layout, &d)?;
    if !p.exists() {
        let tmp = p.with_extension(format!("tmp-{}", JobId::random().to_hex()));
        std::fs::write(&tmp, bytes)?;
        std::fs::rename(&tmp, &p)?;
    }
    Ok(d)
}

/// The tracer layer (an uncompressed tar): `heimdall/heimdall-trace`
/// (0755), plus the `in`, `out`, `tmp`, `proc` and `dev` mount points in
/// case the base image lacks them. Deterministic for one tracer binary.
pub fn tracer_layer(tracer: &[u8]) -> io::Result<Vec<u8>> {
    let mut b = tar::Builder::new(Vec::new());
    for d in ["heimdall", "in", "out", "tmp", "proc", "dev"] {
        let mut h = tar::Header::new_gnu();
        h.set_entry_type(tar::EntryType::Directory);
        h.set_mode(if d == "tmp" { 0o1777 } else { 0o755 });
        h.set_size(0);
        h.set_mtime(0);
        b.append_data(&mut h, d, &[][..])?;
    }
    let mut h = tar::Header::new_gnu();
    h.set_entry_type(tar::EntryType::Regular);
    h.set_mode(0o755);
    h.set_size(tracer.len() as u64);
    h.set_mtime(0);
    b.append_data(&mut h, crate::TRACER_PATH.trim_start_matches('/'), tracer)?;
    b.into_inner()
}

/// Add the tracer at `tracer` (a static binary) to the image `base` in the
/// OCI layout at `layout`; the new image's reference. `base` must be an OCI
/// or Docker v2 image manifest (not an index). The base's blobs stay as they
/// are; three new blobs are added (the layer, the config, the manifest).
pub fn with_tracer(layout: &Path, base: &ImageRef, tracer: &Path) -> io::Result<ImageRef> {
    let m = read_json_blob(layout, &base.digest)?;
    let media = m["mediaType"].as_str().unwrap_or(OCI_MANIFEST).to_string();
    let layer_type = match media.as_str() {
        OCI_MANIFEST => OCI_LAYER_TAR,
        DOCKER_MANIFEST => DOCKER_LAYER_TAR,
        other => {
            return Err(io::Error::other(format!(
                "the base is a {other}, not an image manifest"
            )))
        }
    };
    let cfg_digest = m["config"]["digest"]
        .as_str()
        .ok_or_else(|| io::Error::other("the base manifest has no config"))?
        .to_string();
    let mut cfg = read_json_blob(layout, &cfg_digest)?;
    let layer = tracer_layer(&std::fs::read(tracer)?)?;
    let layer_digest = write_blob(layout, &layer)?;
    // An uncompressed layer: its diff id is its digest.
    match cfg
        .pointer_mut("/rootfs/diff_ids")
        .and_then(Value::as_array_mut)
    {
        Some(ids) => ids.push(json!(layer_digest)),
        None => cfg["rootfs"] = json!({"type": "layers", "diff_ids": [layer_digest]}),
    }
    let cfg_bytes = serde_json::to_vec(&cfg).map_err(io::Error::other)?;
    let new_cfg = write_blob(layout, &cfg_bytes)?;
    let mut nm = m;
    nm["config"]["digest"] = json!(new_cfg);
    nm["config"]["size"] = json!(cfg_bytes.len());
    nm["layers"]
        .as_array_mut()
        .ok_or_else(|| io::Error::other("the base manifest has no layers"))?
        .push(json!({"mediaType": layer_type, "digest": layer_digest, "size": layer.len()}));
    let nm_bytes = serde_json::to_vec(&nm).map_err(io::Error::other)?;
    let digest = write_blob(layout, &nm_bytes)?;
    Ok(ImageRef {
        name: format!("{}+heimdall-trace", base.name),
        digest,
    })
}