From e0c9d78b2b21bd5d03379f3c84a2ef031e7a75cb Mon Sep 17 00:00:00 2001 From: Philipp Schuster Date: Tue, 15 Apr 2025 13:21:16 +0200 Subject: [PATCH 1/5] uefi: allocator: use match{} arms to separate cases This prepares changes in upcoming commits. --- uefi/src/allocator.rs | 104 +++++++++++++++++++++++------------------- 1 file changed, 57 insertions(+), 47 deletions(-) diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs index d030d55ed..deb2b23bf 100644 --- a/uefi/src/allocator.rs +++ b/uefi/src/allocator.rs @@ -63,60 +63,70 @@ unsafe impl GlobalAlloc for Allocator { let align = layout.align(); let memory_type = get_memory_type(); - if align > 8 { - // The requested alignment is greater than 8, but `allocate_pool` is - // 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::allocate_pool(memory_type, size + align) { - ptr.as_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` - // so that we still have space to store the extra pointer described - // below. - let mut offset = full_alloc_ptr.align_offset(align); - if offset == 0 { - offset = align; + match align { + 0..=8 /* UEFI default alignment */ => { + // The requested alignment is less than or equal to eight, and + // `allocate_pool` always provides eight-byte alignment, so we can + // use `allocate_pool` directly. + boot::allocate_pool(memory_type, size) + .map(|ptr| ptr.as_ptr()) + .unwrap_or(ptr::null_mut()) } + 9.. => { + // The requested alignment is greater than 8, but `allocate_pool` is + // 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 = boot::allocate_pool(memory_type, size + align); + let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr + { + ptr.as_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` + // so that we still have space to store the extra pointer described + // below. + let mut offset = full_alloc_ptr.align_offset(align); + if offset == 0 { + offset = align; + } - // Before returning the aligned allocation, store a pointer to the - // full unaligned allocation in the bytes just before the aligned - // allocation. We know we have at least eight bytes there due to - // adding `align` to the memory allocation size. We also know the - // write is appropriately aligned for a `*mut u8` pointer because - // `align_ptr` is aligned, and alignments are always powers of two - // (as enforced by the `Layout` type). - unsafe { - let aligned_ptr = full_alloc_ptr.add(offset); - (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr); - aligned_ptr + // Before returning the aligned allocation, store a pointer to the + // full unaligned allocation in the bytes just before the aligned + // allocation. We know we have at least eight bytes there due to + // adding `align` to the memory allocation size. We also know the + // write is appropriately aligned for a `*mut u8` pointer because + // `align_ptr` is aligned, and alignments are always powers of two + // (as enforced by the `Layout` type). + unsafe { + let aligned_ptr = full_alloc_ptr.add(offset); + (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr); + aligned_ptr + } } - } else { - // The requested alignment is less than or equal to eight, and - // `allocate_pool` always provides eight-byte alignment, so we can - // use `allocate_pool` directly. - boot::allocate_pool(memory_type, size) - .map(|ptr| ptr.as_ptr()) - .unwrap_or(ptr::null_mut()) } } /// Deallocate memory using [`boot::free_pool`]. - unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) { - if layout.align() > 8 { - // Retrieve the pointer to the full allocation that was packed right - // before the aligned allocation in `alloc`. - ptr = unsafe { (ptr as *const *mut u8).sub(1).read() }; + /// + /// This will panic after exiting boot services. + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + match layout.align() { + 0..=8 => { + // OK to unwrap: `ptr` is required to be a valid allocation by the trait API. + let ptr = NonNull::new(ptr).unwrap(); + unsafe { boot::free_pool(ptr) }.unwrap(); + } + 9.. => { + // Retrieve the pointer to the full allocation that was packed right + // before the aligned allocation in `alloc`. + let ptr = unsafe { (ptr as *const *mut u8).sub(1).read() }; + let ptr = NonNull::new(ptr).unwrap(); + unsafe { boot::free_pool(ptr) }.unwrap(); + } } - - // OK to unwrap: `ptr` is required to be a valid allocation by the trait API. - let ptr = NonNull::new(ptr).unwrap(); - - // Warning: this will panic after exiting boot services. - unsafe { boot::free_pool(ptr) }.unwrap(); } } From 2a4b4d7bc6b866a69411394a1dd4f2bc5cd6356b Mon Sep 17 00:00:00 2001 From: Philipp Schuster Date: Tue, 15 Apr 2025 13:22:29 +0200 Subject: [PATCH 2/5] uefi: allocator: improve documentation --- uefi/src/allocator.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs index deb2b23bf..80ee3c5cd 100644 --- a/uefi/src/allocator.rs +++ b/uefi/src/allocator.rs @@ -1,13 +1,11 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -//! This module implements Rust's global allocator interface using UEFI's memory allocation functions. +//! This module exports [`Allocator`]. //! -//! If the `global_allocator` feature is enabled, the [`Allocator`] will be used -//! as the global Rust allocator. +//! The allocator can be used as global Rust allocator using the +//! `global_allocator` crate feature. See [`helpers`] for more info. //! -//! This allocator can only be used while boot services are active. If boot -//! services are not active, `alloc` will return a null pointer, and `dealloc` -//! will panic. +//! [`helpers`]: uefi::helpers use crate::boot; use crate::mem::memory_map::MemoryType; @@ -42,9 +40,13 @@ fn get_memory_type() -> MemoryType { } } -/// Allocator which uses the UEFI pool allocation functions. +/// Allocator using UEFI boot services. /// -/// Only valid for as long as the UEFI boot services are available. +/// This type implements [`GlobalAlloc`] and can be marked with the +/// `#[global_allocator]` attribute to be used as global Rust allocator. +/// +/// Note that if boot services are not active (anymore), [`Allocator::alloc`] +/// will return a null pointer and [`Allocator::dealloc`] will panic. #[derive(Debug)] pub struct Allocator; From a595057228bac0ab1a8b6b11d1995c6bdd0f43b3 Mon Sep 17 00:00:00 2001 From: Philipp Schuster Date: Sun, 6 Apr 2025 23:03:51 +0200 Subject: [PATCH 3/5] uefi: allocator: improve code readability; split code into function --- uefi/src/allocator.rs | 69 ++++++++++++++++++++++--------------------- 1 file changed, 36 insertions(+), 33 deletions(-) diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs index 80ee3c5cd..6f0157529 100644 --- a/uefi/src/allocator.rs +++ b/uefi/src/allocator.rs @@ -40,6 +40,41 @@ fn get_memory_type() -> MemoryType { } } +/// Helper to get a custom alignment out of an allocation with an alignment of +/// eight (UEFI default alignment). This works by allocating extra space and +/// storing a pointer to the actual allocation right above the allocation +/// handed out via the public API. +fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mut u8 { + let full_alloc_ptr = boot::allocate_pool(memory_type, size + align); + let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr { + ptr.as_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` + // so that we still have space to store the extra pointer described + // below. + let mut offset = full_alloc_ptr.align_offset(align); + if offset == 0 { + offset = align; + } + + // Before returning the aligned allocation, store a pointer to the + // full unaligned allocation in the bytes just before the aligned + // allocation. We know we have at least eight bytes there due to + // adding `align` to the memory allocation size. We also know the + // write is appropriately aligned for a `*mut u8` pointer because + // `align_ptr` is aligned, and alignments are always powers of two + // (as enforced by the `Layout` type). + unsafe { + let aligned_ptr = full_alloc_ptr.add(offset); + (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr); + aligned_ptr + } +} + /// Allocator using UEFI boot services. /// /// This type implements [`GlobalAlloc`] and can be marked with the @@ -75,39 +110,7 @@ unsafe impl GlobalAlloc for Allocator { .unwrap_or(ptr::null_mut()) } 9.. => { - // The requested alignment is greater than 8, but `allocate_pool` is - // 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 = boot::allocate_pool(memory_type, size + align); - let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr - { - ptr.as_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` - // so that we still have space to store the extra pointer described - // below. - let mut offset = full_alloc_ptr.align_offset(align); - if offset == 0 { - offset = align; - } - - // Before returning the aligned allocation, store a pointer to the - // full unaligned allocation in the bytes just before the aligned - // allocation. We know we have at least eight bytes there due to - // adding `align` to the memory allocation size. We also know the - // write is appropriately aligned for a `*mut u8` pointer because - // `align_ptr` is aligned, and alignments are always powers of two - // (as enforced by the `Layout` type). - unsafe { - let aligned_ptr = full_alloc_ptr.add(offset); - (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr); - aligned_ptr - } + alloc_pool_aligned(memory_type, size, align) } } } From bd3e22a0439cac6d07064b0c6108e0a089d1af63 Mon Sep 17 00:00:00 2001 From: Philipp Schuster Date: Mon, 7 Apr 2025 10:11:08 +0200 Subject: [PATCH 4/5] uefi-test-runner: restructure alloc test into modules This is a prerequisite for the next commit. --- uefi-test-runner/src/boot/memory.rs | 106 ++++++++++++++++------------ 1 file changed, 60 insertions(+), 46 deletions(-) diff --git a/uefi-test-runner/src/boot/memory.rs b/uefi-test-runner/src/boot/memory.rs index 9dfb9b23b..49a1a3c7c 100644 --- a/uefi-test-runner/src/boot/memory.rs +++ b/uefi-test-runner/src/boot/memory.rs @@ -1,74 +1,88 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use alloc::vec::Vec; -use uefi::boot::{self, AllocateType}; -use uefi::mem::memory_map::{MemoryMap, MemoryMapMut, MemoryType}; +use uefi::boot; +use uefi::mem::memory_map::{MemoryMap, MemoryMapMut}; +use uefi_raw::table::boot::MemoryType; pub fn test() { info!("Testing memory functions"); - test_allocate_pages(); - test_allocate_pool(); + bootservices::allocate_pages(); + bootservices::allocate_pool(); - vec_alloc(); - alloc_alignment(); + global::alloc_vec(); + global::alloc_alignment(); test_memory_map(); } -fn test_allocate_pages() { - let num_pages = 1; - let ptr = - boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages).unwrap(); - let addr = ptr.as_ptr() as usize; - assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned"); +/// Tests that directly use UEFI boot services to allocate memory. +mod bootservices { + use uefi::boot; + use uefi::boot::AllocateType; + use uefi_raw::table::boot::MemoryType; + + /// Tests the `allocate_pages` boot service. + pub fn allocate_pages() { + let num_pages = 1; + let ptr = boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages) + .unwrap(); + let addr = ptr.as_ptr() as usize; + assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned"); + + // Verify the page can be written to. + { + let ptr = ptr.as_ptr(); + unsafe { ptr.write_volatile(0xff) }; + unsafe { ptr.add(4095).write_volatile(0xff) }; + } - // Verify the page can be written to. - { - let ptr = ptr.as_ptr(); - unsafe { ptr.write_volatile(0xff) }; - unsafe { ptr.add(4095).write_volatile(0xff) }; + unsafe { boot::free_pages(ptr, num_pages) }.unwrap(); } - unsafe { boot::free_pages(ptr, num_pages) }.unwrap(); -} + /// Tests the `allocate_pool` boot service. + pub fn allocate_pool() { + let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap(); -fn test_allocate_pool() { - let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap(); - - // Verify the allocation can be written to. - { - let ptr = ptr.as_ptr(); - unsafe { ptr.write_volatile(0xff) }; - unsafe { ptr.add(9).write_volatile(0xff) }; + // Verify the allocation can be written to. + { + let ptr = ptr.as_ptr(); + unsafe { ptr.write_volatile(0xff) }; + unsafe { ptr.add(9).write_volatile(0xff) }; + } + unsafe { boot::free_pool(ptr) }.unwrap(); } - unsafe { boot::free_pool(ptr) }.unwrap(); } -// Simple test to ensure our custom allocator works with the `alloc` crate. -fn vec_alloc() { - info!("Allocating a vector through the `alloc` crate"); +/// Tests that use [`uefi::allocator::Allocator`], which is configured as the +/// global allocator. +mod global { + /// Simple test to ensure our custom allocator works with the `alloc` crate. + pub fn alloc_vec() { + info!("Allocating a vector through the `alloc` crate"); - #[allow(clippy::useless_vec)] - let mut values = vec![-5, 16, 23, 4, 0]; + #[allow(clippy::useless_vec)] + let mut values = vec![-5, 16, 23, 4, 0]; - values.sort_unstable(); + values.sort_unstable(); - assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector"); -} + assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector"); + } -// Simple test to ensure our custom allocator works with correct alignment. -fn alloc_alignment() { - info!("Allocating a structure with alignment to 0x100"); + /// Simple test to ensure our custom allocator works with correct alignment. + pub fn alloc_alignment() { + info!("Allocating a structure with alignment to 0x100"); - #[repr(align(0x100))] - struct Block( - // Ignore warning due to field not being read. - #[allow(dead_code)] [u8; 0x100], - ); + #[repr(align(0x100))] + struct Block( + // Ignore warning due to field not being read. + #[allow(dead_code)] [u8; 0x100], + ); - let value = vec![Block([1; 0x100])]; - assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment"); + let value = vec![Block([1; 0x100])]; + assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment"); + } } fn test_memory_map() { From 2a629c832d15c3e1f106b4fcbdfa101e4ab50692 Mon Sep 17 00:00:00 2001 From: Philipp Schuster Date: Mon, 7 Apr 2025 10:12:36 +0200 Subject: [PATCH 5/5] uefi: allocator: use shortcut for PAGE_SIZE + add test Allocating page-aligned memory via the global allocator is not uncommon for UEFI OS loaders. Therefore, it is feasible to use a shortcut in the allocator, and directly use boot::allocate_pages() rather than boot::allocate_pool(). We can look at the TRACE messages of `cargo xtask run` to verify that the shortcut is taken. --- uefi-test-runner/src/boot/memory.rs | 33 +++++++++---- uefi/CHANGELOG.md | 2 + uefi/src/allocator.rs | 72 +++++++++++++++++++++-------- 3 files changed, 78 insertions(+), 29 deletions(-) diff --git a/uefi-test-runner/src/boot/memory.rs b/uefi-test-runner/src/boot/memory.rs index 49a1a3c7c..c4cea3f40 100644 --- a/uefi-test-runner/src/boot/memory.rs +++ b/uefi-test-runner/src/boot/memory.rs @@ -58,9 +58,12 @@ mod bootservices { /// Tests that use [`uefi::allocator::Allocator`], which is configured as the /// global allocator. mod global { + use alloc::boxed::Box; + use uefi_raw::table::boot::PAGE_SIZE; + /// Simple test to ensure our custom allocator works with the `alloc` crate. pub fn alloc_vec() { - info!("Allocating a vector through the `alloc` crate"); + info!("Allocating a vector using the global allocator"); #[allow(clippy::useless_vec)] let mut values = vec![-5, 16, 23, 4, 0]; @@ -71,17 +74,27 @@ mod global { } /// Simple test to ensure our custom allocator works with correct alignment. + #[allow(dead_code)] // Ignore warning due to field not being read. pub fn alloc_alignment() { - info!("Allocating a structure with alignment to 0x100"); - - #[repr(align(0x100))] - struct Block( - // Ignore warning due to field not being read. - #[allow(dead_code)] [u8; 0x100], - ); + { + info!("Allocating a structure with alignment of 0x100 using the global allocator"); + #[repr(align(0x100))] + struct Block([u8; 0x100]); - let value = vec![Block([1; 0x100])]; - assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment"); + let value = vec![Block([1; 0x100])]; + assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment"); + } + { + info!("Allocating a memory page ({PAGE_SIZE}) using the global allocator"); + #[repr(align(4096))] + struct Page([u8; PAGE_SIZE]); + let value = Box::new(Page([0; PAGE_SIZE])); + assert_eq!( + value.0.as_ptr().align_offset(PAGE_SIZE), + 0, + "Wrong alignment" + ); + } } } diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md index b4edf2c69..5bdb16e4e 100644 --- a/uefi/CHANGELOG.md +++ b/uefi/CHANGELOG.md @@ -37,6 +37,8 @@ - The `Display` impl for `CStr8` now excludes the trailing null character. - `VariableKeys` initializes with a larger name buffer to work around firmware bugs on some devices. +- The UEFI `allocator::Allocator` has been optimized for page-aligned + allocations. # uefi - 0.34.1 (2025-02-07) diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs index 6f0157529..790f313a0 100644 --- a/uefi/src/allocator.rs +++ b/uefi/src/allocator.rs @@ -7,12 +7,13 @@ //! //! [`helpers`]: uefi::helpers -use crate::boot; +use crate::boot::{self, AllocateType}; use crate::mem::memory_map::MemoryType; use crate::proto::loaded_image::LoadedImage; use core::alloc::{GlobalAlloc, Layout}; use core::ptr::{self, NonNull}; use core::sync::atomic::{AtomicU32, Ordering}; +use uefi_raw::table::boot::PAGE_SIZE; /// Get the memory type to use for allocation. /// @@ -75,6 +76,20 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu } } +/// Returns whether the allocation is a multiple of a [`PAGE_SIZE`] and is +/// aligned to [`PAGE_SIZE`]. +/// +/// This does not only check the alignment but also the size. For types +/// allocated by Rust itself (e.g., `Box`), the size is always at least the +/// alignment, as specified in the [Rust type layout]. However, to be also safe +/// when it comes to manual invocations, we additionally check if the size is +/// a multiple of [`PAGE_SIZE`]. +/// +/// [Rust type layout]: https://doc.rust-lang.org/reference/type-layout.html +const fn layout_allows_page_alloc_shortcut(layout: &Layout) -> bool { + layout.size() % PAGE_SIZE == 0 && layout.align() == PAGE_SIZE +} + /// Allocator using UEFI boot services. /// /// This type implements [`GlobalAlloc`] and can be marked with the @@ -86,8 +101,9 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu pub struct Allocator; unsafe impl GlobalAlloc for Allocator { - /// Allocate memory using [`boot::allocate_pool`]. The allocation's [memory - /// type] matches the current image's [data type]. + /// Allocate memory using the UEFI boot services. + /// + /// The allocation's [memory type] matches the current image's [data type]. /// /// [memory type]: MemoryType /// [data type]: LoadedImage::data_type @@ -96,40 +112,58 @@ unsafe impl GlobalAlloc for Allocator { return ptr::null_mut(); } - let size = layout.size(); - let align = layout.align(); let memory_type = get_memory_type(); + let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout); - match align { - 0..=8 /* UEFI default alignment */ => { + match (use_page_shortcut, layout.align()) { + // Allocating pages is actually very expected in UEFI OS loaders, so + // it makes sense to provide this optimization. + (true, _) => { + // To spammy, but useful for manual testing. + // log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}"); + let count = layout.size().div_ceil(PAGE_SIZE); + boot::allocate_pages(AllocateType::AnyPages, memory_type, count) + .map(|ptr| ptr.as_ptr()) + .unwrap_or(ptr::null_mut()) + } + (false, 0..=8 /* UEFI default alignment */) => { // The requested alignment is less than or equal to eight, and // `allocate_pool` always provides eight-byte alignment, so we can // use `allocate_pool` directly. - boot::allocate_pool(memory_type, size) + boot::allocate_pool(memory_type, layout.size()) .map(|ptr| ptr.as_ptr()) .unwrap_or(ptr::null_mut()) } - 9.. => { - alloc_pool_aligned(memory_type, size, align) - } + (false, 9..) => alloc_pool_aligned(memory_type, layout.size(), layout.align()), } } - /// Deallocate memory using [`boot::free_pool`]. + /// Deallocate memory using the UEFI boot services. /// /// This will panic after exiting boot services. unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - match layout.align() { - 0..=8 => { - // OK to unwrap: `ptr` is required to be a valid allocation by the trait API. - let ptr = NonNull::new(ptr).unwrap(); + let ptr = NonNull::new(ptr).unwrap(); + + let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout); + + match (use_page_shortcut, layout.align()) { + (true, _) => { + // To spammy, but useful for manual testing. + // log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}"); + let count = layout.size().div_ceil(PAGE_SIZE); + unsafe { boot::free_pages(ptr, count).unwrap() } + } + (false, 0..=8 /* UEFI default alignment */) => { + // Warning: this will panic after exiting boot services. unsafe { boot::free_pool(ptr) }.unwrap(); } - 9.. => { + (false, 9..) => { + let ptr = ptr.as_ptr().cast::<*mut u8>(); // Retrieve the pointer to the full allocation that was packed right // before the aligned allocation in `alloc`. - let ptr = unsafe { (ptr as *const *mut u8).sub(1).read() }; - let ptr = NonNull::new(ptr).unwrap(); + let actual_alloc_ptr = unsafe { ptr.sub(1).read() }; + let ptr = NonNull::new(actual_alloc_ptr).unwrap(); + // Warning: this will panic after exiting boot services. unsafe { boot::free_pool(ptr) }.unwrap(); } }