lanzatool: init wrapping initrd

rust/lanzatool/src/cli.rs

@@ -37,6 +37,12 @@ impl Commands {
 }
 
 fn install(public_key: &Path, bootspec: &Path) -> Result<()> {
-    let lanzaboote_bin = std::env::var("LANZABOOTE")?;
-    install::install(public_key, bootspec, Path::new(&lanzaboote_bin))
+    let lanzaboote_stub = std::env::var("LANZABOOTE_STUB")?;
+    let initrd_stub = std::env::var("LANZABOOTE_INITRD_STUB")?;
+    install::install(
+        public_key,
+        bootspec,
+        Path::new(&lanzaboote_stub),
+        Path::new(&initrd_stub),
+    )
 }

rust/lanzatool/src/install.rs

@@ -7,15 +7,20 @@ use anyhow::Result;
 
 use crate::bootspec::Bootspec;
 use crate::esp::EspPaths;
-use crate::stub;
+use crate::pe;
 
-pub fn install(_: &Path, bootspec: &Path, lanzaboote_bin: &Path) -> Result<()> {
+pub fn install(
+    _: &Path,
+    bootspec: &Path,
+    lanzaboote_stub: &Path,
+    initrd_stub: &Path,
+) -> Result<()> {
     let bootspec_doc: Bootspec = serde_json::from_slice(&fs::read(bootspec)?)?;
 
     let esp_paths = EspPaths::new(&bootspec_doc.v1.extension.esp);
 
-    let lanzaboote_image = stub::assemble(
-        lanzaboote_bin,
+    let lanzaboote_image = pe::assemble_image(
+        lanzaboote_stub,
         &bootspec_doc.v1.extension.os_release,
         &bootspec_doc.v1.kernel_params,
         &esp_paths.kernel,
@@ -23,10 +28,13 @@ pub fn install(_: &Path, bootspec: &Path, lanzaboote_bin: &Path) -> Result<()> {
     )
     .expect("Failed to assemble stub");
 
+    let wrapped_initrd =
+        pe::wrap_initrd(initrd_stub, &bootspec_doc.v1.initrd).expect("Failed to assemble stub");
+
     // Copy the files to the ESP
     fs::create_dir_all(&esp_paths.nixos)?;
     fs::copy(bootspec_doc.v1.kernel, esp_paths.kernel)?;
-    fs::copy(bootspec_doc.v1.initrd, esp_paths.initrd)?;
+    fs::copy(wrapped_initrd, esp_paths.initrd)?;
 
     fs::create_dir_all(&esp_paths.linux)?;
     fs::copy(lanzaboote_image, esp_paths.lanzaboote_image)?;

rust/lanzatool/src/main.rs

@@ -2,7 +2,7 @@ mod bootspec;
 mod cli;
 mod esp;
 mod install;
-mod stub;
+mod pe;
 
 use std::process;
 

rust/lanzatool/src/pe.rs

@@ -4,10 +4,10 @@ use std::path::{Path, PathBuf};
 use std::process::Command;
 
 use anyhow::Result;
-use goblin;
+use goblin::pe::PE;
 
-pub fn assemble(
-    lanzaboote_bin: &Path,
+pub fn assemble_image(
+    lanzaboote_stub: &Path,
     os_release: &Path,
     kernel_cmdline: &[String],
     kernel_path: &Path,
@@ -23,25 +23,7 @@ pub fn assemble(
     fs::write(kernel_path_file, efi_relative_path_string(kernel_path))?;
     fs::write(initrd_path_file, efi_relative_path_string(initrd_path))?;
 
-    let pe_binary = fs::read(lanzaboote_bin)?;
-    let pe = goblin::pe::PE::parse(&pe_binary)?;
-
-    let image_base = pe
-        .header
-        .optional_header
-        .expect("Failed to find optional header, you're fucked")
-        .windows_fields
-        .image_base;
-
-    let os_release_offs = u64::from(
-        pe.sections
-            .last()
-            .and_then(|s| Some(s.virtual_size + s.virtual_address))
-            .expect("Failed to find the offset"),
-    );
-    // The Virtual Memory Addresss (VMA) is relative ot the image base, aka the image base
-    // needs to be added to the virtual address to get the actual (but still virtual address)
-    let os_release_offs = os_release_offs + image_base;
+    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)?;
@@ -66,7 +48,7 @@ pub fn assemble(
         format!(".kernelp={}", path_to_string(kernel_path_file)),
         String::from("--change-section-vma"),
         format!(".kernelp={:#x}", kernel_path_offs),
-        path_to_string(lanzaboote_bin),
+        path_to_string(lanzaboote_stub),
         path_to_string(&lanzaboote_image),
     ];
 
@@ -78,17 +60,54 @@ pub fn assemble(
     Ok(lanzaboote_image)
 }
 
-fn efi_relative_path_string(path: &Path) -> String {
-    let relative_path = path
-        .strip_prefix("esp")
-        .expect("Failed to make path relative to esp")
-        .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)
+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)?;
+    let pe = PE::parse(&pe_binary)?;
+
+    let image_base = image_base(&pe);
+
+    // The Virtual Memory Addresss (VMA) is relative to the image base, aka the image base
+    // needs to be added to the virtual address to get the actual (but still virtual address)
+    Ok(u64::from(
+        pe.sections
+            .last()
+            .and_then(|s| Some(s.virtual_size + s.virtual_address))
+            .expect("Failed to calculate offset"),
+    ) + image_base)
+}
+
+fn image_base(pe: &PE) -> u64 {
+    pe.header
+        .optional_header
+        .expect("Failed to find optional header, you're fucked")
+        .windows_fields
+        .image_base
 }
 
 // All Linux file paths should be convertable to strings