Improve std::uefi::env APIs

1. Can now only initialize Global variales once.
2. Panics if `image_handle` or `system_table` is called before
   initializing them
3. Special case for `abort_internal()` which can be called before the
   initialization of Globals.

Signed-off-by: Ayush Singh <[email protected]>

Author: Ayush1325

library/std/src/os/uefi/env.rs

@@ -3,9 +3,11 @@
 use crate::ffi::c_void;
 use crate::ptr::NonNull;
 use crate::sync::atomic::{AtomicPtr, Ordering};
+use crate::sync::Once;
 
 static GLOBAL_SYSTEM_TABLE: AtomicPtr<c_void> = AtomicPtr::new(crate::ptr::null_mut());
 static GLOBAL_IMAGE_HANDLE: AtomicPtr<c_void> = AtomicPtr::new(crate::ptr::null_mut());
+pub(crate) static GLOBALS: Once = Once::new();
 
 /// Initializes the global System Table and Image Handle pointers.
 ///
@@ -22,18 +24,22 @@ static GLOBAL_IMAGE_HANDLE: AtomicPtr<c_void> = AtomicPtr::new(crate::ptr::null_
 /// This function must not be called more than once.
 #[unstable(feature = "uefi_std", issue = "100499")]
 pub unsafe fn init_globals(handle: NonNull<c_void>, system_table: NonNull<c_void>) {
-    GLOBAL_SYSTEM_TABLE.store(system_table.as_ptr(), Ordering::Release);
-    GLOBAL_IMAGE_HANDLE.store(handle.as_ptr(), Ordering::Release);
+    GLOBALS.call_once(|| {
+        GLOBAL_SYSTEM_TABLE.store(system_table.as_ptr(), Ordering::Release);
+        GLOBAL_IMAGE_HANDLE.store(handle.as_ptr(), Ordering::Release);
+    })
 }
 
 /// Get the SystemTable Pointer.
 #[unstable(feature = "uefi_std", issue = "100499")]
-pub fn system_table() -> Option<NonNull<c_void>> {
-    NonNull::new(GLOBAL_SYSTEM_TABLE.load(Ordering::Acquire))
+pub fn system_table() -> NonNull<c_void> {
+    assert!(GLOBALS.is_completed());
+    unsafe { NonNull::new_unchecked(GLOBAL_SYSTEM_TABLE.load(Ordering::Acquire)) }
 }
 
 /// Get the SystemHandle Pointer.
 #[unstable(feature = "uefi_std", issue = "100499")]
-pub fn image_handle() -> Option<NonNull<c_void>> {
-    NonNull::new(GLOBAL_IMAGE_HANDLE.load(Ordering::Acquire))
+pub fn image_handle() -> NonNull<c_void> {
+    assert!(GLOBALS.is_completed());
+    unsafe { NonNull::new_unchecked(GLOBAL_IMAGE_HANDLE.load(Ordering::Acquire)) }
 }

library/std/src/sys/uefi/common.rs

@@ -12,15 +12,11 @@ pub(crate) const BOOT_SERVICES_ERROR: io::Error =
     const_io_error!(io::ErrorKind::Other, "failed to acquire boot services",);
 pub(crate) const RUNTIME_SERVICES_ERROR: io::Error =
     const_io_error!(io::ErrorKind::Other, "failed to acquire runtime services",);
-pub(crate) const SYSTEM_TABLE_ERROR: io::Error =
-    const_io_error!(io::ErrorKind::Other, "failed to acquire system table",);
-pub(crate) const IMAGE_HANDLE_ERROR: io::Error =
-    const_io_error!(io::ErrorKind::Other, "failed to acquire system handle",);
 
 /// Get the Protocol for current system handle.
 /// Note: Some protocols need to be manually freed. It is the callers responsibility to do so.
 pub(crate) fn get_current_handle_protocol<T>(protocol_guid: Guid) -> Option<NonNull<T>> {
-    let system_handle = uefi::env::image_handle()?;
+    let system_handle = uefi::env::image_handle();
     open_protocol(system_handle, protocol_guid).ok()
 }
 
@@ -240,7 +236,7 @@ pub(crate) fn open_protocol<T>(
     mut protocol_guid: Guid,
 ) -> io::Result<NonNull<T>> {
     let boot_services = get_boot_services().ok_or(BOOT_SERVICES_ERROR)?;
-    let system_handle = uefi::env::image_handle().ok_or(IMAGE_HANDLE_ERROR)?;
+    let system_handle = uefi::env::image_handle();
     let mut protocol: MaybeUninit<*mut T> = MaybeUninit::uninit();
 
     let r = unsafe {
@@ -469,14 +465,14 @@ where
 
 /// Get the BootServices Pointer.
 pub(crate) fn get_boot_services() -> Option<NonNull<r_efi::efi::BootServices>> {
-    let system_table: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table()?.cast();
+    let system_table: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table().cast();
     let boot_services = unsafe { (*system_table.as_ptr()).boot_services };
     NonNull::new(boot_services)
 }
 
 /// Get the RuntimeServices Pointer.
 pub(crate) fn get_runtime_services() -> Option<NonNull<r_efi::efi::RuntimeServices>> {
-    let system_table: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table()?.cast();
+    let system_table: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table().cast();
     let runtime_services = unsafe { (*system_table.as_ptr()).runtime_services };
     NonNull::new(runtime_services)
 }

library/std/src/sys/uefi/mod.rs

@@ -78,21 +78,22 @@ pub fn decode_error_kind(code: i32) -> crate::io::ErrorKind {
 }
 
 pub fn abort_internal() -> ! {
-    // First try to use EFI_BOOT_SERVICES.Exit()
-    if let (Some(boot_services), Some(handle)) =
-        (common::get_boot_services(), uefi::env::image_handle())
-    {
-        let _ = unsafe {
-            ((*boot_services.as_ptr()).exit)(
-                handle.as_ptr(),
-                r_efi::efi::Status::ABORTED,
-                0,
-                crate::ptr::null_mut(),
-            )
-        };
+    if uefi::env::GLOBALS.is_completed() {
+        let handle = uefi::env::image_handle();
+        // First try to use EFI_BOOT_SERVICES.Exit()
+        if let Some(boot_services) = common::get_boot_services() {
+            let _ = unsafe {
+                ((*boot_services.as_ptr()).exit)(
+                    handle.as_ptr(),
+                    r_efi::efi::Status::ABORTED,
+                    0,
+                    crate::ptr::null_mut(),
+                )
+            };
+        }
     }
 
-    // In case SystemTable and SystemHandle cannot be reached, use `core::intrinsics::abort`
+    // In case SystemTable and ImageHandle cannot be reached, use `core::intrinsics::abort`
     core::intrinsics::abort();
 }
 

library/std/src/sys/uefi/os.rs

@@ -182,12 +182,12 @@ pub fn exit(code: i32) -> ! {
         Err(_) => r_efi::efi::Status::ABORTED,
     };
 
+    let handle = uefi::env::image_handle().cast();
     // First try to use EFI_BOOT_SERVICES.Exit()
-    if let (Some(boot_services), Some(handle)) =
-        (common::get_boot_services(), uefi::env::image_handle())
-    {
-        let _ =
-            unsafe { ((*boot_services.as_ptr()).exit)(handle.as_ptr(), code, 0, [0].as_mut_ptr()) };
+    if let Some(boot_services) = common::get_boot_services() {
+        let _ = unsafe {
+            ((*boot_services.as_ptr()).exit)(handle.as_ptr(), code, 0, crate::ptr::null_mut())
+        };
     }
 
     // If exit is not possible, the call abort

library/std/src/sys/uefi/process.rs

@@ -310,7 +310,7 @@ mod uefi_command {
     impl Command {
         pub(crate) fn load_image(p: &OsStr) -> io::Result<Self> {
             let boot_services = common::get_boot_services().ok_or(common::BOOT_SERVICES_ERROR)?;
-            let system_handle = uefi::env::image_handle().ok_or(common::IMAGE_HANDLE_ERROR)?;
+            let system_handle = uefi::env::image_handle();
             let mut path = super::super::path::device_path_from_os_str(p)?;
             let mut child_handle: MaybeUninit<r_efi::efi::Handle> = MaybeUninit::uninit();
             let r = unsafe {

library/std/src/sys/uefi/stdio.rs

@@ -84,8 +84,7 @@ impl io::Read for Stdin {
                 return pipe::AnonPipe::new(None, None, pipe_protocol).read(buf);
             }
         }
-        let global_system_table =
-            uefi::env::system_table().ok_or(common::SYSTEM_TABLE_ERROR)?.cast();
+        let global_system_table = uefi::env::system_table().cast();
         let con_in = get_con_in(global_system_table)?;
         let con_out = get_con_out(global_system_table)?;
         let wait_for_event = get_wait_for_event()?;
@@ -138,8 +137,7 @@ impl io::Write for Stdout {
             }
         }
 
-        let global_system_table =
-            uefi::env::system_table().ok_or(common::SYSTEM_TABLE_ERROR)?.cast();
+        let global_system_table = uefi::env::system_table().cast();
         let con_out = get_con_out(global_system_table)?;
         unsafe { simple_text_output_write(con_out.as_ptr(), buf) }
     }
@@ -169,8 +167,7 @@ impl io::Write for Stderr {
                 return pipe::AnonPipe::new(None, None, pipe_protocol).write(buf);
             }
         }
-        let global_system_table =
-            uefi::env::system_table().ok_or(common::SYSTEM_TABLE_ERROR)?.cast();
+        let global_system_table = uefi::env::system_table().cast();
         let std_err = get_std_err(global_system_table)?;
         unsafe { simple_text_output_write(std_err.as_ptr(), buf) }
     }