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>

1 files changed, 10 insertions, 0 deletions

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 _,