lanzaboote: load kernel and initrd into memory only once

rust/lanzaboote/src/linux_loader.rs

@@ -7,20 +7,17 @@
 
 use core::{ffi::c_void, ops::Range, pin::Pin, ptr::slice_from_raw_parts_mut};
 
-use alloc::boxed::Box;
+use alloc::{boxed::Box, vec::Vec};
 use uefi::{
     prelude::BootServices,
     proto::{
         device_path::{DevicePath, FfiDevicePath},
-        media::file::RegularFile,
         Protocol,
     },
     table::boot::LoadImageSource,
     unsafe_guid, Handle, Identify, Result, ResultExt, Status,
 };
 
-use crate::uefi_helpers::read_all;
-
 /// The Linux kernel's initrd loading device path.
 ///
 /// The Linux kernel points us to
@@ -65,23 +62,10 @@ struct LoadFile2Protocol {
     ) -> Status,
 
     // This is not part of the official protocol struct.
-    file: RegularFile,
+    initrd_data: Vec<u8>,
-    range: Range<usize>,
 }
 
 impl LoadFile2Protocol {
-    fn initrd_start(&self) -> usize {
-        self.range.start
-    }
-
-    fn initrd_size(&self) -> usize {
-        if self.range.is_empty() {
-            0
-        } else {
-            self.range.end - self.range.start
-        }
-    }
-
     fn load_file(
         &mut self,
         _file_path: *const FfiDevicePath,
@@ -89,24 +73,21 @@ impl LoadFile2Protocol {
         buffer_size: *mut usize,
         buffer: *mut c_void,
     ) -> Result<()> {
-        if buffer.is_null() || unsafe { *buffer_size } < self.initrd_size() {
+        if buffer.is_null() || unsafe { *buffer_size } < self.initrd_data.len() {
             unsafe {
-                *buffer_size = self.initrd_size();
+                *buffer_size = self.initrd_data.len();
             }
             return Err(Status::BUFFER_TOO_SMALL.into());
         };
 
-        self.file
-            .set_position(self.initrd_start().try_into().unwrap())?;
         unsafe {
-            *buffer_size = self.initrd_size();
+            *buffer_size = self.initrd_data.len();
         }
 
         let output_slice: &mut [u8] =
             unsafe { &mut *slice_from_raw_parts_mut(buffer as *mut u8, *buffer_size) };
 
-        let read_bytes = self.file.read(output_slice).map_err(|e| e.status())?;
+        output_slice.copy_from_slice(&self.initrd_data);
-        assert_eq!(read_bytes, unsafe { *buffer_size });
 
         Ok(())
     }
@@ -137,11 +118,8 @@ pub struct InitrdLoader {
 /// Returns the data range of the initrd in the PE binary.
 ///
 /// The initrd has to be embedded in the file as a .initrd PE section.
-fn initrd_location(initrd_efi: &mut RegularFile) -> Result<Range<usize>> {
+fn initrd_location(initrd_efi: &[u8]) -> Result<Range<usize>> {
-    initrd_efi.set_position(0)?;
+    let pe_binary = goblin::pe::PE::parse(initrd_efi).map_err(|_| Status::INVALID_PARAMETER)?;
-
-    let file_data = read_all(initrd_efi)?;
-    let pe_binary = goblin::pe::PE::parse(&file_data).map_err(|_| Status::INVALID_PARAMETER)?;
 
     pe_binary
         .sections
@@ -163,14 +141,11 @@ fn initrd_location(initrd_efi: &mut RegularFile) -> Result<Range<usize>> {
 ///
 /// For this to work, the initrd needs to be a PE binary. We misuse
 /// [`BootServices::load_image`] for this.
-fn initrd_verify(boot_services: &BootServices, initrd_efi: &mut RegularFile) -> Result<()> {
+fn initrd_verify(boot_services: &BootServices, initrd_efi: &[u8]) -> Result<()> {
-    initrd_efi.set_position(0)?;
-    let file_data = read_all(initrd_efi)?;
-
     let initrd_handle = boot_services.load_image(
         boot_services.image_handle(),
         LoadImageSource::FromBuffer {
-            buffer: &file_data,
+            buffer: &initrd_efi,
             file_path: None,
         },
     )?;
@@ -192,15 +167,21 @@ impl InitrdLoader {
     pub fn new(
         boot_services: &BootServices,
         handle: Handle,
-        mut file: RegularFile,
+        mut initrd_data: Vec<u8>,
     ) -> Result<Self> {
-        initrd_verify(boot_services, &mut file)?;
+        initrd_verify(boot_services, &initrd_data)?;
+
+        let range = initrd_location(&initrd_data)?;
+
+        // Remove the PE wrapper from the initrd. We do this in place
+        // to avoid having to keep the initrd in memory twice.
+        initrd_data.drain(0..range.start);
+        initrd_data.resize(range.end - range.start, 0);
+        initrd_data.shrink_to_fit();
 
-        let range = initrd_location(&mut file)?;
         let mut proto = Box::pin(LoadFile2Protocol {
             load_file: raw_load_file,
-            file,
+            initrd_data,
-            range,
         });
 
         // Linux finds the right handle by looking for something that

rust/lanzaboote/src/main.rs

@@ -84,8 +84,8 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
         EmbeddedConfiguration::new(&mut booted_image_file(system_table.boot_services()).unwrap())
             .expect("Failed to extract configuration from binary. Did you run lanzatool?");
 
-    let mut kernel_file;
+    let kernel_data;
-    let initrd_file;
+    let initrd_data;
 
     {
         let mut file_system = system_table
@@ -96,7 +96,7 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
             .open_volume()
             .expect("Failed to find ESP root directory");
 
-        kernel_file = root
+        let mut kernel_file = root
             .open(
                 &config.kernel_filename,
                 FileMode::Read,
@@ -106,7 +106,9 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
             .into_regular_file()
             .expect("Kernel is not a regular file");
 
-        initrd_file = root
+        kernel_data = read_all(&mut kernel_file).expect("Failed to read kernel file into memory");
+
+        let mut initrd_file = root
             .open(
                 &config.initrd_filename,
                 FileMode::Read,
@@ -115,15 +117,14 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
             .expect("Failed to open initrd for reading")
             .into_regular_file()
             .expect("Initrd is not a regular file");
+
+        initrd_data = read_all(&mut initrd_file).expect("Failed to read kernel file into memory");
     }
 
     let kernel_cmdline =
         booted_image_cmdline(system_table.boot_services()).expect("Failed to fetch command line");
 
     let kernel_handle = {
-        let kernel_data =
-            read_all(&mut kernel_file).expect("Failed to read kernel file into memory");
-
         system_table
             .boot_services()
             .load_image(
@@ -148,7 +149,7 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
         );
     }
 
-    let mut initrd_loader = InitrdLoader::new(system_table.boot_services(), handle, initrd_file)
+    let mut initrd_loader = InitrdLoader::new(system_table.boot_services(), handle, initrd_data)
         .expect("Failed to load the initrd. It may not be there or it is not signed");
     let status = system_table
         .boot_services()