Add native rootfs finalization

1 files changed, 211 insertions, 0 deletions

diff --git a/src/rootfs.rs b/src/rootfs.rs
new file mode 100644
index 0000000..06b2557
--- /dev/null
+++ b/src/rootfs.rs
@@ -0,0 +1,211 @@
+use std::collections::BTreeSet;
+use std::fs;
+use std::os::unix::fs::symlink;
+use std::path::{Path, PathBuf};
+
+use anyhow::{bail, Context, Result};
+
+use crate::files::{relative_symlink_is_safe, RootfsPath};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CopyTree {
+    source: PathBuf,
+    destination: RootfsPath,
+}
+
+impl CopyTree {
+    pub fn new(source: impl AsRef<Path>, destination: impl AsRef<Path>) -> Result<Self> {
+        Ok(Self {
+            source: source.as_ref().to_path_buf(),
+            destination: RootfsPath::new(destination)?,
+        })
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct InitrdOem {
+    features: BTreeSet<String>,
+    modules: BTreeSet<String>,
+}
+
+impl InitrdOem {
+    pub fn new(
+        features: impl IntoIterator<Item = impl Into<String>>,
+        modules: impl IntoIterator<Item = impl Into<String>>,
+    ) -> Self {
+        Self {
+            features: features.into_iter().map(Into::into).collect(),
+            modules: modules.into_iter().map(Into::into).collect(),
+        }
+    }
+
+    fn render(&self) -> String {
+        format!(
+            "FEATURES += {}\nMODULES += {}\n",
+            self.features.iter().cloned().collect::<Vec<_>>().join(" "),
+            self.modules.iter().cloned().collect::<Vec<_>>().join(" "),
+        )
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ServiceName(String);
+
+impl ServiceName {
+    pub fn new(name: impl Into<String>) -> Result<Self> {
+        let name = name.into();
+        if !name.ends_with(".service")
+            || name.contains('/')
+            || name.chars().any(char::is_whitespace)
+            || name == ".service"
+        {
+            bail!("invalid systemd service name: {name}");
+        }
+        Ok(Self(name))
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RootfsFinalization {
+    copies: Vec<CopyTree>,
+    initrd_oem: InitrdOem,
+    enabled_services: Vec<ServiceName>,
+}
+
+impl RootfsFinalization {
+    pub fn new(
+        copies: Vec<CopyTree>,
+        initrd_oem: InitrdOem,
+        enabled_services: Vec<ServiceName>,
+    ) -> Self {
+        Self {
+            copies,
+            initrd_oem,
+            enabled_services,
+        }
+    }
+
+    pub fn apply(&self, rootfs: &Path) -> Result<MutationManifest> {
+        if !rootfs.is_dir() {
+            bail!("rootfs is not a directory: {}", rootfs.display());
+        }
+
+        let mut manifest = MutationManifest::default();
+        for copy in &self.copies {
+            copy_tree(rootfs, copy, &mut manifest)?;
+        }
+        write_file(
+            rootfs,
+            &RootfsPath::new("etc/initrd.mk.oem")?,
+            self.initrd_oem.render().as_bytes(),
+            &mut manifest,
+        )?;
+        for service in &self.enabled_services {
+            enable_service(rootfs, service, &mut manifest)?;
+        }
+        manifest.paths.sort();
+        Ok(manifest)
+    }
+}
+
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
+pub struct MutationManifest {
+    paths: Vec<String>,
+}
+
+impl MutationManifest {
+    pub fn created_paths(&self) -> &[String] {
+        &self.paths
+    }
+
+    fn record(&mut self, path: &Path) {
+        self.paths.push(path.display().to_string());
+    }
+}
+
+fn copy_tree(rootfs: &Path, copy: &CopyTree, manifest: &mut MutationManifest) -> Result<()> {
+    if !copy.source.is_dir() {
+        bail!("copy source is not a directory: {}", copy.source.display());
+    }
+    copy_entry(rootfs, &copy.source, copy.destination.as_path(), Path::new(""), manifest)
+}
+
+fn copy_entry(
+    rootfs: &Path,
+    source: &Path,
+    destination: &Path,
+    source_relative: &Path,
+    manifest: &mut MutationManifest,
+) -> Result<()> {
+    let metadata = fs::symlink_metadata(source)
+        .with_context(|| format!("inspect copy source {}", source.display()))?;
+    let output = rootfs.join(destination);
+
+    if metadata.file_type().is_dir() {
+        fs::create_dir_all(&output).with_context(|| format!("create {}", output.display()))?;
+        manifest.record(destination);
+        let mut entries = fs::read_dir(source)
+            .with_context(|| format!("read copy source directory {}", source.display()))?
+            .collect::<std::result::Result<Vec<_>, _>>()?;
+        entries.sort_by_key(|entry| entry.file_name());
+        for entry in entries {
+            let name = entry.file_name();
+            copy_entry(
+                rootfs,
+                &entry.path(),
+                &destination.join(&name),
+                &source_relative.join(&name),
+                manifest,
+            )?;
+        }
+    } else if metadata.file_type().is_symlink() {
+        let target = fs::read_link(source)
+            .with_context(|| format!("read source symlink {}", source.display()))?;
+        if !relative_symlink_is_safe(source_relative.parent().unwrap_or(Path::new("")), &target) {
+            bail!("source symlink escapes copied tree: {}", source.display());
+        }
+        if let Some(parent) = output.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        symlink(&target, &output).with_context(|| format!("create symlink {}", output.display()))?;
+        manifest.record(destination);
+    } else if metadata.file_type().is_file() {
+        if let Some(parent) = output.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        fs::copy(source, &output)
+            .with_context(|| format!("copy {} to {}", source.display(), output.display()))?;
+        fs::set_permissions(&output, metadata.permissions())?;
+        manifest.record(destination);
+    } else {
+        bail!("unsupported copy source type: {}", source.display());
+    }
+    Ok(())
+}
+
+fn write_file(
+    rootfs: &Path,
+    destination: &RootfsPath,
+    contents: &[u8],
+    manifest: &mut MutationManifest,
+) -> Result<()> {
+    let output = rootfs.join(destination.as_path());
+    let parent = output.parent().expect("rootfs-relative path has a parent");
+    fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
+    fs::write(&output, contents).with_context(|| format!("write {}", output.display()))?;
+    manifest.record(destination.as_path());
+    Ok(())
+}
+
+fn enable_service(rootfs: &Path, service: &ServiceName, manifest: &mut MutationManifest) -> Result<()> {
+    let destination = RootfsPath::new(format!(
+        "etc/systemd/system/multi-user.target.wants/{}",
+        service.0
+    ))?;
+    let output = rootfs.join(destination.as_path());
+    fs::create_dir_all(output.parent().expect("service path has a parent"))?;
+    symlink(format!("/usr/lib/systemd/system/{}", service.0), &output)
+        .with_context(|| format!("enable service {}", service.0))?;
+    manifest.record(destination.as_path());
+    Ok(())
+}