lanzatool: implement configuration limit

2022-12-12 02:09:19 +01:00 · 2022-12-12 02:09:19 +01:00 · 9daf9ae0a8
parent 4a8cfa7f7f
commit 9daf9ae0a8
6 changed files with 289 additions and 40 deletions
--- a/rust/lanzatool/src/cli.rs
+++ b/rust/lanzatool/src/cli.rs
@ -27,10 +27,14 @@ struct InstallCommand {
    #[arg(long)]
    private_key: PathBuf,

+    /// Configuration limit
+    #[arg(long, default_value_t = 1)]
+    configuration_limit: usize,
+
    /// EFI system partition mountpoint (e.g. efiSysMountPoint)
    esp: PathBuf,

-    /// List of generations (e.g. /nix/var/nix/profiles/system-*-link)
+    /// List of generation links (e.g. /nix/var/nix/profiles/system-*-link)
    generations: Vec<PathBuf>,
 }

@ -57,6 +61,7 @@ fn install(args: InstallCommand) -> Result<()> {
    install::Installer::new(
        PathBuf::from(lanzaboote_stub),
        key_pair,
+        args.configuration_limit,
        args.esp,
        args.generations,
    )
--- a/rust/lanzatool/src/esp.rs
+++ b/rust/lanzatool/src/esp.rs
@ -1,10 +1,13 @@
-use anyhow::{Context, Result};
+use std::array::IntoIter;
 use std::path::{Path, PathBuf};

+use anyhow::{Context, Result};
+
 use crate::generation::Generation;

 pub struct EspPaths {
    pub esp: PathBuf,
+    pub efi: PathBuf,
    pub nixos: PathBuf,
    pub kernel: PathBuf,
    pub initrd: PathBuf,
@ -19,32 +22,52 @@ pub struct EspPaths {
 impl EspPaths {
    pub fn new(esp: impl AsRef<Path>, generation: &Generation) -> Result<Self> {
        let esp = esp.as_ref();
-        let esp_nixos = esp.join("EFI/nixos");
-        let esp_linux = esp.join("EFI/Linux");
-        let esp_systemd = esp.join("EFI/systemd");
-        let esp_efi_fallback_dir = esp.join("EFI/BOOT");
+        let efi = esp.join("EFI");
+        let efi_nixos = efi.join("nixos");
+        let efi_linux = efi.join("Linux");
+        let efi_systemd = efi.join("systemd");
+        let efi_efi_fallback_dir = efi.join("BOOT");

        let bootspec = &generation.spec.bootspec;

        Ok(Self {
            esp: esp.to_path_buf(),
-            nixos: esp_nixos.clone(),
-            kernel: esp_nixos.join(nixos_path(&bootspec.kernel, "bzImage")?),
-            initrd: esp_nixos.join(nixos_path(
+            efi,
+            nixos: efi_nixos.clone(),
+            kernel: efi_nixos.join(nixos_path(&bootspec.kernel, "bzImage")?),
+            initrd: efi_nixos.join(nixos_path(
                bootspec
                    .initrd
                    .as_ref()
                    .context("Lanzaboote does not support missing initrd yet")?,
                "initrd",
            )?),
-            linux: esp_linux.clone(),
-            lanzaboote_image: esp_linux.join(generation_path(generation)),
-            efi_fallback_dir: esp_efi_fallback_dir.clone(),
-            efi_fallback: esp_efi_fallback_dir.join("BOOTX64.EFI"),
-            systemd: esp_systemd.clone(),
-            systemd_boot: esp_systemd.join("systemd-bootx64.efi"),
+            linux: efi_linux.clone(),
+            lanzaboote_image: efi_linux.join(generation_path(generation)),
+            efi_fallback_dir: efi_efi_fallback_dir.clone(),
+            efi_fallback: efi_efi_fallback_dir.join("BOOTX64.EFI"),
+            systemd: efi_systemd.clone(),
+            systemd_boot: efi_systemd.join("systemd-bootx64.efi"),
        })
    }
+
+    /// Return the used file paths to store as garbage collection roots.
+    pub fn to_iter(&self) -> IntoIter<&PathBuf, 11> {
+        [
+            &self.esp,
+            &self.efi,
+            &self.nixos,
+            &self.kernel,
+            &self.initrd,
+            &self.linux,
+            &self.lanzaboote_image,
+            &self.efi_fallback_dir,
+            &self.efi_fallback,
+            &self.systemd,
+            &self.systemd_boot,
+        ]
+        .into_iter()
+    }
 }

 fn nixos_path(path: impl AsRef<Path>, name: &str) -> Result<PathBuf> {
--- a/rust/lanzatool/src/gc.rs
+++ b/rust/lanzatool/src/gc.rs
@ -0,0 +1,160 @@
+use std::collections::HashSet;
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use anyhow::{Context, Result};
+use walkdir::{DirEntry, WalkDir};
+
+/// Keeps track of the garbage collection roots.
+///
+/// The internal HashSet contains all the paths still in use. These paths
+/// are used to find all **unused** paths and delete them.
+#[derive(Debug)]
+pub struct Roots(HashSet<PathBuf>);
+
+impl Roots {
+    pub fn new() -> Self {
+        Self(HashSet::new())
+    }
+
+    /// Extend the garbage collection roots.
+    ///
+    /// Not only the file paths of roots themselves, but also all parent directories that should
+    /// not be garbage collected need to be **explicitly** added to the roots. For example, if you
+    /// have a path: `rootdir/example/file.txt`, the three paths: `rootdir`, `rootdir/example`, and
+    /// `rootdir/example/file.txt` need to be added for the right files to be garbage collected.
+    pub fn extend<'a>(&mut self, other: impl IntoIterator<Item = &'a PathBuf>) {
+        self.0.extend(other.into_iter().cloned());
+    }
+
+    fn in_use(&self, entry: Option<&DirEntry>) -> bool {
+        match entry {
+            Some(e) => self.0.contains(e.path()),
+            None => false,
+        }
+    }
+
+    pub fn collect_garbage(&self, directory: impl AsRef<Path>) -> Result<()> {
+        // Find all the paths not used anymore.
+        let entries_not_in_use = WalkDir::new(directory.as_ref())
+            .into_iter()
+            .filter(|e| !self.in_use(e.as_ref().ok()));
+
+        // Remove all entries not in use.
+        for e in entries_not_in_use {
+            let entry = e?;
+            let path = entry.path();
+            println!("'{}' not in use anymore. Removing...", path.display());
+
+            if path.is_dir() {
+                // If a directory is marked as unused all its children can be deleted too.
+                fs::remove_dir_all(path)
+                    .with_context(|| format!("Failed to remove directory: {:?}", path))?;
+            } else {
+                // Ignore failing to remove path because the parent directory might have been removed before.
+                fs::remove_file(path).ok();
+            };
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+
+    #[test]
+    fn keep_used_file() -> Result<()> {
+        let tmpdir = tempfile::tempdir()?;
+        let rootdir = create_dir(tmpdir.path().join("root"))?;
+
+        let used_file = create_file(rootdir.join("root_file"))?;
+
+        let mut roots = Roots::new();
+        roots.extend(vec![&rootdir, &used_file]);
+        roots.collect_garbage(&rootdir)?;
+
+        assert!(used_file.exists());
+        Ok(())
+    }
+
+    #[test]
+    fn delete_unused_file() -> Result<()> {
+        let tmpdir = tempfile::tempdir()?;
+        let rootdir = create_dir(tmpdir.path().join("root"))?;
+
+        let unused_file = create_file(rootdir.join("unused_file"))?;
+
+        let mut roots = Roots::new();
+        roots.extend(vec![&rootdir]);
+        roots.collect_garbage(&rootdir)?;
+
+        assert!(!unused_file.exists());
+        Ok(())
+    }
+
+    #[test]
+    fn delete_empty_unused_directory() -> Result<()> {
+        let tmpdir = tempfile::tempdir()?;
+        let rootdir = create_dir(tmpdir.path().join("root"))?;
+
+        let unused_directory = create_dir(rootdir.join("unused_directory"))?;
+
+        let mut roots = Roots::new();
+        roots.extend(vec![&rootdir]);
+        roots.collect_garbage(&rootdir)?;
+
+        assert!(!unused_directory.exists());
+        Ok(())
+    }
+
+    #[test]
+    fn delete_unused_directory_with_unused_file_inside() -> Result<()> {
+        let tmpdir = tempfile::tempdir()?;
+        let rootdir = create_dir(tmpdir.path().join("root"))?;
+
+        let unused_directory = create_dir(rootdir.join("unused_directory"))?;
+        let unused_file_in_directory =
+            create_file(unused_directory.join("unused_file_in_directory"))?;
+
+        let mut roots = Roots::new();
+        roots.extend(vec![&rootdir]);
+        roots.collect_garbage(&rootdir)?;
+
+        assert!(!unused_directory.exists());
+        assert!(!unused_file_in_directory.exists());
+        Ok(())
+    }
+
+    #[test]
+    fn keep_used_dirctory_with_used_and_unused_file() -> Result<()> {
+        let tmpdir = tempfile::tempdir()?;
+        let rootdir = create_dir(tmpdir.path().join("root"))?;
+
+        let used_directory = create_dir(rootdir.join("used_directory"))?;
+        let used_file_in_directory = create_file(used_directory.join("used_file_in_directory"))?;
+        let unused_file_in_directory =
+            create_file(used_directory.join("unused_file_in_directory"))?;
+
+        let mut roots = Roots::new();
+        roots.extend(vec![&rootdir, &used_directory, &used_file_in_directory]);
+        roots.collect_garbage(&rootdir)?;
+
+        assert!(used_directory.exists());
+        assert!(used_file_in_directory.exists());
+        assert!(!unused_file_in_directory.exists());
+        Ok(())
+    }
+
+    fn create_file(path: PathBuf) -> Result<PathBuf> {
+        fs::File::create(&path)?;
+        Ok(path)
+    }
+
+    fn create_dir(path: PathBuf) -> Result<PathBuf> {
+        fs::create_dir(&path)?;
+        Ok(path)
+    }
+}
--- a/rust/lanzatool/src/generation.rs
+++ b/rust/lanzatool/src/generation.rs
@ -1,6 +1,3 @@
-use serde::de::IntoDeserializer;
-use serde::{Deserialize, Serialize};
-
 use std::fmt;
 use std::fs;
 use std::path::{Path, PathBuf};
@ -9,6 +6,8 @@ use anyhow::{anyhow, Context, Result};
 use bootspec::generation::Generation as BootspecGeneration;
 use bootspec::BootJson;
 use bootspec::SpecialisationName;
+use serde::de::IntoDeserializer;
+use serde::{Deserialize, Serialize};

 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SecureBootExtension {
@ -22,6 +21,14 @@ pub struct ExtendedBootJson {
    pub extensions: SecureBootExtension,
 }

+/// A system configuration.
+///
+/// Can be built from a GenerationLink.
+///
+/// NixOS represents a generation as a symlink to a toplevel derivation. This toplevel derivation
+/// contains most of the information necessary to install the generation onto the EFI System
+/// Partition. The only information missing is the version number which is encoded in the file name
+/// of the generation link.
 #[derive(Debug)]
 pub struct Generation {
    /// Profile symlink index
@ -33,17 +40,8 @@ pub struct Generation {
 }

 impl Generation {
-    fn extract_extensions(bootspec: &BootJson) -> Result<SecureBootExtension> {
-        Ok(Deserialize::deserialize(
-            bootspec.extensions.get("lanzaboote")
-            .context("Failed to extract Lanzaboote-specific extension from Bootspec, missing lanzaboote field in `extensions`")?
-            .clone()
-            .into_deserializer()
-        )?)
-    }
-
-    pub fn from_toplevel(toplevel: impl AsRef<Path>) -> Result<Self> {
-        let bootspec_path = toplevel.as_ref().join("boot.json");
+    pub fn from_link(link: &GenerationLink) -> Result<Self> {
+        let bootspec_path = link.path.join("boot.json");
        let generation: BootspecGeneration = serde_json::from_slice(
            &fs::read(bootspec_path).context("Failed to read bootspec file")?,
        )
@ -56,7 +54,7 @@ impl Generation {
        let extensions = Self::extract_extensions(&bootspec)?;

        Ok(Self {
-            version: parse_version(toplevel)?,
+            version: link.version,
            specialisation_name: None,
            spec: ExtendedBootJson {
                bootspec,
@ -65,6 +63,15 @@ impl Generation {
        })
    }

+    fn extract_extensions(bootspec: &BootJson) -> Result<SecureBootExtension> {
+        Ok(Deserialize::deserialize(
+            bootspec.extensions.get("lanzaboote")
+            .context("Failed to extract Lanzaboote-specific extension from Bootspec, missing lanzaboote field in `extensions`")?
+            .clone()
+            .into_deserializer()
+        )?)
+    }
+
    pub fn specialise(&self, name: &SpecialisationName, bootspec: &BootJson) -> Result<Self> {
        Ok(Self {
            version: self.version,
@ -87,6 +94,25 @@ impl fmt::Display for Generation {
    }
 }

+/// A link pointing to a generation.
+///
+/// Can be built from a symlink in /nix/var/nix/profiles/ alone because the name of the
+/// symlink enocdes the version number.
+#[derive(Debug)]
+pub struct GenerationLink {
+    pub version: u64,
+    pub path: PathBuf,
+}
+
+impl GenerationLink {
+    pub fn from_path(path: impl AsRef<Path>) -> Result<Self> {
+        Ok(Self {
+            version: parse_version(&path).context("Failed to parse version")?,
+            path: PathBuf::from(path.as_ref()),
+        })
+    }
+}
+
 /// Parse version number from a path.
 ///
 /// Expects a path in the format of "system-{version}-link".
--- a/rust/lanzatool/src/install.rs
+++ b/rust/lanzatool/src/install.rs
@ -8,37 +8,71 @@ use nix::unistd::sync;
 use tempfile::tempdir;

 use crate::esp::EspPaths;
-use crate::generation::Generation;
+use crate::gc::Roots;
+use crate::generation::{Generation, GenerationLink};
 use crate::pe;
 use crate::signature::KeyPair;

 pub struct Installer {
+    gc_roots: Roots,
    lanzaboote_stub: PathBuf,
    key_pair: KeyPair,
+    configuration_limit: usize,
    esp: PathBuf,
-    generations: Vec<PathBuf>,
+    generation_links: Vec<PathBuf>,
 }

 impl Installer {
    pub fn new(
        lanzaboote_stub: PathBuf,
        key_pair: KeyPair,
+        configuration_limit: usize,
        esp: PathBuf,
-        generations: Vec<PathBuf>,
+        generation_links: Vec<PathBuf>,
    ) -> Self {
        Self {
+            gc_roots: Roots::new(),
            lanzaboote_stub,
            key_pair,
+            configuration_limit,
            esp,
-            generations,
+            generation_links,
        }
    }

-    pub fn install(&self) -> Result<()> {
-        for toplevel in &self.generations {
-            let generation_result = Generation::from_toplevel(toplevel)
-                .with_context(|| format!("Failed to build generation from toplevel: {toplevel:?}"));
+    pub fn install(&mut self) -> Result<()> {
+        let mut links = self
+            .generation_links
+            .iter()
+            .map(GenerationLink::from_path)
+            .collect::<Result<Vec<GenerationLink>>>()?;

+        // A configuration limit of 0 means there is no limit.
+        if self.configuration_limit > 0 {
+            // Sort the links by version.
+            links.sort_by_key(|l| l.version);
+
+            // Only install the number of generations configured.
+            links = links
+                .into_iter()
+                .rev()
+                .take(self.configuration_limit)
+                .collect()
+        };
+        self.install_links(links)?;
+
+        self.gc_roots.collect_garbage(&self.esp)?;
+
+        Ok(())
+    }
+
+    fn install_links(&mut self, links: Vec<GenerationLink>) -> Result<()> {
+        for link in links {
+            let generation_result = Generation::from_link(&link)
+                .with_context(|| format!("Failed to build generation from link: {link:?}"));
+
+            // Ignore failing to read a generation so that old malformed generations do not stop
+            // lanzatool from working.
            let generation = match generation_result {
                Ok(generation) => generation,
                Err(e) => {
@ -61,15 +95,15 @@ impl Installer {
                    .context("Failed to install specialisation")?;
            }
        }
-
        Ok(())
    }

-    fn install_generation(&self, generation: &Generation) -> Result<()> {
+    fn install_generation(&mut self, generation: &Generation) -> Result<()> {
        let bootspec = &generation.spec.bootspec;
        let secureboot_extensions = &generation.spec.extensions;

        let esp_paths = EspPaths::new(&self.esp, generation)?;
+        self.gc_roots.extend(esp_paths.to_iter());

        let kernel_cmdline =
            assemble_kernel_cmdline(&bootspec.init, bootspec.kernel_params.clone());
--- a/rust/lanzatool/src/main.rs
+++ b/rust/lanzatool/src/main.rs
@ -1,5 +1,6 @@
 mod cli;
 mod esp;
+mod gc;
 mod generation;
 mod install;
 mod pe;