Merge pull request #42 from nix-community/lanzatool-small-fixes

Lanzatool: small fixes

rust/lanzatool/src/esp.rs

@@ -53,21 +53,13 @@ fn nixos_path(path: impl AsRef<Path>, name: &str) -> Result<PathBuf> {
         .read_link()
         .unwrap_or_else(|_| path.as_ref().into());
 
-    let parent = resolved.parent().ok_or_else(|| {
-        anyhow::anyhow!(format!(
-            "Path: {} does not have a parent",
-            resolved.display()
-        ))
-    })?;
+    let parent_final_component = resolved
+        .parent()
+        .and_then(|x| x.file_name())
+        .and_then(|x| x.to_str())
+        .with_context(|| format!("Failed to extract final component from: {:?}", resolved))?;
 
-    let without_store = parent.strip_prefix("/nix/store").with_context(|| {
-        format!(
-            "Failed to strip /nix/store from path {}",
-            path.as_ref().display()
-        )
-    })?;
-
-    let nixos_filename = format!("{}-{}.efi", without_store.display(), name);
+    let nixos_filename = format!("{}-{}.efi", parent_final_component, name);
 
     Ok(PathBuf::from(nixos_filename))
 }
@@ -82,3 +74,22 @@ fn generation_path(generation: &Generation) -> PathBuf {
         PathBuf::from(format!("nixos-generation-{}.efi", generation))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn nixos_path_creates_correct_filename_from_nix_store_path() -> Result<()> {
+        let path =
+            Path::new("/nix/store/xqplddjjjy1lhzyzbcv4dza11ccpcfds-initrd-linux-6.1.1/initrd");
+
+        let generated_filename = nixos_path(path, "initrd")?;
+
+        let expected_filename =
+            PathBuf::from("xqplddjjjy1lhzyzbcv4dza11ccpcfds-initrd-linux-6.1.1-initrd.efi");
+
+        assert_eq!(generated_filename, expected_filename);
+        Ok(())
+    }
+}

rust/lanzatool/src/generation.rs

@@ -87,24 +87,29 @@ impl fmt::Display for Generation {
     }
 }
 
-fn parse_version(toplevel: impl AsRef<Path>) -> Result<u64> {
-    let file_name = toplevel
+/// Parse version number from a path.
+///
+/// Expects a path in the format of "system-{version}-link".
+fn parse_version(path: impl AsRef<Path>) -> Result<u64> {
+    let generation_version = path
         .as_ref()
         .file_name()
-        .ok_or_else(|| anyhow::anyhow!("Failed to extract file name from generation"))?;
+        .and_then(|x| x.to_str())
+        .and_then(|x| x.split('-').nth(1))
+        .and_then(|x| x.parse::<u64>().ok())
+        .with_context(|| format!("Failed to extract version from: {:?}", path.as_ref()))?;
 
-    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_else(|| 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(parsed_generation_version)
+    Ok(generation_version)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_version_correctly() {
+        let path = Path::new("system-2-link");
+        let parsed_version = parse_version(path).unwrap();
+        assert_eq!(parsed_version, 2,);
+    }
 }

rust/lanzatool/src/pe.rs

@@ -32,7 +32,7 @@ pub fn lanzaboote_image(
     let kernel_path_file = write_to_tmp(
         target_dir,
         "kernel-esp-path",
-        esp_relative_path_string(esp, kernel_path),
+        esp_relative_path_string(esp, kernel_path)?,
     )?;
     let kernel_hash_file = write_to_tmp(
         target_dir,
@@ -43,7 +43,7 @@ pub fn lanzaboote_image(
     let initrd_path_file = write_to_tmp(
         target_dir,
         "initrd-esp-path",
-        esp_relative_path_string(esp, initrd_path),
+        esp_relative_path_string(esp, initrd_path)?,
     )?;
     let initrd_hash_file = write_to_tmp(
         target_dir,
@@ -167,17 +167,17 @@ fn write_to_tmp(
     Ok(path)
 }
 
-fn esp_relative_path_string(esp: &Path, path: &Path) -> String {
+fn esp_relative_path_string(esp: &Path, path: &Path) -> Result<String> {
     let relative_path = path
         .strip_prefix(esp)
-        .expect("Failed to make path relative to esp")
+        .with_context(|| format!("Failed to strip prefix: {:?} from path: {:?}", esp, path))?
         .to_owned();
     let relative_path_string = relative_path
         .into_os_string()
         .into_string()
         .expect("Failed to convert path '{}' to a relative string path")
         .replace('/', "\\");
-    format!("\\{}", &relative_path_string)
+    Ok(format!("\\{}", &relative_path_string))
 }
 
 fn stub_offset(binary: &Path) -> Result<u64> {
@@ -205,5 +205,13 @@ fn image_base(pe: &PE) -> u64 {
 }
 
 fn file_size(path: impl AsRef<Path>) -> Result<u64> {
-    Ok(fs::File::open(path)?.metadata()?.size())
+    Ok(fs::File::open(&path)
+        .with_context(|| {
+            format!(
+                "Failed to read file to calculate its size: {:?}",
+                path.as_ref()
+            )
+        })?
+        .metadata()?
+        .size())
 }

rust/lanzatool/src/signature.rs

@@ -1,4 +1,5 @@
 use std::ffi::OsString;
+use std::io::Write;
 use std::path::{Path, PathBuf};
 use std::process::Command;
 
@@ -31,7 +32,7 @@ impl KeyPair {
         let output = Command::new("sbsign").args(&args).output()?;
 
         if !output.status.success() {
-            print!("{:?}", output.stderr);
+            std::io::stderr().write_all(&output.stderr).unwrap();
             return Err(anyhow::anyhow!(
                 "Failed to sign file using sbsign with args `{:?}`",
                 &args