diff --git a/library/std/src/sys/pal/uefi/helpers.rs b/library/std/src/sys/pal/uefi/helpers.rs index ba53ed88f37..23aa4da14a7 100644 --- a/library/std/src/sys/pal/uefi/helpers.rs +++ b/library/std/src/sys/pal/uefi/helpers.rs @@ -10,14 +10,16 @@ //! - 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::{self, Guid}; +use r_efi::protocols::{device_path, device_path_to_text}; +use crate::ffi::OsString; +use crate::io::{self, const_io_error}; use crate::mem::{size_of, MaybeUninit}; -use crate::os::uefi; +use crate::os::uefi::{self, env::boot_services, ffi::OsStringExt}; use crate::ptr::NonNull; -use crate::{ - io::{self, const_io_error}, - os::uefi::env::boot_services, -}; +use crate::slice; +use crate::sync::atomic::{AtomicPtr, Ordering}; +use crate::sys_common::wstr::WStrUnits; const BOOT_SERVICES_UNAVAILABLE: io::Error = const_io_error!(io::ErrorKind::Other, "Boot Services are no longer available"); @@ -142,9 +144,74 @@ pub(crate) unsafe fn close_event(evt: NonNull) -> io::Result /// 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 image_handle_protocol(protocol_guid: Guid) -> Option> { - let system_handle = uefi::env::try_image_handle()?; - open_protocol(system_handle, protocol_guid).ok() +pub(crate) fn image_handle_protocol(protocol_guid: Guid) -> io::Result> { + let system_handle = uefi::env::try_image_handle().ok_or(io::const_io_error!( + io::ErrorKind::NotFound, + "Protocol not found in Image handle" + ))?; + open_protocol(system_handle, protocol_guid) +} + +pub(crate) fn device_path_to_text(path: NonNull) -> io::Result { + fn path_to_text( + protocol: NonNull, + path: NonNull, + ) -> io::Result { + let path_ptr: *mut r_efi::efi::Char16 = unsafe { + ((*protocol.as_ptr()).convert_device_path_to_text)( + path.as_ptr(), + // DisplayOnly + r_efi::efi::Boolean::FALSE, + // AllowShortcuts + r_efi::efi::Boolean::FALSE, + ) + }; + + // SAFETY: `convert_device_path_to_text` returns a pointer to a null-terminated UTF-16 + // string, and that string cannot be deallocated prior to dropping the `WStrUnits`, so + // it's safe for `WStrUnits` to use. + let path_len = unsafe { + WStrUnits::new(path_ptr) + .ok_or(io::const_io_error!(io::ErrorKind::InvalidData, "Invalid path"))? + .count() + }; + + let path = OsString::from_wide(unsafe { slice::from_raw_parts(path_ptr.cast(), path_len) }); + + if let Some(boot_services) = crate::os::uefi::env::boot_services() { + let boot_services: NonNull = boot_services.cast(); + unsafe { + ((*boot_services.as_ptr()).free_pool)(path_ptr.cast()); + } + } + + Ok(path) + } + + static LAST_VALID_HANDLE: AtomicPtr = + AtomicPtr::new(crate::ptr::null_mut()); + + if let Some(handle) = NonNull::new(LAST_VALID_HANDLE.load(Ordering::Acquire)) { + if let Ok(protocol) = open_protocol::( + handle, + device_path_to_text::PROTOCOL_GUID, + ) { + return path_to_text(protocol, path); + } + } + + let device_path_to_text_handles = locate_handles(device_path_to_text::PROTOCOL_GUID)?; + for handle in device_path_to_text_handles { + if let Ok(protocol) = open_protocol::( + handle, + device_path_to_text::PROTOCOL_GUID, + ) { + LAST_VALID_HANDLE.store(handle.as_ptr(), Ordering::Release); + return path_to_text(protocol, path); + } + } + + Err(io::const_io_error!(io::ErrorKind::NotFound, "No device path to text protocol found")) } /// Get RuntimeServices diff --git a/library/std/src/sys/pal/uefi/os.rs b/library/std/src/sys/pal/uefi/os.rs index e6693db68e6..58838c5876e 100644 --- a/library/std/src/sys/pal/uefi/os.rs +++ b/library/std/src/sys/pal/uefi/os.rs @@ -1,4 +1,4 @@ -use super::{unsupported, RawOsError}; +use super::{helpers, unsupported, RawOsError}; use crate::error::Error as StdError; use crate::ffi::{OsStr, OsString}; use crate::fmt; @@ -7,6 +7,7 @@ use crate::os::uefi; use crate::path::{self, PathBuf}; use crate::ptr::NonNull; +use r_efi::efi::protocols::{device_path, loaded_image_device_path}; use r_efi::efi::Status; pub fn errno() -> RawOsError { @@ -164,7 +165,10 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { impl StdError for JoinPathsError {} pub fn current_exe() -> io::Result { - unsupported() + let protocol = helpers::image_handle_protocol::( + loaded_image_device_path::PROTOCOL_GUID, + )?; + helpers::device_path_to_text(protocol).map(PathBuf::from) } pub struct Env(!);