diff --git a/CHANGELOG.md b/CHANGELOG.md index ffd058c89..7cb2c206d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,11 @@ - Added `LoadedImageDevicePath` protocol. - Added `FileAttribute::is_directory(&self)` and `FileAttribute::is_regular_file(&self)` +- Added `LoadedImage::code_type()` and `LoadedImage::data_type()` +- `Allocator` will now use the memory type of the running UEFI binary: + - `MemoryType::LOADER_DATA` for UEFI applications + - `MemoryType::BOOT_SERVICES_DATA` for UEFI boot drivers + - `MemoryType::RUNTIME_SERVICES_DATA` for UEFI runtime drivers ### Changed diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs index 207711b22..cf9153373 100644 --- a/uefi/src/allocator.rs +++ b/uefi/src/allocator.rs @@ -14,6 +14,7 @@ use core::alloc::{GlobalAlloc, Layout}; use core::ptr::{self, NonNull}; +use crate::proto::loaded_image::LoadedImage; use crate::table::boot::{BootServices, MemoryType}; /// Reference to the boot services table, used to call the pool memory allocation functions. @@ -22,6 +23,10 @@ use crate::table::boot::{BootServices, MemoryType}; /// exited by the host application yet. static mut BOOT_SERVICES: Option> = None; +/// The memory type used for pool memory allocations. +/// TODO: Use OnceCell when stablilized. +static mut MEMORY_TYPE: MemoryType = MemoryType::LOADER_DATA; + /// Initializes the allocator. /// /// # Safety @@ -30,6 +35,12 @@ static mut BOOT_SERVICES: Option> = None; /// will be called when UEFI boot services will be exited. pub unsafe fn init(boot_services: &BootServices) { BOOT_SERVICES = NonNull::new(boot_services as *const _ as *mut _); + + if let Ok(loaded_image) = + boot_services.open_protocol_exclusive::(boot_services.image_handle()) + { + MEMORY_TYPE = loaded_image.data_type() + } } /// Access the boot services @@ -54,9 +65,9 @@ pub struct Allocator; unsafe impl GlobalAlloc for Allocator { /// Allocate memory using [`BootServices::allocate_pool`]. The allocation is - /// of type [`MemoryType::LOADER_DATA`]. + /// of type [`MemoryType::LOADER_DATA`] for UEFI applications, [`MemoryType::BOOT_SERVICES_DATA`] + /// for UEFI boot drivers and [`MemoryType::RUNTIME_SERVICES_DATA`] for UEFI runtime drivers. unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - let mem_ty = MemoryType::LOADER_DATA; let size = layout.size(); let align = layout.align(); @@ -65,12 +76,14 @@ unsafe impl GlobalAlloc for Allocator { // only guaranteed to provide eight-byte alignment. Allocate extra // space so that we can return an appropriately-aligned pointer // within the allocation. - let full_alloc_ptr = - if let Ok(ptr) = boot_services().as_ref().allocate_pool(mem_ty, size + align) { - ptr - } else { - return ptr::null_mut(); - }; + let full_alloc_ptr = if let Ok(ptr) = boot_services() + .as_ref() + .allocate_pool(MEMORY_TYPE, size + align) + { + ptr + } else { + return ptr::null_mut(); + }; // Calculate the offset needed to get an aligned pointer within the // full allocation. If that offset is zero, increase it to `align` @@ -97,7 +110,7 @@ unsafe impl GlobalAlloc for Allocator { // use `allocate_pool` directly. boot_services() .as_ref() - .allocate_pool(mem_ty, size) + .allocate_pool(MEMORY_TYPE, size) .unwrap_or(ptr::null_mut()) } } diff --git a/uefi/src/proto/loaded_image.rs b/uefi/src/proto/loaded_image.rs index 279902dcc..29ca2b42c 100644 --- a/uefi/src/proto/loaded_image.rs +++ b/uefi/src/proto/loaded_image.rs @@ -170,4 +170,26 @@ impl LoadedImage { pub const fn info(&self) -> (*const c_void, u64) { (self.image_base, self.image_size) } + + /// Get the memory type of the image's code sections. + /// + /// Normally the returned value is one of: + /// - `MemoryType::LOADER_CODE` for UEFI applications + /// - `MemoryType::BOOT_SERVICES_CODE` for UEFI boot drivers + /// - `MemoryType::RUNTIME_SERVICES_CODE` for UEFI runtime drivers + #[must_use] + pub fn code_type(&self) -> MemoryType { + self.image_code_type + } + + /// Get the memory type of the image's data sections. + /// + /// Normally the returned value is one of: + /// - `MemoryType::LOADER_DATA` for UEFI applications + /// - `MemoryType::BOOT_SERVICES_DATA` for UEFI boot drivers + /// - `MemoryType::RUNTIME_SERVICES_DATA` for UEFI runtime drivers + #[must_use] + pub fn data_type(&self) -> MemoryType { + self.image_data_type + } }