diff --git a/rust/stub/src/main.rs b/rust/stub/src/main.rs index 0458444..5d5dc6f 100644 --- a/rust/stub/src/main.rs +++ b/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 { - let string = pe_section_as_string(file_data, section).ok_or(Status::INVALID_PARAMETER)?; +fn extract_string(pe_data: &[u8], section: &str) -> Result { + 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 { - let array: [u8; 32] = pe_section(file_data, section) +fn extract_hash(pe_data: &[u8], section: &str) -> Result { + 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 { } impl EmbeddedConfiguration { - fn new(file: &mut RegularFile) -> Result { - file.set_position(0)?; - let file_data = read_all(file)?; - + fn new(file_data: &[u8]) -> Result { 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) -> 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; diff --git a/rust/stub/src/pe_section.rs b/rust/stub/src/pe_section.rs index 0ef7cce..ba9541b 100644 --- a/rust/stub/src/pe_section.rs +++ b/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 { - 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 { + pe_section(pe_data, section_name).map(|data| core::str::from_utf8(data).unwrap().to_owned()) } diff --git a/rust/stub/src/uefi_helpers.rs b/rust/stub/src/uefi_helpers.rs index 7bb374b..1b430cc 100644 --- a/rust/stub/src/uefi_helpers.rs +++ b/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> { 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 { - 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 { let loaded_image = boot_services.open_protocol_exclusive::(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::( - boot_services.get_handle_for_protocol::()?, - )?; - - 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)?, + }) }