From 587e388364516df878e84dbfc7d86813788a1d21 Mon Sep 17 00:00:00 2001
From: nikstur <nikstur@outlook.com>
Date: Thu, 24 Nov 2022 13:33:01 +0100
Subject: [PATCH] lanzatool: improve error handling

---
 rust/lanzatool/src/install.rs | 53 ++++++++++++++++++++++++-----------
 rust/lanzatool/src/main.rs    |  9 +-----
 rust/lanzatool/src/pe.rs      | 11 +++++---
 3 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/rust/lanzatool/src/install.rs b/rust/lanzatool/src/install.rs
index 4351757..a582632 100644
--- a/rust/lanzatool/src/install.rs
+++ b/rust/lanzatool/src/install.rs
@@ -1,8 +1,7 @@
 use std::fs;
-
 use std::path::Path;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 
 use crate::bootspec::Bootspec;
 use crate::esp::EspPaths;
@@ -14,10 +13,16 @@ pub fn install(
     lanzaboote_stub: &Path,
     initrd_stub: &Path,
 ) -> Result<()> {
-    let bootspec_doc: Bootspec = serde_json::from_slice(&fs::read(bootspec)?)?;
+    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);
 
+    println!("Assembling lanzaboote image...");
+
     let lanzaboote_image = pe::assemble_image(
         lanzaboote_stub,
         &bootspec_doc.extension.os_release,
@@ -25,29 +30,45 @@ pub fn install(
         &esp_paths.kernel,
         &esp_paths.initrd,
     )
-    .expect("Failed to assemble stub");
+    .context("Failed to assemble stub")?;
+
+    println!("Wrapping initrd into a PE binary...");
 
     let wrapped_initrd =
-        pe::wrap_initrd(initrd_stub, &bootspec_doc.initrd).expect("Failed to assemble stub");
+        pe::wrap_initrd(initrd_stub, &bootspec_doc.initrd).context("Failed to assemble stub")?;
 
-    // Copy the files to the ESP
-    fs::create_dir_all(&esp_paths.nixos)?;
-    fs::copy(bootspec_doc.kernel, esp_paths.kernel)?;
-    fs::copy(wrapped_initrd, esp_paths.initrd)?;
-
-    fs::create_dir_all(&esp_paths.linux)?;
-    fs::copy(lanzaboote_image, esp_paths.lanzaboote_image)?;
+    println!("Copy files to EFI system partition...");
 
     let systemd_boot = bootspec_doc
         .extension
         .systemd
         .join("lib/systemd/boot/efi/systemd-bootx64.efi");
 
-    fs::create_dir_all(esp_paths.efi_fallback_dir)?;
-    fs::copy(&systemd_boot, esp_paths.efi_fallback)?;
+    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),
+    ];
 
-    fs::create_dir_all(&esp_paths.systemd)?;
-    fs::copy(&systemd_boot, esp_paths.systemd_boot)?;
+    for (source, target) in files_to_copy {
+        copy(&source, &target)?;
+    }
 
+    println!(
+        "Succesfully installed lanzaboote to '{}'",
+        esp_paths.esp.display()
+    );
+    Ok(())
+}
+
+fn copy(from: &Path, to: &Path) -> Result<()> {
+    match to.parent() {
+        Some(parent) => fs::create_dir_all(parent).unwrap_or(()),
+        _ => (),
+    };
+    fs::copy(from, to)
+        .with_context(|| format!("Failed to copy from {} to {}", from.display(), to.display()))?;
     Ok(())
 }
diff --git a/rust/lanzatool/src/main.rs b/rust/lanzatool/src/main.rs
index f99e4a5..b75dad1 100644
--- a/rust/lanzatool/src/main.rs
+++ b/rust/lanzatool/src/main.rs
@@ -4,18 +4,11 @@ mod esp;
 mod install;
 mod pe;
 
-use std::process;
-
 use anyhow::Result;
 use clap::Parser;
 
 use cli::Cli;
 
 fn main() -> Result<()> {
-    let cli = Cli::parse();
-    if let Err(e) = cli.call() {
-        eprintln!("{}", e);
-        process::exit(1)
-    };
-    Ok(())
+    Cli::parse().call()
 }
diff --git a/rust/lanzatool/src/pe.rs b/rust/lanzatool/src/pe.rs
index 2ffc89b..6ae8479 100644
--- a/rust/lanzatool/src/pe.rs
+++ b/rust/lanzatool/src/pe.rs
@@ -3,7 +3,7 @@ use std::os::unix::fs::MetadataExt;
 use std::path::{Path, PathBuf};
 use std::process::Command;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use goblin::pe::PE;
 
 pub fn assemble_image(
@@ -52,7 +52,10 @@ pub fn assemble_image(
         path_to_string(&lanzaboote_image),
     ];
 
-    let status = Command::new("objcopy").args(&args).status()?;
+    let status = Command::new("objcopy")
+        .args(&args)
+        .status()
+        .context("Failed to run objcopy command")?;
     if !status.success() {
         return Err(anyhow::anyhow!("Failed to build stub with args `{:?}`", &args).into());
     }
@@ -100,8 +103,8 @@ pub fn wrap_initrd(initrd_stub: &Path, initrd: &Path) -> Result<PathBuf> {
 }
 
 fn stub_offset(binary: &Path) -> Result<u64> {
-    let pe_binary = fs::read(binary)?;
-    let pe = PE::parse(&pe_binary)?;
+    let pe_binary = fs::read(binary).context("Failed to read PE binary file")?;
+    let pe = PE::parse(&pe_binary).context("Failed to parse PE binary file")?;
 
     let image_base = image_base(&pe);