Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit f1141b9

Browse files
Merge pull request #1611 from rust-osdev/globalalloc
UEFI Allocator: add PAGE_SIZE shortcut
2 parents 8ee0b6c + 2a629c8 commit f1141b9

File tree

3 files changed

+185
-107
lines changed

3 files changed

+185
-107
lines changed

‎uefi-test-runner/src/boot/memory.rs

Lines changed: 73 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,74 +1,101 @@
11
// SPDX-License-Identifier: MIT OR Apache-2.0
22

33
use alloc::vec::Vec;
4-
use uefi::boot::{self, AllocateType};
5-
use uefi::mem::memory_map::{MemoryMap, MemoryMapMut, MemoryType};
4+
use uefi::boot;
5+
use uefi::mem::memory_map::{MemoryMap, MemoryMapMut};
6+
use uefi_raw::table::boot::MemoryType;
67

78
pub fn test() {
89
info!("Testing memory functions");
910

10-
test_allocate_pages();
11-
test_allocate_pool();
11+
bootservices::allocate_pages();
12+
bootservices::allocate_pool();
1213

13-
vec_alloc();
14-
alloc_alignment();
14+
global::alloc_vec();
15+
global::alloc_alignment();
1516

1617
test_memory_map();
1718
}
1819

19-
fn test_allocate_pages() {
20-
let num_pages = 1;
21-
let ptr =
22-
boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages).unwrap();
23-
let addr = ptr.as_ptr() as usize;
24-
assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned");
20+
/// Tests that directly use UEFI boot services to allocate memory.
21+
mod bootservices {
22+
use uefi::boot;
23+
use uefi::boot::AllocateType;
24+
use uefi_raw::table::boot::MemoryType;
25+
26+
/// Tests the `allocate_pages` boot service.
27+
pub fn allocate_pages() {
28+
let num_pages = 1;
29+
let ptr = boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages)
30+
.unwrap();
31+
let addr = ptr.as_ptr() as usize;
32+
assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned");
33+
34+
// Verify the page can be written to.
35+
{
36+
let ptr = ptr.as_ptr();
37+
unsafe { ptr.write_volatile(0xff) };
38+
unsafe { ptr.add(4095).write_volatile(0xff) };
39+
}
2540

26-
// Verify the page can be written to.
27-
{
28-
let ptr = ptr.as_ptr();
29-
unsafe { ptr.write_volatile(0xff) };
30-
unsafe { ptr.add(4095).write_volatile(0xff) };
41+
unsafe { boot::free_pages(ptr, num_pages) }.unwrap();
3142
}
3243

33-
unsafe { boot::free_pages(ptr, num_pages) }.unwrap();
34-
}
35-
36-
fn test_allocate_pool() {
37-
let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
44+
/// Tests the `allocate_pool` boot service.
45+
pub fn allocate_pool() {
46+
let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
3847

39-
// Verify the allocation can be written to.
40-
{
41-
let ptr = ptr.as_ptr();
42-
unsafe { ptr.write_volatile(0xff) };
43-
unsafe { ptr.add(9).write_volatile(0xff) };
48+
// Verify the allocation can be written to.
49+
{
50+
let ptr = ptr.as_ptr();
51+
unsafe { ptr.write_volatile(0xff) };
52+
unsafe { ptr.add(9).write_volatile(0xff) };
53+
}
54+
unsafe { boot::free_pool(ptr) }.unwrap();
4455
}
45-
unsafe { boot::free_pool(ptr) }.unwrap();
4656
}
4757

48-
// Simple test to ensure our custom allocator works with the `alloc` crate.
49-
fn vec_alloc() {
50-
info!("Allocating a vector through the `alloc` crate");
58+
/// Tests that use [`uefi::allocator::Allocator`], which is configured as the
59+
/// global allocator.
60+
mod global {
61+
use alloc::boxed::Box;
62+
use uefi_raw::table::boot::PAGE_SIZE;
5163

52-
#[allow(clippy::useless_vec)]
53-
let mut values = vec![-5, 16, 23, 4, 0];
64+
/// Simple test to ensure our custom allocator works with the `alloc` crate.
65+
pub fn alloc_vec() {
66+
info!("Allocating a vector using the global allocator");
5467

55-
values.sort_unstable();
68+
#[allow(clippy::useless_vec)]
69+
let mut values = vec![-5, 16, 23, 4, 0];
5670

57-
assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector");
58-
}
71+
values.sort_unstable();
5972

60-
// Simple test to ensure our custom allocator works with correct alignment.
61-
fn alloc_alignment() {
62-
info!("Allocating a structure with alignment to 0x100");
73+
assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector");
74+
}
6375

64-
#[repr(align(0x100))]
65-
struct Block(
66-
// Ignore warning due to field not being read.
67-
#[allow(dead_code)] [u8; 0x100],
68-
);
76+
/// Simple test to ensure our custom allocator works with correct alignment.
77+
#[allow(dead_code)] // Ignore warning due to field not being read.
78+
pub fn alloc_alignment() {
79+
{
80+
info!("Allocating a structure with alignment of 0x100 using the global allocator");
81+
#[repr(align(0x100))]
82+
struct Block([u8; 0x100]);
6983

70-
let value = vec![Block([1; 0x100])];
71-
assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
84+
let value = vec![Block([1; 0x100])];
85+
assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
86+
}
87+
{
88+
info!("Allocating a memory page ({PAGE_SIZE}) using the global allocator");
89+
#[repr(align(4096))]
90+
struct Page([u8; PAGE_SIZE]);
91+
let value = Box::new(Page([0; PAGE_SIZE]));
92+
assert_eq!(
93+
value.0.as_ptr().align_offset(PAGE_SIZE),
94+
0,
95+
"Wrong alignment"
96+
);
97+
}
98+
}
7299
}
73100

74101
fn test_memory_map() {

‎uefi/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
- The `Display` impl for `CStr8` now excludes the trailing null character.
4040
- `VariableKeys` initializes with a larger name buffer to work around firmware
4141
bugs on some devices.
42+
- The UEFI `allocator::Allocator` has been optimized for page-aligned
43+
allocations.
4244

4345

4446
# uefi - 0.34.1 (2025年02月07日)

‎uefi/src/allocator.rs

Lines changed: 110 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,19 @@
11
// SPDX-License-Identifier: MIT OR Apache-2.0
22

3-
//! This module implements Rust's global allocator interface using UEFI's memory allocation functions.
3+
//! This module exports [`Allocator`].
44
//!
5-
//! If the `global_allocator` feature is enabled, the [`Allocator`] will be used
6-
//! as the global Rust allocator.
5+
//! The allocator can be used as global Rust allocator using the
6+
//! `global_allocator` crate feature. See [`helpers`] for more info.
77
//!
8-
//! This allocator can only be used while boot services are active. If boot
9-
//! services are not active, `alloc` will return a null pointer, and `dealloc`
10-
//! will panic.
8+
//! [`helpers`]: uefi::helpers
119
12-
use crate::boot;
10+
use crate::boot::{self,AllocateType};
1311
use crate::mem::memory_map::MemoryType;
1412
use crate::proto::loaded_image::LoadedImage;
1513
use core::alloc::{GlobalAlloc, Layout};
1614
use core::ptr::{self, NonNull};
1715
use core::sync::atomic::{AtomicU32, Ordering};
16+
use uefi_raw::table::boot::PAGE_SIZE;
1817

1918
/// Get the memory type to use for allocation.
2019
///
@@ -42,15 +41,69 @@ fn get_memory_type() -> MemoryType {
4241
}
4342
}
4443

45-
/// Allocator which uses the UEFI pool allocation functions.
44+
/// Helper to get a custom alignment out of an allocation with an alignment of
45+
/// eight (UEFI default alignment). This works by allocating extra space and
46+
/// storing a pointer to the actual allocation right above the allocation
47+
/// handed out via the public API.
48+
fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mut u8 {
49+
let full_alloc_ptr = boot::allocate_pool(memory_type, size + align);
50+
let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr {
51+
ptr.as_ptr()
52+
} else {
53+
return ptr::null_mut();
54+
};
55+
56+
// Calculate the offset needed to get an aligned pointer within the
57+
// full allocation. If that offset is zero, increase it to `align`
58+
// so that we still have space to store the extra pointer described
59+
// below.
60+
let mut offset = full_alloc_ptr.align_offset(align);
61+
if offset == 0 {
62+
offset = align;
63+
}
64+
65+
// Before returning the aligned allocation, store a pointer to the
66+
// full unaligned allocation in the bytes just before the aligned
67+
// allocation. We know we have at least eight bytes there due to
68+
// adding `align` to the memory allocation size. We also know the
69+
// write is appropriately aligned for a `*mut u8` pointer because
70+
// `align_ptr` is aligned, and alignments are always powers of two
71+
// (as enforced by the `Layout` type).
72+
unsafe {
73+
let aligned_ptr = full_alloc_ptr.add(offset);
74+
(aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
75+
aligned_ptr
76+
}
77+
}
78+
79+
/// Returns whether the allocation is a multiple of a [`PAGE_SIZE`] and is
80+
/// aligned to [`PAGE_SIZE`].
81+
///
82+
/// This does not only check the alignment but also the size. For types
83+
/// allocated by Rust itself (e.g., `Box<T>`), the size is always at least the
84+
/// alignment, as specified in the [Rust type layout]. However, to be also safe
85+
/// when it comes to manual invocations, we additionally check if the size is
86+
/// a multiple of [`PAGE_SIZE`].
87+
///
88+
/// [Rust type layout]: https://doc.rust-lang.org/reference/type-layout.html
89+
const fn layout_allows_page_alloc_shortcut(layout: &Layout) -> bool {
90+
layout.size() % PAGE_SIZE == 0 && layout.align() == PAGE_SIZE
91+
}
92+
93+
/// Allocator using UEFI boot services.
94+
///
95+
/// This type implements [`GlobalAlloc`] and can be marked with the
96+
/// `#[global_allocator]` attribute to be used as global Rust allocator.
4697
///
47-
/// Only valid for as long as the UEFI boot services are available.
98+
/// Note that if boot services are not active (anymore), [`Allocator::alloc`]
99+
/// will return a null pointer and [`Allocator::dealloc`] will panic.
48100
#[derive(Debug)]
49101
pub struct Allocator;
50102

51103
unsafe impl GlobalAlloc for Allocator {
52-
/// Allocate memory using [`boot::allocate_pool`]. The allocation's [memory
53-
/// type] matches the current image's [data type].
104+
/// Allocate memory using the UEFI boot services.
105+
///
106+
/// The allocation's [memory type] matches the current image's [data type].
54107
///
55108
/// [memory type]: MemoryType
56109
/// [data type]: LoadedImage::data_type
@@ -59,64 +112,60 @@ unsafe impl GlobalAlloc for Allocator {
59112
return ptr::null_mut();
60113
}
61114

62-
let size = layout.size();
63-
let align = layout.align();
64115
let memory_type = get_memory_type();
116+
let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
65117

66-
if align > 8 {
67-
// The requested alignment is greater than 8, but `allocate_pool` is
68-
// only guaranteed to provide eight-byte alignment. Allocate extra
69-
// space so that we can return an appropriately-aligned pointer
70-
// within the allocation.
71-
let full_alloc_ptr = if let Ok(ptr) = boot::allocate_pool(memory_type, size + align) {
72-
ptr.as_ptr()
73-
} else {
74-
return ptr::null_mut();
75-
};
76-
77-
// Calculate the offset needed to get an aligned pointer within the
78-
// full allocation. If that offset is zero, increase it to `align`
79-
// so that we still have space to store the extra pointer described
80-
// below.
81-
let mut offset = full_alloc_ptr.align_offset(align);
82-
if offset == 0 {
83-
offset = align;
118+
match (use_page_shortcut, layout.align()) {
119+
// Allocating pages is actually very expected in UEFI OS loaders, so
120+
// it makes sense to provide this optimization.
121+
(true, _) => {
122+
// To spammy, but useful for manual testing.
123+
// log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
124+
let count = layout.size().div_ceil(PAGE_SIZE);
125+
boot::allocate_pages(AllocateType::AnyPages, memory_type, count)
126+
.map(|ptr| ptr.as_ptr())
127+
.unwrap_or(ptr::null_mut())
84128
}
85-
86-
// Before returning the aligned allocation, store a pointer to the
87-
// full unaligned allocation in the bytes just before the aligned
88-
// allocation. We know we have at least eight bytes there due to
89-
// adding `align` to the memory allocation size. We also know the
90-
// write is appropriately aligned for a `*mut u8` pointer because
91-
// `align_ptr` is aligned, and alignments are always powers of two
92-
// (as enforced by the `Layout` type).
93-
unsafe {
94-
let aligned_ptr = full_alloc_ptr.add(offset);
95-
(aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
96-
aligned_ptr
129+
(false, 0..=8 /* UEFI default alignment */) => {
130+
// The requested alignment is less than or equal to eight, and
131+
// `allocate_pool` always provides eight-byte alignment, so we can
132+
// use `allocate_pool` directly.
133+
boot::allocate_pool(memory_type, layout.size())
134+
.map(|ptr| ptr.as_ptr())
135+
.unwrap_or(ptr::null_mut())
97136
}
98-
} else {
99-
// The requested alignment is less than or equal to eight, and
100-
// `allocate_pool` always provides eight-byte alignment, so we can
101-
// use `allocate_pool` directly.
102-
boot::allocate_pool(memory_type, size)
103-
.map(|ptr| ptr.as_ptr())
104-
.unwrap_or(ptr::null_mut())
137+
(false, 9..) => alloc_pool_aligned(memory_type, layout.size(), layout.align()),
105138
}
106139
}
107140

108-
/// Deallocate memory using [`boot::free_pool`].
109-
unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
110-
if layout.align() > 8 {
111-
// Retrieve the pointer to the full allocation that was packed right
112-
// before the aligned allocation in `alloc`.
113-
ptr = unsafe { (ptr as *const *mut u8).sub(1).read() };
114-
}
115-
116-
// OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
141+
/// Deallocate memory using the UEFI boot services.
142+
///
143+
/// This will panic after exiting boot services.
144+
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
117145
let ptr = NonNull::new(ptr).unwrap();
118146

119-
// Warning: this will panic after exiting boot services.
120-
unsafe { boot::free_pool(ptr) }.unwrap();
147+
let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
148+
149+
match (use_page_shortcut, layout.align()) {
150+
(true, _) => {
151+
// To spammy, but useful for manual testing.
152+
// log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
153+
let count = layout.size().div_ceil(PAGE_SIZE);
154+
unsafe { boot::free_pages(ptr, count).unwrap() }
155+
}
156+
(false, 0..=8 /* UEFI default alignment */) => {
157+
// Warning: this will panic after exiting boot services.
158+
unsafe { boot::free_pool(ptr) }.unwrap();
159+
}
160+
(false, 9..) => {
161+
let ptr = ptr.as_ptr().cast::<*mut u8>();
162+
// Retrieve the pointer to the full allocation that was packed right
163+
// before the aligned allocation in `alloc`.
164+
let actual_alloc_ptr = unsafe { ptr.sub(1).read() };
165+
let ptr = NonNull::new(actual_alloc_ptr).unwrap();
166+
// Warning: this will panic after exiting boot services.
167+
unsafe { boot::free_pool(ptr) }.unwrap();
168+
}
169+
}
121170
}
122171
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /