Exploration on implementation of RFC #893.

uefi-macros/src/lib.rs

@@ -209,6 +209,17 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
         );
     }
 
+    if let Some(ref system_table_ident) = system_table_ident {
+        f.block.stmts.insert(
+            0,
+            parse_quote! {
+                unsafe {
+                    #system_table_ident.set_global_system_table();
+                }
+            },
+        );
+    }
+
     let fn_ident = &f.sig.ident;
     // Get an iterator of the function inputs types. This is needed instead of
     // directly using `sig.inputs` because patterns you can use in fn items like

uefi-test-runner/examples/shell_params.rs

@@ -6,7 +6,8 @@
 
 use log::error;
 // ANCHOR: use
-use uefi::{prelude::*, proto::shell_params::ShellParameters};
+use uefi::prelude::*;
+use uefi::proto::shell_params::ShellParameters;
 use uefi_services::println;
 
 extern crate alloc;

uefi-test-runner/src/main.rs

@@ -59,6 +59,8 @@ fn efi_main(image: Handle, mut st: SystemTable<Boot>) -> Status {
     // probably want to test them after exit_boot_services. However,
     // exit_boot_services is currently called during shutdown.
 
+    log::info!("Boot handle count: {}", uefi::boot::check_count());
+
     runtime::test(st.runtime_services());
 
     shutdown(st);

uefi-test-runner/src/proto/console/gop.rs

@@ -22,7 +22,7 @@ pub unsafe fn test(image: Handle, bt: &BootServices) {
         )
         .expect("failed to open Graphics Output Protocol");
 
-    set_graphics_mode(gop, bt);
+    set_graphics_mode(gop);
     fill_color(gop);
     draw_fb(gop);
 
@@ -33,10 +33,10 @@ pub unsafe fn test(image: Handle, bt: &BootServices) {
 }
 
 // Set a larger graphics mode.
-fn set_graphics_mode(gop: &mut GraphicsOutput, bs: &BootServices) {
+fn set_graphics_mode(gop: &mut GraphicsOutput) {
     // We know for sure QEMU has a 1024x768 mode.
     let mode = gop
-        .modes(bs)
+        .modes()
         .find(|mode| {
             let info = mode.info();
             info.resolution() == (1024, 768)

uefi-test-runner/src/proto/pi/mp.rs

@@ -27,9 +27,16 @@ pub fn test(bt: &BootServices) {
     test_get_number_of_processors(mp_support);
     test_get_processor_info(mp_support);
     test_startup_all_aps(mp_support, bt);
+    log::info!("After startup all: {}", uefi::boot::check_count());
     test_startup_this_ap(mp_support, bt);
+    log::info!("After startup this: {}", uefi::boot::check_count());
     test_enable_disable_ap(mp_support);
+    log::info!("Enable/Disable: {}", uefi::boot::check_count());
     test_switch_bsp_and_who_am_i(mp_support);
+    log::info!(
+        "Boot Services Processor & who am i: {}",
+        uefi::boot::check_count()
+    );
 }
 
 fn test_get_number_of_processors(mps: &MpServices) {
@@ -88,6 +95,8 @@ fn test_startup_all_aps(mps: &MpServices, bt: &BootServices) {
         .unwrap();
     assert_eq!(counter.load(Ordering::Relaxed), NUM_CPUS - 1);
 
+    log::info!("After proc_increment atomic: {}", uefi::boot::check_count());
+
     // Make sure that timeout works
     let bt_ptr: *mut c_void = bt as *const _ as *mut _;
     let ret = mps.startup_all_aps(

uefi/src/allocator.rs

@@ -2,60 +2,30 @@
 //!
 //! If the `global_allocator` feature is enabled, the [`Allocator`] will be used
 //! as the global Rust allocator.
-//!
-//! # Usage
-//!
-//! Call the `init` function with a reference to the boot services table.
-//! Failure to do so before calling a memory allocating function will panic.
-//!
-//! Call the `exit_boot_services` function before exiting UEFI boot services.
-//! Failure to do so will turn subsequent allocation into undefined behaviour.
 
 use core::alloc::{GlobalAlloc, Layout};
-use core::ptr::{self, NonNull};
+use core::ptr;
 
+use crate::boot::{allocate_pool, free_pool, image_handle, open_protocol_exclusive};
 use crate::proto::loaded_image::LoadedImage;
 use crate::table::boot::{BootServices, MemoryType};
 
-/// Reference to the boot services table, used to call the pool memory allocation functions.
-///
-/// The inner 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;
-
 /// The memory type used for pool memory allocations.
 /// TODO: Use OnceCell when stablilized.
 static mut MEMORY_TYPE: MemoryType = MemoryType::LOADER_DATA;
 
 /// 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 _);
-
-    if let Ok(loaded_image) =
-        boot_services.open_protocol_exclusive::<LoadedImage>(boot_services.image_handle())
-    {
-        MEMORY_TYPE = loaded_image.data_type()
+pub fn init(_boot_services: &BootServices) {
+    if let Ok(loaded_image) = open_protocol_exclusive::<LoadedImage>(image_handle()) {
+        unsafe { MEMORY_TYPE = loaded_image.data_type() }
     }
 }
 
-/// Access the boot services
-fn boot_services() -> NonNull<BootServices> {
-    unsafe { BOOT_SERVICES.expect("Boot services are unavailable or have been exited") }
-}
-
 /// 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;
-    }
-}
+/// No longer needs to be called.
+#[deprecated]
+pub fn exit_boot_services() {}
 
 /// Allocator which uses the UEFI pool allocation functions.
 ///
@@ -64,7 +34,7 @@ pub fn exit_boot_services() {
 pub struct Allocator;
 
 unsafe impl GlobalAlloc for Allocator {
-    /// Allocate memory using [`BootServices::allocate_pool`]. The allocation is
+    /// Allocate memory using [`allocate_pool`]. The allocation is
     /// of type [`MemoryType::LOADER_DATA`] for UEFI applications, [`MemoryType::BOOT_SERVICES_DATA`]
     /// for UEFI boot drivers and [`MemoryType::RUNTIME_SERVICES_DATA`] for UEFI runtime drivers.
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
@@ -76,10 +46,7 @@ unsafe impl GlobalAlloc for Allocator {
             // 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)
-            {
+            let full_alloc_ptr = if let Ok(ptr) = allocate_pool(MEMORY_TYPE, size + align) {
                 ptr
             } else {
                 return ptr::null_mut();
@@ -108,21 +75,18 @@ 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)
-                .unwrap_or(ptr::null_mut())
+            allocate_pool(MEMORY_TYPE, size).unwrap_or(ptr::null_mut())
         }
     }
 
-    /// Deallocate memory using [`BootServices::free_pool`].
+    /// Deallocate memory using [`free_pool`].
     unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
         if layout.align() > 8 {
             // Retrieve the pointer to the full allocation that was packed right
             // before the aligned allocation in `alloc`.
             ptr = (ptr as *const *mut u8).sub(1).read();
         }
-        boot_services().as_ref().free_pool(ptr).unwrap();
+        free_pool(ptr).unwrap();
     }
 }