Merge pull request #14 from blitz/lanzatool-make-it-more-typedriven

lanzatool: make it more typedriven

rust/lanzatool/src/cli.rs

@@ -1,4 +1,4 @@
-use std::path::{Path, PathBuf};
+use std::path::PathBuf;
 
 use anyhow::{Context, Result};
 use clap::{Parser, Subcommand};
@@ -8,33 +8,36 @@ use crate::install;
 #[derive(Parser)]
 pub struct Cli {
     #[clap(subcommand)]
-    pub commands: Commands,
+    commands: Commands,
 }
 
 #[derive(Subcommand)]
-pub enum Commands {
-    Install {
+enum Commands {
+    Install(InstallCommand),
+}
+
+#[derive(Parser)]
+struct InstallCommand {
     // Secure Boot Public Key
-        #[clap(long)]
+    #[arg(long)]
     public_key: PathBuf,
 
     // Secure Boot Private Key
-        #[clap(long)]
+    #[arg(long)]
     private_key: PathBuf,
 
     // Secure Boot PKI Bundle for auto enrolling key
-        #[clap(long)]
+    #[arg(long)]
     pki_bundle: Option<PathBuf>,
 
     // Enable auto enrolling your keys in UEFI
     // Be aware that this might irrevocably brick your device
-        #[clap(long, default_value = "false")]
+    #[arg(long, default_value = "false")]
     auto_enroll: bool,
 
     bootspec: PathBuf,
 
     generations: Vec<PathBuf>,
-    },
 }
 
 impl Cli {
@@ -46,46 +49,26 @@ impl Cli {
 impl Commands {
     pub fn call(self) -> Result<()> {
         match self {
-            Commands::Install {
-                public_key,
-                private_key,
-                pki_bundle,
-                auto_enroll,
-                bootspec,
-                generations,
-            } => install(
-                &public_key,
-                &private_key,
-                &pki_bundle,
-                auto_enroll,
-                &bootspec,
-                generations,
-            ),
+            Commands::Install(args) => install(args),
         }
     }
 }
 
-fn install(
-    public_key: &Path,
-    private_key: &Path,
-    pki_bundle: &Option<PathBuf>,
-    auto_enroll: bool,
-    bootspec: &Path,
-    generations: Vec<PathBuf>,
-) -> Result<()> {
+fn install(args: InstallCommand) -> Result<()> {
     let lanzaboote_stub =
         std::env::var("LANZABOOTE_STUB").context("Failed to read LANZABOOTE_STUB env variable")?;
     let initrd_stub = std::env::var("LANZABOOTE_INITRD_STUB")
         .context("Failed to read LANZABOOTE_INITRD_STUB env variable")?;
 
-    install::install(
-        public_key,
-        private_key,
-        pki_bundle,
-        auto_enroll,
-        bootspec,
-        generations,
-        Path::new(&lanzaboote_stub),
-        Path::new(&initrd_stub),
+    install::Installer::new(
+        PathBuf::from(lanzaboote_stub),
+        PathBuf::from(initrd_stub),
+        args.public_key,
+        args.private_key,
+        args.pki_bundle,
+        args.auto_enroll,
+        args.bootspec,
+        args.generations,
     )
+    .install()
 }

rust/lanzatool/src/esp.rs

@@ -2,6 +2,7 @@ use anyhow::{Context, Result};
 use std::path::{Path, PathBuf};
 
 use crate::bootspec::Bootspec;
+use crate::generation::Generation;
 
 pub struct EspPaths {
     pub esp: PathBuf,
@@ -17,7 +18,7 @@ pub struct EspPaths {
 }
 
 impl EspPaths {
-    pub fn new(esp: &str, generation: u64, bootspec: &Bootspec) -> Result<Self> {
+    pub fn new(esp: &str, generation: Generation, bootspec: &Bootspec) -> Result<Self> {
         let esp = Path::new(esp);
         let esp_nixos = esp.join("EFI/nixos");
         let esp_linux = esp.join("EFI/Linux");
@@ -59,6 +60,6 @@ fn nixos_path(path: impl AsRef<Path>, name: &str) -> Result<PathBuf> {
     Ok(PathBuf::from(nixos_filename))
 }
 
-fn generation_path(generation: u64) -> PathBuf {
+fn generation_path(generation: Generation) -> PathBuf {
     PathBuf::from(format!("nixos-generation-{}.efi", generation))
 }

rust/lanzatool/src/generation.rs

@@ -0,0 +1,36 @@
+use std::fmt;
+use std::path::Path;
+
+use anyhow::{Context, Result};
+
+#[derive(Debug)]
+pub struct Generation(u64);
+
+impl Generation {
+    pub fn from_toplevel(toplevel: impl AsRef<Path>) -> Result<Self> {
+        let file_name = toplevel.as_ref().file_name().ok_or(anyhow::anyhow!(
+            "Failed to extract file name from generation"
+        ))?;
+
+        let file_name_str = file_name
+            .to_str()
+            .with_context(|| "Failed to convert file name of generation to string")?;
+
+        let generation_version = file_name_str
+            .split("-")
+            .nth(1)
+            .ok_or(anyhow::anyhow!("Failed to extract version from generation"))?;
+
+        let parsed_generation_version = generation_version.parse().with_context(|| {
+            format!("Failed to parse generation version: {}", generation_version)
+        })?;
+
+        Ok(Self(parsed_generation_version))
+    }
+}
+
+impl fmt::Display for Generation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}

rust/lanzatool/src/install.rs

@@ -1,86 +1,72 @@
 use std::fs;
 use std::os::unix::prelude::PermissionsExt;
 use std::path::{Path, PathBuf};
-use nix::unistd::sync;
+use std::process::Command;
 
 use anyhow::{Context, Result};
+use nix::unistd::sync;
+use tempfile::tempdir;
 
 use crate::bootspec::Bootspec;
 use crate::esp::EspPaths;
+use crate::generation::Generation;
 use crate::pe;
 use crate::signer::Signer;
 
-use tempfile::tempdir;
-
-use std::process::Command;
-
-pub fn install(
-    public_key: &Path,
-    private_key: &Path,
-    pki_bundle: &Option<PathBuf>,
-    auto_enroll: bool,
-    bootspec: &Path,
+pub struct Installer {
+    lanzaboote_stub: PathBuf,
+    initrd_stub: PathBuf,
+    public_key: PathBuf,
+    private_key: PathBuf,
+    _pki_bundle: Option<PathBuf>,
+    _auto_enroll: bool,
+    bootspec: PathBuf,
     generations: Vec<PathBuf>,
-    lanzaboote_stub: &Path,
-    initrd_stub: &Path,
-) -> Result<()> {
-    for generation in generations {
-        let generation_version = extract_generation_version(&generation).with_context(|| {
-            format!(
-                "Failed to extract generation version from generation: {}",
-                generation.display()
-            )
-        })?;
+}
 
-        println!("Installing generation {generation_version}");
-
-        install_generation(
-            generation_version,
-            public_key,
-            private_key,
-            pki_bundle,
-            auto_enroll,
-            bootspec,
+impl Installer {
+    pub fn new(
+        lanzaboote_stub: PathBuf,
+        initrd_stub: PathBuf,
+        public_key: PathBuf,
+        private_key: PathBuf,
+        _pki_bundle: Option<PathBuf>,
+        _auto_enroll: bool,
+        bootspec: PathBuf,
+        generations: Vec<PathBuf>,
+    ) -> Self {
+        Self {
             lanzaboote_stub,
             initrd_stub,
-        )?;
+            public_key,
+            private_key,
+            _pki_bundle,
+            _auto_enroll,
+            bootspec,
+            generations,
+        }
+    }
+
+    pub fn install(&self) -> Result<()> {
+        for toplevel in &self.generations {
+            let generation = Generation::from_toplevel(toplevel).with_context(|| {
+                format!("Failed to extract generation version from: {toplevel:?}")
+            })?;
+
+            println!("Installing generation {generation}");
+
+            self.install_generation(generation)?
         }
 
         Ok(())
     }
 
-fn extract_generation_version(path: impl AsRef<Path>) -> Result<u64> {
-    let file_name = path.as_ref().file_name().ok_or(anyhow::anyhow!(
-        "Failed to extract file name from generation"
-    ))?;
-    let file_name_str = file_name
-        .to_str()
-        .with_context(|| "Failed to convert file name of generation to string")?;
-
-    let generation_version = file_name_str
-        .split("-")
-        .nth(1)
-        .ok_or(anyhow::anyhow!("Failed to extract version from generation"))?;
-
-    Ok(generation_version
-        .parse()
-        .with_context(|| format!("Failed to parse generation version: {}", generation_version))?)
-}
-
-fn install_generation(
-    generation: u64,
-    public_key: &Path,
-    private_key: &Path,
-    _pki_bundle: &Option<PathBuf>,
-    _auto_enroll: bool,
-    bootspec: &Path,
-    lanzaboote_stub: &Path,
-    initrd_stub: &Path,
-) -> Result<()> {
+    pub fn install_generation(&self, generation: Generation) -> Result<()> {
         println!("Reading bootspec...");
 
-    let bootspec_doc: Bootspec =
-        serde_json::from_slice(&fs::read(bootspec).context("Failed to read bootspec file")?)
+        let bootspec_doc: Bootspec = serde_json::from_slice(
+            &fs::read(&self.bootspec).context("Failed to read bootspec file")?,
+        )
         .context("Failed to parse bootspec json")?;
 
         let esp_paths = EspPaths::new(&bootspec_doc.extension.esp, generation, &bootspec_doc)?;
@@ -99,7 +85,7 @@ fn install_generation(
 
         let lanzaboote_image = pe::lanzaboote_image(
             &secure_temp_dir,
-        lanzaboote_stub,
+            &self.lanzaboote_stub,
             &bootspec_doc.extension.os_release,
             &kernel_cmdline,
             &esp_paths.kernel,
@@ -113,11 +99,10 @@ fn install_generation(
         let initrd_location = secure_temp_dir.path().join("initrd");
         copy(&bootspec_doc.initrd, &initrd_location)?;
         if let Some(initrd_secrets_script) = bootspec_doc.initrd_secrets {
-        append_initrd_secrets(&initrd_secrets_script,
-            &initrd_location)?;
+            append_initrd_secrets(&initrd_secrets_script, &initrd_location)?;
         }
-    let wrapped_initrd =
-        pe::wrap_initrd(&secure_temp_dir, initrd_stub, &initrd_location).context("Failed to assemble stub")?;
+        let wrapped_initrd = pe::wrap_initrd(&secure_temp_dir, &self.initrd_stub, &initrd_location)
+            .context("Failed to assemble stub")?;
 
         println!("Copy files to EFI system partition...");
 
@@ -143,7 +128,7 @@ fn install_generation(
 
         println!("Signing files...");
 
-    let signer = Signer::new(&public_key, &private_key);
+        let signer = Signer::new(&self.public_key, &self.private_key);
 
         let files_to_sign = [
             &esp_paths.efi_fallback,
@@ -168,16 +153,22 @@ fn install_generation(
 
         Ok(())
     }
+}
 
-pub fn append_initrd_secrets(append_initrd_secrets_path: &Path, initrd_path: &PathBuf) -> Result<()> {
+pub fn append_initrd_secrets(
+    append_initrd_secrets_path: &Path,
+    initrd_path: &PathBuf,
+) -> Result<()> {
     let status = Command::new(append_initrd_secrets_path)
-        .args(vec![
-            initrd_path
-        ])
+        .args(vec![initrd_path])
         .status()
         .context("Failed to append initrd secrets")?;
     if !status.success() {
-        return Err(anyhow::anyhow!("Failed to append initrd secrets with args `{:?}`", vec![append_initrd_secrets_path, initrd_path]).into());
+        return Err(anyhow::anyhow!(
+            "Failed to append initrd secrets with args `{:?}`",
+            vec![append_initrd_secrets_path, initrd_path]
+        )
+        .into());
     }
 
     Ok(())

rust/lanzatool/src/main.rs

@@ -1,6 +1,7 @@
 mod bootspec;
 mod cli;
 mod esp;
+mod generation;
 mod install;
 mod pe;
 mod signer;