tartan_uefi/
allocator.rs

1//! [`BootServices`](super::BootServices)-based heap allocator to support the standard
2//! [alloc] crate.
3
4use super::global::SYSTEM_TABLE;
5use super::MemoryType;
6use core::alloc::{GlobalAlloc, Layout};
7use core::ffi::c_void;
8use core::ptr;
9
10#[allow(clippy::module_name_repetitions)]
11pub struct BootAllocator;
12
13impl BootAllocator {
14    // UEFI spec guarantees 8-byte alignment from allocate_pool()
15    const UEFI_ALIGNMENT: usize = 8;
16
17    const VOID_PTR_LAYOUT: Layout = Layout::new::<*const c_void>();
18
19    fn extend_layout_for_alignment(orig_layout: Layout) -> (Layout, usize) {
20        // To manually align the memory region, so we need extra space for:
21        //  * A place to store the original pointer that allocate_pool() gave us, so that
22        //    we can recover it for free_pool()
23        //  * One extra unit of alignment to accommodate any shifting we need to do to
24        //    align the pointer
25        let (saved_ptr_layout, saved_ptr_offset) = orig_layout
26            .extend(Self::VOID_PTR_LAYOUT)
27            .expect("Could not construct extended layout for alignment structure");
28        let adjusted_size = saved_ptr_layout.size() + saved_ptr_layout.align();
29        let adjusted_layout =
30            Layout::from_size_align(adjusted_size, saved_ptr_layout.align())
31                .expect("Could not construct extended layout for alignable allocation");
32        (adjusted_layout, saved_ptr_offset)
33    }
34
35    unsafe fn alloc(layout: Layout) -> *mut u8 {
36        let system_table = SYSTEM_TABLE.expect("System table not initialized");
37        let boot_services =
38            (*system_table).boot_services.expect("Boot services unavailable");
39
40        let (adjusted_layout, saved_ptr_offset) =
41            if layout.align() <= Self::UEFI_ALIGNMENT {
42                // allocate_pool() will give the proper alignment. No adjustment needed.
43                (layout, 0)
44            } else {
45                Self::extend_layout_for_alignment(layout)
46            };
47
48        let mut buffer: *mut c_void = ptr::null_mut();
49        (boot_services.allocate_pool)(
50            MemoryType::LoaderData,
51            adjusted_layout.size(),
52            &raw mut buffer,
53        )
54        .into_result()
55        .expect("allocate_pool() failed");
56
57        if adjusted_layout.align() > Self::UEFI_ALIGNMENT {
58            let orig_address = buffer as usize;
59            let align_offset = orig_address % adjusted_layout.align();
60            if align_offset != 0 {
61                // Shift to accommodate alignment
62                let shifted_ptr = orig_address + (adjusted_layout.align() - align_offset);
63                buffer = shifted_ptr as *mut c_void;
64            }
65            // Store the original pointer after the object we are allocating. This happens
66            // even if we didn't have make an alignment shift, because dealloc() has no
67            // way of telling whether we did.
68            let saved_ptr_ptr = (buffer as usize + saved_ptr_offset) as *mut *mut c_void;
69            *saved_ptr_ptr = orig_address as *mut c_void;
70        }
71
72        buffer.cast()
73    }
74
75    unsafe fn dealloc(buffer: *mut u8, layout: Layout) {
76        let system_table = SYSTEM_TABLE.expect("System table not initialized");
77        let boot_services =
78            (*system_table).boot_services.expect("Boot services unavailable");
79
80        let original_ptr: *mut c_void = if layout.align() <= Self::UEFI_ALIGNMENT {
81            // We gave Rust the pointer that came from allocate_pool() directly
82            buffer.cast()
83        } else {
84            // This pointer was manually aligned, which means it isn't the same as the
85            // pointer that allocate_pool() gave us. We have to recover the original
86            // pointer to pass to free_pool().
87            let (_, saved_ptr_offset) = Self::extend_layout_for_alignment(layout);
88            let saved_ptr_ptr = (buffer as usize + saved_ptr_offset) as *mut *mut c_void;
89            *saved_ptr_ptr
90        };
91
92        (boot_services.free_pool)(original_ptr)
93            .into_result()
94            .expect("Could not free memory");
95    }
96}
97
98unsafe impl GlobalAlloc for BootAllocator {
99    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
100        BootAllocator::alloc(layout)
101    }
102
103    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
104        BootAllocator::dealloc(ptr, layout);
105    }
106}