Extract Linux kernel and initrd filenames from PE binary

rust/lanzaboote/src/main.rs

@@ -9,18 +9,20 @@ mod linux_loader;
 mod pe_section;
 mod uefi_helpers;
 
+use pe_section::pe_section_as_string;
 use uefi::{
     prelude::*,
     proto::{
         console::text::Output,
         loaded_image::LoadedImage,
-        media::file::{File, FileAttribute, FileMode},
+        media::file::{File, FileAttribute, FileMode, RegularFile},
     },
+    CString16,
 };
 
 use crate::{
     linux_loader::InitrdLoader,
-    uefi_helpers::{booted_image_cmdline, read_all},
+    uefi_helpers::{booted_image_cmdline, booted_image_file, read_all},
 };
 
 fn print_logo(output: &mut Output) {
@@ -41,27 +43,61 @@ fn print_logo(output: &mut Output) {
         .unwrap();
 }
 
+struct EmbeddedConfiguration {
+    kernel_filename: CString16,
+    initrd_filename: CString16,
+}
+
+impl TryFrom<&mut RegularFile> for EmbeddedConfiguration {
+    type Error = uefi::Error;
+
+    fn try_from(file: &mut RegularFile) -> Result<Self, Self::Error> {
+        file.set_position(0)?;
+        let file_data = read_all(file)?;
+
+        let kernel_filename =
+            pe_section_as_string(&file_data, ".kernelp").ok_or(Status::INVALID_PARAMETER)?;
+        let initrd_filename =
+            pe_section_as_string(&file_data, ".initrdp").ok_or(Status::INVALID_PARAMETER)?;
+
+        Ok(Self {
+            kernel_filename: CString16::try_from(kernel_filename.as_str()).unwrap(),
+            initrd_filename: CString16::try_from(initrd_filename.as_str()).unwrap(),
+        })
+    }
+}
+
 #[entry]
 fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
     uefi_services::init(&mut system_table).unwrap();
 
     print_logo(system_table.stdout());
 
+    let config: EmbeddedConfiguration = {
+        let mut booted_image = booted_image_file(system_table.boot_services()).unwrap();
+
+        (&mut booted_image).try_into().unwrap()
+    };
+
     let mut file_system = system_table
         .boot_services()
         .get_image_file_system(handle)
         .unwrap();
     let mut root = file_system.open_volume().unwrap();
 
-    let mut file = root
-        .open(cstr16!("linux.efi"), FileMode::Read, FileAttribute::empty())
+    let mut kernel_file = root
+        .open(
+            &config.kernel_filename,
+            FileMode::Read,
+            FileAttribute::empty(),
+        )
         .unwrap()
         .into_regular_file()
         .unwrap();
 
     let initrd = root
         .open(
-            cstr16!("initrd.efi"),
+            &config.initrd_filename,
             FileMode::Read,
             FileAttribute::empty(),
         )
@@ -75,7 +111,7 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
 
     let kernel_cmdline = booted_image_cmdline(system_table.boot_services()).unwrap();
 
-    let kernel_data = read_all(&mut file).unwrap();
+    let kernel_data = read_all(&mut kernel_file).unwrap();
     let kernel_handle = system_table
         .boot_services()
         .load_image(

rust/lanzaboote/src/pe_section.rs

@@ -1,3 +1,6 @@
+use alloc::{borrow::ToOwned, string::String};
+
+/// Extracts the data of a section of a PE file.
 pub fn pe_section<'a>(file_data: &'a [u8], section_name: &str) -> Option<&'a [u8]> {
     let pe_binary = goblin::pe::PE::parse(file_data).ok()?;
 
@@ -12,3 +15,9 @@ pub fn pe_section<'a>(file_data: &'a [u8], section_name: &str) -> Option<&'a [u8
             Some(&file_data[section_start..section_end])
         })
 }
+
+/// Extracts the data of a section of a PE file and returns it as a string.
+pub fn pe_section_as_string<'a>(file_data: &'a [u8], section_name: &str) -> Option<String> {
+    pe_section(file_data, section_name)
+        .and_then(|data| Some(core::str::from_utf8(data).unwrap().to_owned()))
+}