lanzatool: improve tempfiles and error handling in pe

rust/lanzatool/Cargo.lock

@@ -68,6 +68,15 @@ dependencies = [
  "os_str_bytes",
 ]
 
+[[package]]
+name = "fastrand"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7a407cfaa3385c4ae6b23e84623d48c2798d06e3e6a1878f7f59f17b3f86499"
+dependencies = [
+ "instant",
+]
+
 [[package]]
 name = "goblin"
 version = "0.6.0"
@@ -94,6 +103,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "instant"
+version = "0.1.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.4"
@@ -109,6 +127,7 @@ dependencies = [
  "goblin",
  "serde",
  "serde_json",
+ "tempfile",
 ]
 
 [[package]]
@@ -186,6 +205,24 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "redox_syscall"
+version = "0.2.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a"
+dependencies = [
+ "bitflags",
+]
+
+[[package]]
+name = "remove_dir_all"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "ryu"
 version = "1.0.11"
@@ -260,6 +297,20 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "tempfile"
+version = "3.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5cdb1ef4eaeeaddc8fbd371e5017057064af0911902ef36b39801f67cc6d79e4"
+dependencies = [
+ "cfg-if",
+ "fastrand",
+ "libc",
+ "redox_syscall",
+ "remove_dir_all",
+ "winapi",
+]
+
 [[package]]
 name = "termcolor"
 version = "1.1.3"

rust/lanzatool/Cargo.toml

@@ -11,3 +11,4 @@ clap = { version = "4.0.26", features = ["derive"] }
 goblin = "0.6.0"
 serde = { version = "1.0.147", features = ["derive"] }
 serde_json = "1.0.89"
+tempfile = "3.3.0"

rust/lanzatool/src/install.rs

@@ -29,7 +29,7 @@ pub fn install(
 
     let kernel_cmdline = assemble_kernel_cmdline(bootspec_doc.init, bootspec_doc.kernel_params);
 
-    let lanzaboote_image = pe::assemble_image(
+    let lanzaboote_image = pe::lanzaboote_image(
         lanzaboote_stub,
         &bootspec_doc.extension.os_release,
         &kernel_cmdline,

rust/lanzatool/src/pe.rs

@@ -1,72 +1,108 @@
 use std::fs;
+use std::io::Write;
 use std::os::unix::fs::MetadataExt;
 use std::path::{Path, PathBuf};
 use std::process::Command;
 
 use anyhow::{Context, Result};
 use goblin::pe::PE;
+use tempfile::NamedTempFile;
 
-pub fn assemble_image(
+pub fn lanzaboote_image(
     lanzaboote_stub: &Path,
     os_release: &Path,
     kernel_cmdline: &[String],
     kernel_path: &Path,
     initrd_path: &Path,
-    esp_mountpoint: &Path
+    esp: &Path,
 ) -> Result<PathBuf> {
     // objcopy copies files into the PE binary. That's why we have to write the contents
-    // of some bootspec properties to disk
+    // of some bootspec properties to disks
-    let kernel_cmdline_file = Path::new("/tmp/kernel_cmdline");
+    let kernel_cmdline_file = write_to_tmp(kernel_cmdline.join(" "))?;
-    let kernel_path_file = Path::new("/tmp/kernel_path");
+    let kernel_path_file = write_to_tmp(esp_relative_path_string(esp, kernel_path))?;
-    let initrd_path_file = Path::new("/tmp/initrd_path");
+    let initrd_path_file = write_to_tmp(esp_relative_path_string(esp, initrd_path))?;
-
-    fs::write(kernel_cmdline_file, kernel_cmdline.join(" "))?;
-    fs::write(kernel_path_file, efi_relative_path_string(esp_mountpoint, kernel_path))?;
-    fs::write(initrd_path_file, efi_relative_path_string(esp_mountpoint, initrd_path))?;
 
     let os_release_offs = stub_offset(lanzaboote_stub)?;
+    let kernel_cmdline_offs = os_release_offs + file_size(&os_release)?;
+    let initrd_path_offs = kernel_cmdline_offs + file_size(&kernel_cmdline_file)?;
+    let kernel_path_offs = initrd_path_offs + file_size(&initrd_path_file)?;
 
-    let kernel_cmdline_offs = os_release_offs + file_size(os_release)?;
+    let sections = vec![
-    let initrd_path_offs = kernel_cmdline_offs + file_size(kernel_cmdline_file)?;
+        s(".osrel", os_release, os_release_offs),
-    let kernel_path_offs = initrd_path_offs + file_size(initrd_path_file)?;
+        s(".cmdline", kernel_cmdline_file, kernel_cmdline_offs),
-
+        s(".initrdp", initrd_path_file, initrd_path_offs),
-    let lanzaboote_image = PathBuf::from("/tmp/lanzaboote-image.efi");
+        s(".kernelp", kernel_path_file, kernel_path_offs),
-
-    let args = vec![
-        String::from("--add-section"),
-        format!(".osrel={}", path_to_string(os_release)),
-        String::from("--change-section-vma"),
-        format!(".osrel={:#x}", os_release_offs),
-        String::from("--add-section"),
-        format!(".cmdline={}", path_to_string(kernel_cmdline_file)),
-        String::from("--change-section-vma"),
-        format!(".cmdline={:#x}", kernel_cmdline_offs),
-        String::from("--add-section"),
-        format!(".initrdp={}", path_to_string(initrd_path_file)),
-        String::from("--change-section-vma"),
-        format!(".initrdp={:#x}", initrd_path_offs),
-        String::from("--add-section"),
-        format!(".kernelp={}", path_to_string(kernel_path_file)),
-        String::from("--change-section-vma"),
-        format!(".kernelp={:#x}", kernel_path_offs),
-        path_to_string(lanzaboote_stub),
-        path_to_string(&lanzaboote_image),
     ];
 
+    wrap_in_pe(&lanzaboote_stub, sections)
+}
+
+pub fn wrap_initrd(initrd_stub: &Path, initrd: &Path) -> Result<PathBuf> {
+    let initrd_offs = stub_offset(initrd_stub)?;
+    let sections = vec![s(".initrd", initrd, initrd_offs)];
+    wrap_in_pe(initrd_stub, sections)
+}
+
+fn wrap_in_pe(stub: &Path, sections: Vec<Section>) -> Result<PathBuf> {
+    let image = NamedTempFile::new().context("Failed to generate named temp file")?;
+
+    let mut args: Vec<String> = sections.iter().flat_map(Section::to_objcopy).collect();
+    let extra_args = vec![path_to_string(stub), path_to_string(&image)];
+    args.extend(extra_args);
+
     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());
+        return Err(anyhow::anyhow!("Failed to wrap in pe with args `{:?}`", &args).into());
     }
 
-    Ok(lanzaboote_image)
+    let (_, persistent_image) = image.keep().with_context(|| {
+        format!(
+            "Failed to persist image with stub: {} from temporary file",
+            stub.display()
+        )
+    })?;
+    Ok(persistent_image)
 }
 
-fn efi_relative_path_string(esp_mountpoint: &Path, path: &Path) -> String {
+struct Section {
+    name: &'static str,
+    file_path: PathBuf,
+    offset: u64,
+}
+
+impl Section {
+    fn to_objcopy(&self) -> Vec<String> {
+        vec![
+            String::from("--add-section"),
+            format!("{}={}", self.name, path_to_string(&self.file_path)),
+            String::from("--change-section-vma"),
+            format!("{}={:#x}", self.name, self.offset),
+        ]
+    }
+}
+
+fn s(name: &'static str, file_path: impl AsRef<Path>, offset: u64) -> Section {
+    Section {
+        name,
+        file_path: file_path.as_ref().into(),
+        offset,
+    }
+}
+
+fn write_to_tmp(contents: impl AsRef<[u8]>) -> Result<PathBuf> {
+    let mut tmpfile = NamedTempFile::new().context("Failed to create tempfile")?;
+    tmpfile
+        .write_all(contents.as_ref())
+        .context("Failed to write to tempfile")?;
+    Ok(tmpfile.keep()?.1)
+}
+
+fn esp_relative_path_string(esp: &Path, path: &Path) -> String {
     let relative_path = path
-        .strip_prefix(esp_mountpoint)
+        .strip_prefix(esp)
         .expect("Failed to make path relative to esp")
         .to_owned();
     let relative_path_string = relative_path
@@ -77,32 +113,6 @@ fn efi_relative_path_string(esp_mountpoint: &Path, path: &Path) -> String {
     format!("\\{}", &relative_path_string)
 }
 
-pub fn wrap_initrd(initrd_stub: &Path, initrd: &Path) -> Result<PathBuf> {
-    let initrd_offs = stub_offset(initrd_stub)?;
-
-    let wrapped_initrd = PathBuf::from("/tmp/initrd.efi");
-
-    let args = vec![
-        String::from("--add-section"),
-        format!(".initrd={}", path_to_string(initrd)),
-        String::from("--change-section-vma"),
-        format!(".initrd={:#x}", initrd_offs),
-        path_to_string(initrd_stub),
-        path_to_string(&wrapped_initrd),
-    ];
-
-    let status = Command::new("objcopy").args(&args).status()?;
-    if !status.success() {
-        return Err(anyhow::anyhow!(
-            "Failed to wrap initrd into a PE binary with args `{:?}`",
-            &args
-        )
-        .into());
-    }
-
-    Ok(wrapped_initrd)
-}
-
 fn stub_offset(binary: &Path) -> Result<u64> {
     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")?;
@@ -128,16 +138,17 @@ fn image_base(pe: &PE) -> u64 {
 }
 
 // All Linux file paths should be convertable to strings
-fn path_to_string(path: &Path) -> String {
+fn path_to_string(path: impl AsRef<Path>) -> String {
-    path.to_owned()
+    path.as_ref()
+        .to_owned()
         .into_os_string()
         .into_string()
         .expect(&format!(
             "Failed to convert path '{}' to a string",
-            path.display()
+            path.as_ref().display()
         ))
 }
 
-fn file_size(path: &Path) -> Result<u64> {
+fn file_size(path: impl AsRef<Path>) -> Result<u64> {
     Ok(fs::File::open(path)?.metadata()?.size())
 }

rust/lanzatool/src/signer.rs

@@ -30,9 +30,11 @@ impl Signer {
         let status = Command::new("sbsign").args(&args).status()?;
 
         if !status.success() {
-            return Err(
+            return Err(anyhow::anyhow!(
-                anyhow::anyhow!("Failed to sign with  sbsign with args `{:?}`", &args).into(),
+                "Failed to sign file using sbsign with args `{:?}`",
-            );
+                &args
+            )
+            .into());
         }
 
         Ok(())