path: root/src/initramfs.rs

use std::ffi::OsString;
use std::fs;
use std::os::unix::fs::symlink;
use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{Context, Result, bail};
use sha2::{Digest, Sha256};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KernelVersion(String);

impl KernelVersion {
    pub fn discover_rt(rootfs: &Path) -> Result<Self> {
        let boot = rootfs.join("boot");
        let entries = fs::read_dir(&boot)
            .with_context(|| format!("read boot directory {}", boot.display()))?;
        let mut kernels = entries
            .filter_map(|entry| entry.ok())
            .filter_map(|entry| {
                let file_type = entry.file_type().ok()?;
                if !file_type.is_file() {
                    return None;
                }
                let name = entry.file_name();
                let name = name.to_str()?;
                let version = name.strip_prefix("vmlinuz-")?;
                (version.contains("rt")).then(|| version.to_owned())
            })
            .collect::<Vec<_>>();
        kernels.sort();
        match kernels.as_slice() {
            [] => bail!("no RT kernel image found in {}", boot.display()),
            [version] => Ok(Self(version.clone())),
            _ => bail!(
                "multiple RT kernel images found in {}: {}",
                boot.display(),
                kernels.join(", ")
            ),
        }
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsRecipe {
    features: Vec<String>,
    modules: Vec<String>,
}

impl InitramfsRecipe {
    pub fn new(
        features: impl IntoIterator<Item = impl Into<String>>,
        modules: impl IntoIterator<Item = impl Into<String>>,
    ) -> Result<Self> {
        let mut features = normalize("initrd feature", features)?;
        let mut modules = normalize("initrd module", modules)?;
        features.sort();
        features.dedup();
        modules.sort();
        modules.dedup();
        Ok(Self { features, modules })
    }

    pub fn render(&self) -> String {
        format!(
            "FEATURES += {}\nMODULES += {}\n",
            self.features.join(" "),
            self.modules.join(" ")
        )
    }
}

fn normalize(
    kind: &str,
    values: impl IntoIterator<Item = impl Into<String>>,
) -> Result<Vec<String>> {
    values
        .into_iter()
        .map(Into::into)
        .map(|value: String| {
            if value.trim().is_empty() || value.chars().any(char::is_whitespace) {
                bail!("{kind} cannot be empty or contain whitespace");
            }
            Ok(value)
        })
        .collect()
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsRequest {
    rootfs: PathBuf,
    kernel: KernelVersion,
}

impl InitramfsRequest {
    pub fn discover(rootfs: impl AsRef<Path>) -> Result<Self> {
        let rootfs = rootfs.as_ref();
        if !rootfs.is_dir() {
            bail!("rootfs is not a directory: {}", rootfs.display());
        }
        Ok(Self {
            rootfs: rootfs.to_path_buf(),
            kernel: KernelVersion::discover_rt(rootfs)?,
        })
    }

    pub fn rootfs(&self) -> &Path {
        &self.rootfs
    }

    pub fn kernel(&self) -> &KernelVersion {
        &self.kernel
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Invocation {
    program: OsString,
    arguments: Vec<OsString>,
}

impl Invocation {
    pub fn new(
        program: impl Into<OsString>,
        arguments: impl IntoIterator<Item = impl Into<OsString>>,
    ) -> Self {
        Self {
            program: program.into(),
            arguments: arguments.into_iter().map(Into::into).collect(),
        }
    }
}

pub trait CommandRunner {
    fn run(&mut self, invocation: Invocation) -> Result<()>;
}

#[derive(Debug, Default)]
pub struct ProcessRunner;

impl CommandRunner for ProcessRunner {
    fn run(&mut self, invocation: Invocation) -> Result<()> {
        let status = Command::new(&invocation.program)
            .args(&invocation.arguments)
            .status()
            .with_context(|| format!("run {}", invocation.program.to_string_lossy()))?;
        if !status.success() {
            bail!(
                "{} exited with {status}",
                invocation.program.to_string_lossy()
            );
        }
        Ok(())
    }
}

pub trait InitramfsBuilder {
    fn build(&mut self, request: &InitramfsRequest) -> Result<InitramfsResult>;
}

#[derive(Debug)]
pub struct MakeInitrdBuilder<R> {
    runner: R,
}

impl<R> MakeInitrdBuilder<R> {
    pub fn new(runner: R) -> Self {
        Self { runner }
    }

    pub fn runner(&self) -> &R {
        &self.runner
    }
}

impl<R: CommandRunner> InitramfsBuilder for MakeInitrdBuilder<R> {
    fn build(&mut self, request: &InitramfsRequest) -> Result<InitramfsResult> {
        self.runner.run(Invocation::new(
            "chroot",
            [
                request.rootfs().as_os_str().to_owned(),
                "make-initrd".into(),
                "-k".into(),
                request.kernel().as_str().into(),
            ],
        ))?;
        let result = InitramfsResult::from_rootfs(request.rootfs(), request.kernel())?;
        write_boot_alias(
            request.rootfs(),
            "vmlinuz",
            Path::new(&format!("vmlinuz-{}", request.kernel().as_str())),
        )?;
        write_boot_alias(
            request.rootfs(),
            "initrd.img",
            Path::new(result.initrd_path().file_name().expect("initrd filename")),
        )?;
        Ok(result)
    }
}

/// Set a conventional boot alias without ever following or replacing a
/// pre-existing non-symlink filesystem entry.
fn write_boot_alias(rootfs: &Path, alias: &str, target: &Path) -> Result<()> {
    let path = rootfs.join("boot").join(alias);
    match fs::symlink_metadata(&path) {
        Ok(metadata) if metadata.file_type().is_symlink() => {
            fs::remove_file(&path)
                .with_context(|| format!("remove existing boot alias {}", path.display()))?;
        }
        Ok(_) => bail!(
            "refusing to replace non-symlink boot alias: {}",
            path.display()
        ),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
        Err(error) => {
            return Err(error).with_context(|| format!("inspect boot alias {}", path.display()));
        }
    }
    symlink(target, &path).with_context(|| format!("create boot alias {}", path.display()))
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsResult {
    initrd_path: PathBuf,
    sha256: String,
}

impl InitramfsResult {
    /// Record the generated initrd as a rootfs-relative semantic fact.
    pub fn from_rootfs(rootfs: &Path, kernel: &KernelVersion) -> Result<Self> {
        let initrd_path = PathBuf::from(format!("boot/initrd-{}.img", kernel.as_str()));
        let contents = fs::read(rootfs.join(&initrd_path))
            .with_context(|| format!("read generated initrd {}", initrd_path.display()))?;
        Ok(Self {
            initrd_path,
            sha256: format!("{:x}", Sha256::digest(contents)),
        })
    }

    pub fn initrd_path(&self) -> &Path {
        &self.initrd_path
    }

    pub fn sha256(&self) -> &str {
        &self.sha256
    }

    /// Ensure the initrd recorded by the platform adapter still exists in the
    /// assembled rootfs and has not changed before native packaging begins.
    pub fn verify_in_rootfs(&self, rootfs: &Path) -> Result<()> {
        let contents = fs::read(rootfs.join(&self.initrd_path)).with_context(|| {
            format!(
                "read generated initrd {}",
                rootfs.join(&self.initrd_path).display()
            )
        })?;
        let digest = format!("{:x}", Sha256::digest(contents));
        if digest != self.sha256 {
            bail!(
                "generated initrd digest changed for {}: expected {}, found {digest}",
                self.initrd_path.display(),
                self.sha256
            );
        }
        Ok(())
    }
}