From 441068b613b48c8f6cb251530ef1c56039720fd3 Mon Sep 17 00:00:00 2001 From: Ayush Singh Date: Wed, 25 Oct 2023 22:38:39 +0530 Subject: [PATCH] Use ImageDataType for allocation type Signed-off-by: Ayush Singh --- library/std/src/sys/uefi/alloc.rs | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/library/std/src/sys/uefi/alloc.rs b/library/std/src/sys/uefi/alloc.rs index 789e3cbd81ade..ad3904d82f3ff 100644 --- a/library/std/src/sys/uefi/alloc.rs +++ b/library/std/src/sys/uefi/alloc.rs @@ -1,13 +1,17 @@ //! Global Allocator for UEFI. //! Uses [r-efi-alloc](https://crates.io/crates/r-efi-alloc) -use crate::alloc::{GlobalAlloc, Layout, System}; +use r_efi::protocols::loaded_image; -const MEMORY_TYPE: u32 = r_efi::efi::LOADER_DATA; +use crate::alloc::{GlobalAlloc, Layout, System}; +use crate::sync::OnceLock; +use crate::sys::uefi::helpers; #[stable(feature = "alloc_system_type", since = "1.28.0")] unsafe impl GlobalAlloc for System { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + static EFI_MEMORY_TYPE: OnceLock = OnceLock::new(); + // Return null pointer if boot services are not available if crate::os::uefi::env::boot_services().is_none() { return crate::ptr::null_mut(); @@ -15,8 +19,20 @@ unsafe impl GlobalAlloc for System { // If boot services is valid then SystemTable is not null. let system_table = crate::os::uefi::env::system_table().as_ptr().cast(); + + // Each loaded image has an image handle that supports `EFI_LOADED_IMAGE_PROTOCOL`. Thus, this + // will never fail. + let mem_type = EFI_MEMORY_TYPE.get_or_init(|| { + let protocol = helpers::image_handle_protocol::( + loaded_image::PROTOCOL_GUID, + ) + .unwrap(); + // Gives allocations the memory type that the data sections were loaded as. + unsafe { (*protocol.as_ptr()).image_data_type } + }); + // The caller must ensure non-0 layout - unsafe { r_efi_alloc::raw::alloc(system_table, layout, MEMORY_TYPE) } + unsafe { r_efi_alloc::raw::alloc(system_table, layout, *mem_type) } } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {