1//! This module implements Rust's global allocator interface using UEFI's memory allocation functions.
2//!
3//! Enabling the `alloc` optional feature in your app will allow you to use Rust's higher-level data structures,
4//! like boxes, vectors, hash maps, linked lists and so on.
5//!
6//! # Usage
7//!
8//! Call the `init` function with a reference to the boot services table.
9//! Failure to do so before calling a memory allocating function will panic.
10//!
11//! Call the `exit_boot_services` function before exiting UEFI boot services.
12//! Failure to do so will turn subsequent allocation into undefined behaviour.
13
14use core::alloc::{GlobalAlloc, Layout};
15use core::ptr::{self, NonNull};
16
17use crate::table::boot::{BootServices, MemoryType};
18
19/// Reference to the boot services table, used to call the pool memory allocation functions.
20///
21/// The inner pointer is only safe to dereference if UEFI boot services have not been
22/// exited by the host application yet.
23static mut BOOT_SERVICES: Option<NonNull<BootServices>> = None;
24
25/// Initializes the allocator.
26///
27/// # Safety
28///
29/// This function is unsafe because you _must_ make sure that exit_boot_services
30/// will be called when UEFI boot services will be exited.
31pub unsafe fn init(boot_services: &BootServices) {
32 BOOT_SERVICES = NonNull::new(ptr:boot_services as *const _ as *mut _);
33}
34
35/// Access the boot services
36fn boot_services() -> NonNull<BootServices> {
37 unsafe { BOOT_SERVICES.expect(msg:"Boot services are unavailable or have been exited") }
38}
39
40/// Notify the allocator library that boot services are not safe to call anymore
41///
42/// You must arrange for this function to be called on exit from UEFI boot services
43pub fn exit_boot_services() {
44 unsafe {
45 BOOT_SERVICES = None;
46 }
47}
48
49/// Allocator which uses the UEFI pool allocation functions.
50///
51/// Only valid for as long as the UEFI boot services are available.
52pub struct Allocator;
53
54unsafe impl GlobalAlloc for Allocator {
55 /// Allocate memory using [`BootServices::allocate_pool`]. The allocation is
56 /// of type [`MemoryType::LOADER_DATA`].
57 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
58 let mem_ty = MemoryType::LOADER_DATA;
59 let size = layout.size();
60 let align = layout.align();
61
62 if align > 8 {
63 // The requested alignment is greater than 8, but `allocate_pool` is
64 // only guaranteed to provide eight-byte alignment. Allocate extra
65 // space so that we can return an appropriately-aligned pointer
66 // within the allocation.
67 let full_alloc_ptr =
68 if let Ok(ptr) = boot_services().as_ref().allocate_pool(mem_ty, size + align) {
69 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_services()
98 .as_ref()
99 .allocate_pool(mem_ty, size)
100 .unwrap_or(ptr::null_mut())
101 }
102 }
103
104 /// Deallocate memory using [`BootServices::free_pool`].
105 unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
106 if layout.align() > 8 {
107 // Retrieve the pointer to the full allocation that was packed right
108 // before the aligned allocation in `alloc`.
109 ptr = (ptr as *const *mut u8).sub(1).read();
110 }
111 boot_services().as_ref().free_pool(ptr).unwrap();
112 }
113}
114
115#[global_allocator]
116static ALLOCATOR: Allocator = Allocator;
117