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 {
+enum Commands {
-    Install {
+    Install(InstallCommand),
-        // Secure Boot Public Key
+}
-        #[clap(long)]
-        public_key: PathBuf,
 
-        // Secure Boot Private Key
+#[derive(Parser)]
-        #[clap(long)]
+struct InstallCommand {
-        private_key: PathBuf,
+    // Secure Boot Public Key
+    #[arg(long)]
+    public_key: PathBuf,
 
-        // Secure Boot PKI Bundle for auto enrolling key
+    // Secure Boot Private Key
-        #[clap(long)]
+    #[arg(long)]
-        pki_bundle: Option<PathBuf>,
+    private_key: PathBuf,
 
-        // Enable auto enrolling your keys in UEFI
+    // Secure Boot PKI Bundle for auto enrolling key
-        // Be aware that this might irrevocably brick your device
+    #[arg(long)]
-        #[clap(long, default_value = "false")]
+    pki_bundle: Option<PathBuf>,
-        auto_enroll: bool,
 
-        bootspec: PathBuf,
+    // Enable auto enrolling your keys in UEFI
+    // Be aware that this might irrevocably brick your device
+    #[arg(long, default_value = "false")]
+    auto_enroll: bool,
 
-        generations: Vec<PathBuf>,
+    bootspec: PathBuf,
-    },
+
+    generations: Vec<PathBuf>,
 }
 
 impl Cli {
@@ -46,46 +49,26 @@ impl Cli {
 impl Commands {
     pub fn call(self) -> Result<()> {
         match self {
-            Commands::Install {
+            Commands::Install(args) => install(args),
-                public_key,
-                private_key,
-                pki_bundle,
-                auto_enroll,
-                bootspec,
-                generations,
-            } => install(
-                &public_key,
-                &private_key,
-                &pki_bundle,
-                auto_enroll,
-                &bootspec,
-                generations,
-            ),
         }
     }
 }
 
-fn install(
+fn install(args: InstallCommand) -> Result<()> {
-    public_key: &Path,
-    private_key: &Path,
-    pki_bundle: &Option<PathBuf>,
-    auto_enroll: bool,
-    bootspec: &Path,
-    generations: Vec<PathBuf>,
-) -> 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(
+    install::Installer::new(
-        public_key,
+        PathBuf::from(lanzaboote_stub),
-        private_key,
+        PathBuf::from(initrd_stub),
-        pki_bundle,
+        args.public_key,
-        auto_enroll,
+        args.private_key,
-        bootspec,
+        args.pki_bundle,
-        generations,
+        args.auto_enroll,
-        Path::new(&lanzaboote_stub),
+        args.bootspec,
-        Path::new(&initrd_stub),
+        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,183 +1,174 @@
 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;
+pub struct Installer {
-
+    lanzaboote_stub: PathBuf,
-use std::process::Command;
+    initrd_stub: PathBuf,
-
+    public_key: PathBuf,
-pub fn install(
+    private_key: PathBuf,
-    public_key: &Path,
+    _pki_bundle: Option<PathBuf>,
-    private_key: &Path,
+    _auto_enroll: bool,
-    pki_bundle: &Option<PathBuf>,
+    bootspec: PathBuf,
-    auto_enroll: bool,
-    bootspec: &Path,
     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}");
+impl Installer {
-
+    pub fn new(
-        install_generation(
+        lanzaboote_stub: PathBuf,
-            generation_version,
+        initrd_stub: PathBuf,
-            public_key,
+        public_key: PathBuf,
-            private_key,
+        private_key: PathBuf,
-            pki_bundle,
+        _pki_bundle: Option<PathBuf>,
-            auto_enroll,
+        _auto_enroll: bool,
-            bootspec,
+        bootspec: PathBuf,
+        generations: Vec<PathBuf>,
+    ) -> Self {
+        Self {
             lanzaboote_stub,
             initrd_stub,
-        )?;
+            public_key,
+            private_key,
+            _pki_bundle,
+            _auto_enroll,
+            bootspec,
+            generations,
+        }
     }
 
-    Ok(())
+    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(())
+    }
+
+    pub fn install_generation(&self, generation: Generation) -> Result<()> {
+        println!("Reading bootspec...");
+
+        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)?;
+
+        println!("Assembling lanzaboote image...");
+
+        let kernel_cmdline = assemble_kernel_cmdline(bootspec_doc.init, bootspec_doc.kernel_params);
+
+        // prepare a secure temporary directory
+        // permission bits are not set, because when files below
+        // are opened, they are opened with 600 mode bits.
+        // hence, they cannot be read except by the current user
+        // which is assumed to be root in most cases.
+        // TODO(Raito): prove to niksnur this is actually acceptable.
+        let secure_temp_dir = tempdir()?;
+
+        let lanzaboote_image = pe::lanzaboote_image(
+            &secure_temp_dir,
+            &self.lanzaboote_stub,
+            &bootspec_doc.extension.os_release,
+            &kernel_cmdline,
+            &esp_paths.kernel,
+            &esp_paths.initrd,
+            &esp_paths.esp,
+        )
+        .context("Failed to assemble stub")?;
+
+        println!("Wrapping initrd into a PE binary...");
+
+        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)?;
+        }
+        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...");
+
+        let systemd_boot = bootspec_doc
+            .extension
+            .systemd
+            .join("lib/systemd/boot/efi/systemd-bootx64.efi");
+
+        let files_to_copy = [
+            (bootspec_doc.kernel, &esp_paths.kernel),
+            (wrapped_initrd, &esp_paths.initrd),
+            (lanzaboote_image, &esp_paths.lanzaboote_image),
+            (systemd_boot.clone(), &esp_paths.efi_fallback),
+            (systemd_boot, &esp_paths.systemd_boot),
+        ];
+
+        for (source, target) in files_to_copy {
+            copy(&source, &target)?;
+        }
+
+        // TODO: we should implement sign_and_copy which would be secure
+        // by construction for TOCTOU.
+
+        println!("Signing files...");
+
+        let signer = Signer::new(&self.public_key, &self.private_key);
+
+        let files_to_sign = [
+            &esp_paths.efi_fallback,
+            &esp_paths.systemd_boot,
+            &esp_paths.lanzaboote_image,
+            &esp_paths.kernel,
+            &esp_paths.initrd,
+        ];
+
+        for file in files_to_sign {
+            println!("Signing {}...", file.display());
+            signer
+                .sign_file(&file)
+                .with_context(|| format!("Failed to sign file {}", &file.display()))?;
+            sync();
+        }
+
+        println!(
+            "Successfully installed lanzaboote to '{}'",
+            esp_paths.esp.display()
+        );
+
+        Ok(())
+    }
 }
 
-fn extract_generation_version(path: impl AsRef<Path>) -> Result<u64> {
+pub fn append_initrd_secrets(
-    let file_name = path.as_ref().file_name().ok_or(anyhow::anyhow!(
+    append_initrd_secrets_path: &Path,
-        "Failed to extract file name from generation"
+    initrd_path: &PathBuf,
-    ))?;
-    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<()> {
-    println!("Reading bootspec...");
-
-    let bootspec_doc: Bootspec =
-        serde_json::from_slice(&fs::read(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)?;
-
-    println!("Assembling lanzaboote image...");
-
-    let kernel_cmdline = assemble_kernel_cmdline(bootspec_doc.init, bootspec_doc.kernel_params);
-
-    // prepare a secure temporary directory
-    // permission bits are not set, because when files below
-    // are opened, they are opened with 600 mode bits.
-    // hence, they cannot be read except by the current user
-    // which is assumed to be root in most cases.
-    // TODO(Raito): prove to niksnur this is actually acceptable.
-    let secure_temp_dir = tempdir()?;
-
-    let lanzaboote_image = pe::lanzaboote_image(
-        &secure_temp_dir,
-        lanzaboote_stub,
-        &bootspec_doc.extension.os_release,
-        &kernel_cmdline,
-        &esp_paths.kernel,
-        &esp_paths.initrd,
-        &esp_paths.esp,
-    )
-    .context("Failed to assemble stub")?;
-
-    println!("Wrapping initrd into a PE binary...");
-
-    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)?;
-    }
-    let wrapped_initrd =
-        pe::wrap_initrd(&secure_temp_dir, initrd_stub, &initrd_location).context("Failed to assemble stub")?;
-
-    println!("Copy files to EFI system partition...");
-
-    let systemd_boot = bootspec_doc
-        .extension
-        .systemd
-        .join("lib/systemd/boot/efi/systemd-bootx64.efi");
-
-    let files_to_copy = [
-        (bootspec_doc.kernel, &esp_paths.kernel),
-        (wrapped_initrd, &esp_paths.initrd),
-        (lanzaboote_image, &esp_paths.lanzaboote_image),
-        (systemd_boot.clone(), &esp_paths.efi_fallback),
-        (systemd_boot, &esp_paths.systemd_boot),
-    ];
-
-    for (source, target) in files_to_copy {
-        copy(&source, &target)?;
-    }
-
-    // TODO: we should implement sign_and_copy which would be secure
-    // by construction for TOCTOU.
-
-    println!("Signing files...");
-
-    let signer = Signer::new(&public_key, &private_key);
-
-    let files_to_sign = [
-        &esp_paths.efi_fallback,
-        &esp_paths.systemd_boot,
-        &esp_paths.lanzaboote_image,
-        &esp_paths.kernel,
-        &esp_paths.initrd,
-    ];
-
-    for file in files_to_sign {
-        println!("Signing {}...", file.display());
-        signer
-            .sign_file(&file)
-            .with_context(|| format!("Failed to sign file {}", &file.display()))?;
-        sync();
-    }
-
-    println!(
-        "Successfully installed lanzaboote to '{}'",
-        esp_paths.esp.display()
-    );
-
-    Ok(())
-}
-
-pub fn append_initrd_secrets(append_initrd_secrets_path: &Path, initrd_path: &PathBuf) -> Result<()> {
     let status = Command::new(append_initrd_secrets_path)
-        .args(vec![
+        .args(vec![initrd_path])
-            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;