stub: do not read loaded image again from ESP

... because this might not work, if we were not loaded from a file
system. It also removes the issue where we might not load the signed
image that was actually loaded.

Fixes #123

rust/stub/src/main.rs

@@ -20,7 +20,7 @@ use uefi::{
     prelude::*,
     proto::{
         loaded_image::LoadedImage,
-        media::file::{File, FileAttribute, FileMode, RegularFile},
+        media::file::{File, FileAttribute, FileMode},
     },
     CStr16, CString16, Result,
 };
@@ -73,15 +73,15 @@ struct EmbeddedConfiguration {
 }
 
 /// Extract a string, stored as UTF-8, from a PE section.
-fn extract_string(file_data: &[u8], section: &str) -> Result<CString16> {
-    let string = pe_section_as_string(file_data, section).ok_or(Status::INVALID_PARAMETER)?;
+fn extract_string(pe_data: &[u8], section: &str) -> Result<CString16> {
+    let string = pe_section_as_string(pe_data, section).ok_or(Status::INVALID_PARAMETER)?;
 
     Ok(CString16::try_from(string.as_str()).map_err(|_| Status::INVALID_PARAMETER)?)
 }
 
 /// Extract a Blake3 hash from a PE section.
-fn extract_hash(file_data: &[u8], section: &str) -> Result<Hash> {
-    let array: [u8; 32] = pe_section(file_data, section)
+fn extract_hash(pe_data: &[u8], section: &str) -> Result<Hash> {
+    let array: [u8; 32] = pe_section(pe_data, section)
         .ok_or(Status::INVALID_PARAMETER)?
         .try_into()
         .map_err(|_| Status::INVALID_PARAMETER)?;
@@ -90,18 +90,15 @@ fn extract_hash(file_data: &[u8], section: &str) -> Result<Hash> {
 }
 
 impl EmbeddedConfiguration {
-    fn new(file: &mut RegularFile) -> Result<Self> {
-        file.set_position(0)?;
-        let file_data = read_all(file)?;
-
+    fn new(file_data: &[u8]) -> Result<Self> {
         Ok(Self {
-            kernel_filename: extract_string(&file_data, ".kernelp")?,
-            kernel_hash: extract_hash(&file_data, ".kernelh")?,
+            kernel_filename: extract_string(file_data, ".kernelp")?,
+            kernel_hash: extract_hash(file_data, ".kernelh")?,
 
-            initrd_filename: extract_string(&file_data, ".initrdp")?,
-            initrd_hash: extract_hash(&file_data, ".initrdh")?,
+            initrd_filename: extract_string(file_data, ".initrdp")?,
+            initrd_hash: extract_hash(file_data, ".initrdh")?,
 
-            cmdline: extract_string(&file_data, ".cmdline")?,
+            cmdline: extract_string(file_data, ".cmdline")?,
         })
     }
 }
@@ -180,9 +177,14 @@ fn main(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
 
     print_logo();
 
-    let config: EmbeddedConfiguration =
-        EmbeddedConfiguration::new(&mut booted_image_file(system_table.boot_services()).unwrap())
-            .expect("Failed to extract configuration from binary. Did you run lzbt?");
+    let config: EmbeddedConfiguration = unsafe {
+        EmbeddedConfiguration::new(
+            booted_image_file(system_table.boot_services())
+                .unwrap()
+                .as_slice(),
+        )
+        .expect("Failed to extract configuration from binary. Did you run lzbt?")
+    };
 
     let kernel_data;
     let initrd_data;

rust/stub/src/pe_section.rs

@@ -6,25 +6,25 @@
 
 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()?;
+/// Extracts the data of a section of a loaded PE file.
+pub fn pe_section<'a>(pe_data: &'a [u8], section_name: &str) -> Option<&'a [u8]> {
+    let pe_binary = goblin::pe::PE::parse(pe_data).ok()?;
 
     pe_binary
         .sections
         .iter()
-        .find(|s| s.name().unwrap() == section_name)
+        .find(|s| s.name().map(|n| n == section_name).unwrap_or(false))
         .and_then(|s| {
-            let section_start: usize = s.pointer_to_raw_data.try_into().ok()?;
+            let section_start: usize = s.virtual_address.try_into().ok()?;
 
             assert!(s.virtual_size <= s.size_of_raw_data);
             let section_end: usize = section_start + usize::try_from(s.virtual_size).ok()?;
 
-            Some(&file_data[section_start..section_end])
+            Some(&pe_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).map(|data| core::str::from_utf8(data).unwrap().to_owned())
+/// Extracts the data of a section of a loaded PE image and returns it as a string.
+pub fn pe_section_as_string<'a>(pe_data: &'a [u8], section_name: &str) -> Option<String> {
+    pe_section(pe_data, section_name).map(|data| core::str::from_utf8(data).unwrap().to_owned())
 }

rust/stub/src/uefi_helpers.rs

@@ -1,12 +1,10 @@
+use core::ffi::c_void;
+
 use alloc::vec::Vec;
 use uefi::{
     prelude::BootServices,
-    proto::{
-        device_path::text::{AllowShortcuts, DevicePathToText, DisplayOnly},
-        loaded_image::LoadedImage,
-        media::file::{File, FileAttribute, FileMode, RegularFile},
-    },
-    Result, Status,
+    proto::{loaded_image::LoadedImage, media::file::RegularFile},
+    Result,
 };
 
 /// Read the whole file into a vector.
@@ -27,32 +25,36 @@ pub fn read_all(file: &mut RegularFile) -> Result<Vec<u8>> {
     Ok(buf)
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct PeInMemory {
+    image_base: *const c_void,
+    image_size: usize,
+}
+
+impl PeInMemory {
+    /// Return a reference to the currently running image.
+    ///
+    /// # Safety
+    ///
+    /// The returned slice covers the whole loaded image in which we
+    /// currently execute. This means the safety guarantees of
+    /// [`core::slice::from_raw_parts`] that we use in this function
+    /// are only guaranteed, if we we don't mutate anything in this
+    /// range. This means no modification of global variables or
+    /// anything.
+    pub unsafe fn as_slice(&self) -> &'static [u8] {
+        unsafe { core::slice::from_raw_parts(self.image_base as *const u8, self.image_size) }
+    }
+}
+
 /// Open the currently executing image as a file.
-pub fn booted_image_file(boot_services: &BootServices) -> Result<RegularFile> {
-    let mut file_system = boot_services.get_image_file_system(boot_services.image_handle())?;
-
-    // The root directory of the volume where our binary lies.
-    let mut root = file_system.open_volume()?;
-
+pub fn booted_image_file(boot_services: &BootServices) -> Result<PeInMemory> {
     let loaded_image =
         boot_services.open_protocol_exclusive::<LoadedImage>(boot_services.image_handle())?;
+    let (image_base, image_size) = loaded_image.info();
 
-    let file_path = loaded_image.file_path().ok_or(Status::NOT_FOUND)?;
-
-    let to_text = boot_services.open_protocol_exclusive::<DevicePathToText>(
-        boot_services.get_handle_for_protocol::<DevicePathToText>()?,
-    )?;
-
-    let file_path = to_text.convert_device_path_to_text(
-        boot_services,
-        file_path,
-        DisplayOnly(false),
-        AllowShortcuts(false),
-    )?;
-
-    let our_image = root.open(&file_path, FileMode::Read, FileAttribute::empty())?;
-
-    Ok(our_image
-        .into_regular_file()
-        .ok_or(Status::INVALID_PARAMETER)?)
+    Ok(PeInMemory {
+        image_base,
+        image_size: usize::try_from(image_size).map_err(|_| uefi::Status::INVALID_PARAMETER)?,
+    })
 }