1 // Copyright 2023, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 use crate::defs::{EFI_MEMORY_TYPE_LOADER_DATA, EFI_STATUS_ALREADY_STARTED};
16 use crate::{EfiEntry, EfiResult};
17 
18 use core::alloc::{GlobalAlloc, Layout};
19 use core::ptr::null_mut;
20 
21 /// Implement a global allocator using `EFI_BOOT_SERVICES.AllocatePool()/FreePool()`
22 pub enum EfiAllocator {
23     Uninitialized,
24     Initialized(EfiEntry),
25     Exited,
26 }
27 
28 #[global_allocator]
29 static mut EFI_GLOBAL_ALLOCATOR: EfiAllocator = EfiAllocator::Uninitialized;
30 
31 /// An internal API to obtain library internal global EfiEntry.
internal_efi_entry() -> Option<&'static EfiEntry>32 pub(crate) fn internal_efi_entry() -> Option<&'static EfiEntry> {
33     // SAFETY:
34     // For now, `EfiAllocator` is only modified in `init_efi_global_alloc()` when `EfiAllocator` is
35     // being initialized or in `exit_efi_global_alloc` after `EFI_BOOT_SERVICES.
36     // ExitBootServices()` is called, where there should be no event/notification function that can
37     // be triggered. Therefore, it should be safe from race condition.
38     unsafe { EFI_GLOBAL_ALLOCATOR.get_efi_entry() }
39 }
40 
41 /// Initializes global allocator.
42 ///
43 /// # Safety
44 ///
45 /// This function modifies global variable `EFI_GLOBAL_ALLOCATOR`. It should only be called when
46 /// there is no event/notification function that can be triggered or modify it. Otherwise there
47 /// is a risk of race condition.
init_efi_global_alloc(efi_entry: EfiEntry) -> EfiResult<()>48 pub(crate) unsafe fn init_efi_global_alloc(efi_entry: EfiEntry) -> EfiResult<()> {
49     // SAFETY: See SAFETY of `internal_efi_entry()`
50     unsafe {
51         match EFI_GLOBAL_ALLOCATOR {
52             EfiAllocator::Uninitialized => {
53                 EFI_GLOBAL_ALLOCATOR = EfiAllocator::Initialized(efi_entry);
54                 Ok(())
55             }
56             _ => Err(EFI_STATUS_ALREADY_STARTED.into()),
57         }
58     }
59 }
60 
61 /// Internal API to invalidate global allocator after ExitBootService().
62 ///
63 /// # Safety
64 ///
65 /// This function modifies global variable `EFI_GLOBAL_ALLOCATOR`. It should only be called when
66 /// there is no event/notification function that can be triggered or modify it. Otherwise there
67 /// is a risk of race condition.
exit_efi_global_alloc()68 pub(crate) unsafe fn exit_efi_global_alloc() {
69     // SAFETY: See SAFETY of `internal_efi_entry()`
70     unsafe {
71         EFI_GLOBAL_ALLOCATOR = EfiAllocator::Exited;
72     }
73 }
74 
75 impl EfiAllocator {
76     /// Returns a reference to the EfiEntry.
get_efi_entry(&self) -> Option<&EfiEntry>77     fn get_efi_entry(&self) -> Option<&EfiEntry> {
78         match self {
79             EfiAllocator::Initialized(ref entry) => Some(entry),
80             _ => None,
81         }
82     }
83 
84     /// Allocate memory via EFI_BOOT_SERVICES.
allocate(&self, size: usize) -> *mut u885     fn allocate(&self, size: usize) -> *mut u8 {
86         match self
87             .get_efi_entry()
88             .unwrap()
89             .system_table()
90             .boot_services()
91             .allocate_pool(EFI_MEMORY_TYPE_LOADER_DATA, size)
92         {
93             Ok(p) => p as *mut _,
94             _ => null_mut(),
95         }
96     }
97 
98     /// Deallocate memory previously allocated by `Self::allocate()`. Passing invalid pointer will
99     /// cause the method to panic.
deallocate(&self, ptr: *mut u8)100     fn deallocate(&self, ptr: *mut u8) {
101         match self.get_efi_entry() {
102             Some(ref entry) => {
103                 entry.system_table().boot_services().free_pool(ptr as *mut _).unwrap();
104             }
105             // After EFI_BOOT_SERVICES.ExitBootServices(), all allocated memory is considered
106             // leaked and under full ownership of subsequent OS loader code.
107             _ => {}
108         }
109     }
110 }
111 
112 unsafe impl GlobalAlloc for EfiAllocator {
alloc(&self, layout: Layout) -> *mut u8113     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
114         let size = layout.size();
115         let align = layout.align();
116         // TODO(300168989): `EFI_BOOT_SERVICES.AllocatePool()` is only 8-byte aligned. Add support
117         // for arbitrary alignment.
118         // `AllocatePool()` can be slow for allocating large buffers. In this case,
119         // `AllocatePages()` is recommended.
120         assert_eq!(8usize.checked_rem(align).unwrap(), 0);
121         self.allocate(size)
122     }
123 
dealloc(&self, ptr: *mut u8, _layout: Layout)124     unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
125         self.deallocate(ptr);
126     }
127 }
128 
129 /// API for allocating raw memory via EFI_BOOT_SERVICES
efi_malloc(size: usize) -> *mut u8130 pub fn efi_malloc(size: usize) -> *mut u8 {
131     // SAFETY: See SAFETY of `internal_efi_entry()`.
132     unsafe { EFI_GLOBAL_ALLOCATOR.allocate(size) }
133 }
134 
135 /// API for deallocating raw memory previously allocated by `efi_malloc()`. Passing invalid
136 /// pointer will cause the function to panic.
efi_free(ptr: *mut u8)137 pub fn efi_free(ptr: *mut u8) {
138     // SAFETY: See SAFETY of `internal_efi_entry()`.
139     unsafe { EFI_GLOBAL_ALLOCATOR.deallocate(ptr) }
140 }
141