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, AtomicPtr)> = 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 = Cell::new(false); + // Position 0 = SystemTable + // Position 1 = ImageHandle + static GLOBALS: Cell, NonNull)>> = Cell::new(None); +} /// Initializes the global System Table and Image Handle pointers. /// @@ -25,7 +28,7 @@ static GLOBALS: OnceLock<(AtomicPtr, AtomicPtr)> = OnceLock::new /// /// This function must not be called more than once. pub unsafe fn init_globals(handle: NonNull, system_table: NonNull) { - 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 { /// 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::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::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> { + if BOOT_SERVICES_FLAG.get() { + let system_table: NonNull = 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( handle: NonNull, mut protocol_guid: Guid, ) -> io::Result> { - 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 { - try_boot_services().unwrap() +pub(crate) fn create_event( + signal: u32, + tpl: efi::Tpl, + handler: Option, + context: *mut crate::ffi::c_void, +) -> io::Result> { + 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> { - let system_table: NonNull = 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) -> 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>> = 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() -> std_io::Result { @@ -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(); +}