Merge pull request rust-osdev#979 from nicholasbishop/bishop-atomic-allocator

uefi: Use atomics instead of `static mut` in allocator

Author: phip1611

CHANGELOG.md

@@ -14,6 +14,8 @@
 - `Logger` no longer requires exterior mutability. `Logger::new` is now `const`,
   takes no arguments, and creates the logger in a disabled state. Call
   `Logger::set_output` to enable it.
+- `uefi::allocator::init` now takes a `&mut SystemTable<Boot>` instead of
+  `&BootServices`.
 
 ## uefi-macros - [Unreleased]
 

uefi-services/src/lib.rs

@@ -101,10 +101,10 @@ pub fn init(st: &mut SystemTable<Boot>) -> Result<Option<Event>> {
         #[cfg(feature = "logger")]
         init_logger(st);
 
-        let boot_services = st.boot_services();
-        uefi::allocator::init(boot_services);
+        uefi::allocator::init(st);
 
         // Schedule these tools to be disabled on exit from UEFI boot services
+        let boot_services = st.boot_services();
         boot_services
             .create_event(
                 EventType::SIGNAL_EXIT_BOOT_SERVICES,

uefi/src/allocator.rs

@@ -12,49 +12,54 @@
 //! Failure to do so will turn subsequent allocation into undefined behaviour.
 
 use core::alloc::{GlobalAlloc, Layout};
-use core::ptr::{self, NonNull};
+use core::ffi::c_void;
+use core::ptr;
+use core::sync::atomic::{AtomicPtr, AtomicU32, Ordering};
 
 use crate::proto::loaded_image::LoadedImage;
 use crate::table::boot::{BootServices, MemoryType};
+use crate::table::{Boot, SystemTable};
 
-/// Reference to the boot services table, used to call the pool memory allocation functions.
+/// Reference to the system table, used to call the boot services pool memory
+/// allocation functions.
 ///
-/// The inner pointer is only safe to dereference if UEFI boot services have not been
+/// The pointer is only safe to dereference if UEFI boot services have not been
 /// exited by the host application yet.
-static mut BOOT_SERVICES: Option<NonNull<BootServices>> = None;
+static SYSTEM_TABLE: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());
 
 /// The memory type used for pool memory allocations.
-/// TODO: Use OnceCell when stablilized.
-static mut MEMORY_TYPE: MemoryType = MemoryType::LOADER_DATA;
+static MEMORY_TYPE: AtomicU32 = AtomicU32::new(MemoryType::LOADER_DATA.0);
 
 /// Initializes the allocator.
 ///
 /// # Safety
 ///
 /// This function is unsafe because you _must_ make sure that exit_boot_services
 /// will be called when UEFI boot services will be exited.
-pub unsafe fn init(boot_services: &BootServices) {
-    BOOT_SERVICES = NonNull::new(boot_services as *const _ as *mut _);
+pub unsafe fn init(system_table: &mut SystemTable<Boot>) {
+    SYSTEM_TABLE.store(system_table.as_ptr().cast_mut(), Ordering::Release);
 
+    let boot_services = system_table.boot_services();
     if let Ok(loaded_image) =
         boot_services.open_protocol_exclusive::<LoadedImage>(boot_services.image_handle())
     {
-        MEMORY_TYPE = loaded_image.data_type()
+        MEMORY_TYPE.store(loaded_image.data_type().0, Ordering::Release);
     }
 }
 
 /// Access the boot services
-fn boot_services() -> NonNull<BootServices> {
-    unsafe { BOOT_SERVICES.expect("Boot services are unavailable or have been exited") }
+fn boot_services() -> *const BootServices {
+    let ptr = SYSTEM_TABLE.load(Ordering::Acquire);
+    let system_table =
+        unsafe { SystemTable::from_ptr(ptr) }.expect("The system table handle is not available");
+    system_table.boot_services()
 }
 
 /// Notify the allocator library that boot services are not safe to call anymore
 ///
 /// You must arrange for this function to be called on exit from UEFI boot services
 pub fn exit_boot_services() {
-    unsafe {
-        BOOT_SERVICES = None;
-    }
+    SYSTEM_TABLE.store(ptr::null_mut(), Ordering::Release);
 }
 
 /// Allocator which uses the UEFI pool allocation functions.
@@ -70,20 +75,21 @@ unsafe impl GlobalAlloc for Allocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         let size = layout.size();
         let align = layout.align();
+        let memory_type = MemoryType(MEMORY_TYPE.load(Ordering::Acquire));
+
+        let boot_services = &*boot_services();
 
         if align > 8 {
             // The requested alignment is greater than 8, but `allocate_pool` is
             // only guaranteed to provide eight-byte alignment. Allocate extra
             // space so that we can return an appropriately-aligned pointer
             // within the allocation.
-            let full_alloc_ptr = if let Ok(ptr) = boot_services()
-                .as_ref()
-                .allocate_pool(MEMORY_TYPE, size + align)
-            {
-                ptr
-            } else {
-                return ptr::null_mut();
-            };
+            let full_alloc_ptr =
+                if let Ok(ptr) = boot_services.allocate_pool(memory_type, size + align) {
+                    ptr
+                } else {
+                    return ptr::null_mut();
+                };
 
             // Calculate the offset needed to get an aligned pointer within the
             // full allocation. If that offset is zero, increase it to `align`
@@ -108,9 +114,8 @@ unsafe impl GlobalAlloc for Allocator {
             // The requested alignment is less than or equal to eight, and
             // `allocate_pool` always provides eight-byte alignment, so we can
             // use `allocate_pool` directly.
-            boot_services()
-                .as_ref()
-                .allocate_pool(MEMORY_TYPE, size)
+            boot_services
+                .allocate_pool(memory_type, size)
                 .unwrap_or(ptr::null_mut())
         }
     }
@@ -122,7 +127,7 @@ unsafe impl GlobalAlloc for Allocator {
             // before the aligned allocation in `alloc`.
             ptr = (ptr as *const *mut u8).sub(1).read();
         }
-        boot_services().as_ref().free_pool(ptr).unwrap();
+        (*boot_services()).free_pool(ptr).unwrap();
     }
 }