1 //! This module implements Rust's global allocator interface using UEFI's memory allocation functions.
2 //!
3 //! If the `global_allocator` feature is enabled, the [`Allocator`] will be used
4 //! as the global Rust allocator.
5 //!
6 //! This allocator can only be used while boot services are active. If boot
7 //! services are not active, `alloc` will return a null pointer, and `dealloc`
8 //! will panic.
9 
10 use core::alloc::{GlobalAlloc, Layout};
11 use core::ptr::{self, NonNull};
12 use core::sync::atomic::{AtomicU32, Ordering};
13 
14 use crate::boot;
15 use crate::mem::memory_map::MemoryType;
16 use crate::proto::loaded_image::LoadedImage;
17 
18 /// Get the memory type to use for allocation.
19 ///
20 /// The first time this is called, the data type of the loaded image will be
21 /// retrieved. That value is cached in a static and reused on subsequent
22 /// calls. If the memory type of the loaded image cannot be retrieved for some
23 /// reason, a default of `LOADER_DATA` is used.
get_memory_type() -> MemoryType24 fn get_memory_type() -> MemoryType {
25     // Initialize to a `RESERVED` to indicate the actual value hasn't been set yet.
26     static MEMORY_TYPE: AtomicU32 = AtomicU32::new(MemoryType::RESERVED.0);
27 
28     let memory_type = MEMORY_TYPE.load(Ordering::Acquire);
29     if memory_type == MemoryType::RESERVED.0 {
30         let memory_type = if let Ok(loaded_image) =
31             boot::open_protocol_exclusive::<LoadedImage>(boot::image_handle())
32         {
33             loaded_image.data_type()
34         } else {
35             MemoryType::LOADER_DATA
36         };
37         MEMORY_TYPE.store(memory_type.0, Ordering::Release);
38         memory_type
39     } else {
40         MemoryType(memory_type)
41     }
42 }
43 
44 /// Allocator which uses the UEFI pool allocation functions.
45 ///
46 /// Only valid for as long as the UEFI boot services are available.
47 #[derive(Debug)]
48 pub struct Allocator;
49 
50 unsafe impl GlobalAlloc for Allocator {
51     /// Allocate memory using [`boot::allocate_pool`]. The allocation is
52     /// of type [`MemoryType::LOADER_DATA`] for UEFI applications, [`MemoryType::BOOT_SERVICES_DATA`]
53     /// for UEFI boot drivers and [`MemoryType::RUNTIME_SERVICES_DATA`] for UEFI runtime drivers.
alloc(&self, layout: Layout) -> *mut u854     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
55         if !boot::are_boot_services_active() {
56             return ptr::null_mut();
57         }
58 
59         let size = layout.size();
60         let align = layout.align();
61         let memory_type = get_memory_type();
62 
63         if align > 8 {
64             // The requested alignment is greater than 8, but `allocate_pool` is
65             // only guaranteed to provide eight-byte alignment. Allocate extra
66             // space so that we can return an appropriately-aligned pointer
67             // within the allocation.
68             let full_alloc_ptr = if let Ok(ptr) = boot::allocate_pool(memory_type, size + align) {
69                 ptr.as_ptr()
70             } else {
71                 return ptr::null_mut();
72             };
73 
74             // Calculate the offset needed to get an aligned pointer within the
75             // full allocation. If that offset is zero, increase it to `align`
76             // so that we still have space to store the extra pointer described
77             // below.
78             let mut offset = full_alloc_ptr.align_offset(align);
79             if offset == 0 {
80                 offset = align;
81             }
82 
83             // Before returning the aligned allocation, store a pointer to the
84             // full unaligned allocation in the bytes just before the aligned
85             // allocation. We know we have at least eight bytes there due to
86             // adding `align` to the memory allocation size. We also know the
87             // write is appropriately aligned for a `*mut u8` pointer because
88             // `align_ptr` is aligned, and alignments are always powers of two
89             // (as enforced by the `Layout` type).
90             let aligned_ptr = full_alloc_ptr.add(offset);
91             (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
92             aligned_ptr
93         } else {
94             // The requested alignment is less than or equal to eight, and
95             // `allocate_pool` always provides eight-byte alignment, so we can
96             // use `allocate_pool` directly.
97             boot::allocate_pool(memory_type, size)
98                 .map(|ptr| ptr.as_ptr())
99                 .unwrap_or(ptr::null_mut())
100         }
101     }
102 
103     /// Deallocate memory using [`boot::free_pool`].
dealloc(&self, mut ptr: *mut u8, layout: Layout)104     unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
105         if layout.align() > 8 {
106             // Retrieve the pointer to the full allocation that was packed right
107             // before the aligned allocation in `alloc`.
108             ptr = (ptr as *const *mut u8).sub(1).read();
109         }
110 
111         // OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
112         let ptr = NonNull::new(ptr).unwrap();
113 
114         // Warning: this will panic after exiting boot services.
115         boot::free_pool(ptr).unwrap();
116     }
117 }
118