Handle ExitBootServices

- Make BootServices unavailable if ExitBootServices event is signaled.
- Use thread locals for SystemTable and ImageHandle

Signed-off-by: Ayush Singh <ayushsingh1325@gmail.com>

4 files changed, 116 insertions, 26 deletions

diff --git a/library/std/src/os/uefi/env.rs b/library/std/src/os/uefi/env.rs
index 71bf0ff8ff3..e62441596f3 100644
--- a/library/std/src/os/uefi/env.rs
+++ b/library/std/src/os/uefi/env.rs
@@ -2,14 +2,17 @@
 
 #![unstable(feature = "uefi_std", issue = "100499")]
 
-use crate::ffi::c_void;
-use crate::ptr::NonNull;
-use crate::sync::atomic::{AtomicPtr, Ordering};
-use crate::sync::OnceLock;
+use crate::{cell::Cell, ffi::c_void, ptr::NonNull};
 
-// Position 0 = SystemTable
-// Position 1 = ImageHandle
-static GLOBALS: OnceLock<(AtomicPtr<c_void>, AtomicPtr<c_void>)> = OnceLock::new();
+// Since UEFI is single-threaded, making the global variables thread local should be safe.
+thread_local! {
+    // Flag to check if BootServices are still valid.
+    // Start with assuming that they are not available
+    static BOOT_SERVICES_FLAG: Cell<bool> = Cell::new(false);
+    // Position 0 = SystemTable
+    // Position 1 = ImageHandle
+    static GLOBALS: Cell<Option<(NonNull<c_void>, NonNull<c_void>)>> = Cell::new(None);
+}
 
 /// Initializes the global System Table and Image Handle pointers.
 ///
@@ -25,7 +28,7 @@ static GLOBALS: OnceLock<(AtomicPtr<c_void>, AtomicPtr<c_void>)> = OnceLock::new
 ///
 /// This function must not be called more than once.
 pub unsafe fn init_globals(handle: NonNull<c_void>, system_table: NonNull<c_void>) {
-    GLOBALS.set((AtomicPtr::new(system_table.as_ptr()), AtomicPtr::new(handle.as_ptr()))).unwrap()
+    GLOBALS.set(Some((system_table, handle)));
 }
 
 /// Get the SystemTable Pointer.
@@ -43,11 +46,31 @@ pub fn image_handle() -> NonNull<c_void> {
 /// Get the SystemTable Pointer.
 /// This function is mostly intended for places where panic is not an option
 pub(crate) fn try_system_table() -> Option<NonNull<crate::ffi::c_void>> {
-    NonNull::new(GLOBALS.get()?.0.load(Ordering::Acquire))
+    GLOBALS.get().map(|x| x.0)
 }
 
 /// Get the SystemHandle Pointer.
 /// This function is mostly intended for places where panic is not an option
 pub(crate) fn try_image_handle() -> Option<NonNull<crate::ffi::c_void>> {
-    NonNull::new(GLOBALS.get()?.1.load(Ordering::Acquire))
+    GLOBALS.get().map(|x| x.1)
+}
+
+/// Get the BootServices Pointer.
+/// This function also checks if `ExitBootServices` has already been called.
+pub(crate) fn boot_services() -> Option<NonNull<r_efi::efi::BootServices>> {
+    if BOOT_SERVICES_FLAG.get() {
+        let system_table: NonNull<r_efi::efi::SystemTable> = try_system_table()?.cast();
+        let boot_services = unsafe { (*system_table.as_ptr()).boot_services };
+        NonNull::new(boot_services)
+    } else {
+        None
+    }
+}
+
+pub(crate) fn enable_boot_services() {
+    BOOT_SERVICES_FLAG.set(true);
+}
+
+pub(crate) fn disable_boot_services() {
+    BOOT_SERVICES_FLAG.set(false);
 }
diff --git a/library/std/src/sys/uefi/alloc.rs b/library/std/src/sys/uefi/alloc.rs
index 54d86d0c8a7..d6eb371033c 100644
--- a/library/std/src/sys/uefi/alloc.rs
+++ b/library/std/src/sys/uefi/alloc.rs
@@ -8,6 +8,11 @@ const MEMORY_TYPE: u32 = r_efi::efi::LOADER_DATA;
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // Return null pointer if boot services are not available
+        if crate::os::uefi::env::boot_services().is_none() {
+            return crate::ptr::null_mut();
+        }
+
         let system_table = match crate::os::uefi::env::try_system_table() {
             None => return crate::ptr::null_mut(),
             Some(x) => x.as_ptr() as *mut _,
@@ -18,6 +23,11 @@ unsafe impl GlobalAlloc for System {
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Do nothing if boot services are not available
+        if crate::os::uefi::env::boot_services().is_none() {
+            return;
+        }
+
         let system_table = match crate::os::uefi::env::try_system_table() {
             None => handle_alloc_error(layout),
             Some(x) => x.as_ptr() as *mut _,
diff --git a/library/std/src/sys/uefi/helpers.rs b/library/std/src/sys/uefi/helpers.rs
index e83adabaf16..28a82394226 100644
--- a/library/std/src/sys/uefi/helpers.rs
+++ b/library/std/src/sys/uefi/helpers.rs
@@ -9,12 +9,18 @@
 //! - Protocols are produced and consumed.
 //! - More information about protocols can be found [here](https://edk2-docs.gitbook.io/edk-ii-uefi-driver-writer-s-guide/3_foundation/36_protocols_and_handles)
 
-use r_efi::efi::Guid;
+use r_efi::efi::{self, Guid};
 
-use crate::io::{self, const_io_error};
 use crate::mem::{size_of, MaybeUninit};
 use crate::os::uefi;
 use crate::ptr::NonNull;
+use crate::{
+    io::{self, const_io_error},
+    os::uefi::env::boot_services,
+};
+
+const BOOT_SERVICES_UNAVAILABLE: io::Error =
+    const_io_error!(io::ErrorKind::Other, "Boot Services are no longer available");
 
 /// Locate Handles with a particular Protocol GUID
 /// Implemented using `EFI_BOOT_SERVICES.LocateHandles()`
@@ -40,7 +46,7 @@ pub(crate) fn locate_handles(mut guid: Guid) -> io::Result<Vec<NonNull<crate::ff
         if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
     }
 
-    let boot_services = boot_services();
+    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
     let mut buf_len = 0usize;
 
     // This should always fail since the size of buffer is 0. This call should update the buf_len
@@ -76,7 +82,7 @@ pub(crate) fn open_protocol<T>(
     handle: NonNull<crate::ffi::c_void>,
     mut protocol_guid: Guid,
 ) -> io::Result<NonNull<T>> {
-    let boot_services = boot_services();
+    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
     let system_handle = uefi::env::image_handle();
     let mut protocol: MaybeUninit<*mut T> = MaybeUninit::uninit();
 
@@ -267,14 +273,32 @@ pub(crate) fn status_to_io_error(s: r_efi::efi::Status) -> io::Error {
     }
 }
 
-/// Get the BootServices Pointer.
-pub(crate) fn boot_services() -> NonNull<r_efi::efi::BootServices> {
-    try_boot_services().unwrap()
+pub(crate) fn create_event(
+    signal: u32,
+    tpl: efi::Tpl,
+    handler: Option<efi::EventNotify>,
+    context: *mut crate::ffi::c_void,
+) -> io::Result<NonNull<crate::ffi::c_void>> {
+    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let mut exit_boot_service_event: r_efi::efi::Event = crate::ptr::null_mut();
+    let r = unsafe {
+        let create_event = (*boot_services.as_ptr()).create_event;
+        (create_event)(signal, tpl, handler, context, &mut exit_boot_service_event)
+    };
+    if r.is_error() {
+        Err(status_to_io_error(r))
+    } else {
+        NonNull::new(exit_boot_service_event)
+            .ok_or(const_io_error!(io::ErrorKind::Other, "null protocol"))
+    }
 }
-/// Get the BootServices Pointer.
-/// This function is mostly intended for places where panic is not an option
-pub(crate) fn try_boot_services() -> Option<NonNull<r_efi::efi::BootServices>> {
-    let system_table: NonNull<r_efi::efi::SystemTable> = uefi::env::try_system_table()?.cast();
-    let boot_services = unsafe { (*system_table.as_ptr()).boot_services };
-    NonNull::new(boot_services)
+
+pub(crate) fn close_event(evt: NonNull<crate::ffi::c_void>) -> io::Result<()> {
+    let boot_services = boot_services().ok_or(BOOT_SERVICES_UNAVAILABLE)?;
+    let r = unsafe {
+        let close_event = (*boot_services.as_ptr()).close_event;
+        (close_event)(evt.as_ptr())
+    };
+
+    if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
 }
diff --git a/library/std/src/sys/uefi/mod.rs b/library/std/src/sys/uefi/mod.rs
index fed44b4f49d..fa33b681a19 100644
--- a/library/std/src/sys/uefi/mod.rs
+++ b/library/std/src/sys/uefi/mod.rs
@@ -52,6 +52,7 @@ pub(crate) mod helpers;
 #[cfg(test)]
 mod tests;
 
+use crate::cell::Cell;
 use crate::io as std_io;
 use crate::os::uefi;
 use crate::ptr::NonNull;
@@ -60,6 +61,10 @@ pub mod memchr {
     pub use core::slice::memchr::{memchr, memrchr};
 }
 
+thread_local! {
+    static EXIT_BOOT_SERVICE_EVENT: Cell<Option<NonNull<crate::ffi::c_void>>> = Cell::new(None);
+}
+
 /// # SAFETY
 /// - must be called only once during runtime initialization.
 /// - argc must be 2.
@@ -68,13 +73,32 @@ pub(crate) unsafe fn init(argc: isize, argv: *const *const u8, _sigpipe: u8) {
     assert_eq!(argc, 2);
     let image_handle = unsafe { NonNull::new(*argv as *mut crate::ffi::c_void).unwrap() };
     let system_table = unsafe { NonNull::new(*argv.add(1) as *mut crate::ffi::c_void).unwrap() };
-    unsafe { crate::os::uefi::env::init_globals(image_handle, system_table) };
+    unsafe { uefi::env::init_globals(image_handle, system_table) };
+    // Enable boot services once GLOBALS are initialized
+    uefi::env::enable_boot_services();
+
+    // Register exit boot services handler
+    match helpers::create_event(
+        r_efi::efi::EVT_SIGNAL_EXIT_BOOT_SERVICES,
+        r_efi::efi::TPL_NOTIFY,
+        Some(exit_boot_service_handler),
+        crate::ptr::null_mut(),
+    ) {
+        Ok(x) => {
+            EXIT_BOOT_SERVICE_EVENT.set(Some(x));
+        }
+        Err(_) => abort_internal(),
+    }
 }
 
 /// # SAFETY
 /// this is not guaranteed to run, for example when the program aborts.
 /// - must be called only once during runtime cleanup.
-pub unsafe fn cleanup() {}
+pub unsafe fn cleanup() {
+    if let Some(exit_boot_service_event) = EXIT_BOOT_SERVICE_EVENT.take() {
+        let _ = helpers::close_event(exit_boot_service_event);
+    }
+}
 
 #[inline]
 pub const fn unsupported<T>() -> std_io::Result<T> {
@@ -98,8 +122,12 @@ pub fn decode_error_kind(code: i32) -> crate::io::ErrorKind {
 }
 
 pub fn abort_internal() -> ! {
+    if let Some(exit_boot_service_event) = EXIT_BOOT_SERVICE_EVENT.take() {
+        let _ = helpers::close_event(exit_boot_service_event);
+    }
+
     if let (Some(boot_services), Some(handle)) =
-        (helpers::try_boot_services(), uefi::env::try_image_handle())
+        (uefi::env::boot_services(), uefi::env::try_image_handle())
     {
         let _ = unsafe {
             ((*boot_services.as_ptr()).exit)(
@@ -155,3 +183,8 @@ fn get_random() -> Option<(u64, u64)> {
     }
     None
 }
+
+/// Disable access to BootServices if `EVT_SIGNAL_EXIT_BOOT_SERVICES` is signaled
+extern "efiapi" fn exit_boot_service_handler(_e: r_efi::efi::Event, _ctx: *mut crate::ffi::c_void) {
+    uefi::env::disable_boot_services();
+}