diff --git a/uefi-test-runner/src/boot/mod.rs b/uefi-test-runner/src/boot/mod.rs index e23d08118..86b524070 100644 --- a/uefi-test-runner/src/boot/mod.rs +++ b/uefi-test-runner/src/boot/mod.rs @@ -15,12 +15,12 @@ pub fn test(st: &SystemTable) { info!("Testing boot services"); memory::test(bt); misc::test(st); - test_locate_handle_buffer(bt); + test_locate_handles(bt); test_load_image(bt); } -fn test_locate_handle_buffer(bt: &BootServices) { - info!("Testing the `locate_handle_buffer` function"); +fn test_locate_handles(bt: &BootServices) { + info!("Testing the `locate_handle_buffer`/`find_handles` functions"); { // search all handles @@ -51,6 +51,11 @@ fn test_locate_handle_buffer(bt: &BootServices) { *handles, *boot::locate_handle_buffer(SearchType::ByProtocol(&Output::GUID)).unwrap() ); + + // Compare with `boot::find_handles`. This implicitly tests + // `boot::locate_handle` as well. + let handles_vec = boot::find_handles::().unwrap(); + assert_eq!(*handles, handles_vec); } } diff --git a/uefi/src/boot.rs b/uefi/src/boot.rs index cd1fb8211..111e8af98 100644 --- a/uefi/src/boot.rs +++ b/uefi/src/boot.rs @@ -4,11 +4,13 @@ use crate::data_types::PhysicalAddress; use crate::mem::memory_map::{MemoryMapBackingMemory, MemoryMapKey, MemoryMapMeta, MemoryMapOwned}; +use crate::polyfill::maybe_uninit_slice_assume_init_ref; use crate::proto::device_path::DevicePath; use crate::proto::{Protocol, ProtocolPointer}; use crate::table::Revision; use crate::util::opt_nonnull_to_ptr; use core::ffi::c_void; +use core::mem::MaybeUninit; use core::ops::{Deref, DerefMut}; use core::ptr::{self, NonNull}; use core::sync::atomic::{AtomicPtr, Ordering}; @@ -16,6 +18,9 @@ use core::{mem, slice}; use uefi::{table, Char16, Error, Event, Guid, Handle, Result, Status, StatusExt}; use uefi_raw::table::boot::InterfaceType; +#[cfg(feature = "alloc")] +use {alloc::vec::Vec, uefi::ResultExt}; + #[cfg(doc)] use { crate::proto::device_path::LoadedImageDevicePath, crate::proto::loaded_image::LoadedImage, @@ -667,6 +672,47 @@ pub fn locate_device_path( } } +/// Enumerates all handles installed on the system which match a certain query. +/// +/// # Errors +/// +/// * [`Status::NOT_FOUND`]: no matching handles found. +/// * [`Status::BUFFER_TOO_SMALL`]: the buffer is not large enough. The required +/// size (in number of handles, not bytes) will be returned in the error data. +pub fn locate_handle<'buf>( + search_ty: SearchType, + buffer: &'buf mut [MaybeUninit], +) -> Result<&'buf [Handle], Option> { + let bt = boot_services_raw_panicking(); + let bt = unsafe { bt.as_ref() }; + + // Obtain the needed data from the parameters. + let (ty, guid, key) = match search_ty { + SearchType::AllHandles => (0, ptr::null(), ptr::null()), + SearchType::ByRegisterNotify(registration) => { + (1, ptr::null(), registration.0.as_ptr().cast_const()) + } + SearchType::ByProtocol(guid) => (2, guid as *const Guid, ptr::null()), + }; + + let mut buffer_size = buffer.len() * mem::size_of::(); + let status = + unsafe { (bt.locate_handle)(ty, guid, key, &mut buffer_size, buffer.as_mut_ptr().cast()) }; + + let num_handles = buffer_size / mem::size_of::(); + + match status { + Status::SUCCESS => { + let buffer = &buffer[..num_handles]; + // SAFETY: the entries up to `num_handles` have been initialized. + let handles = unsafe { maybe_uninit_slice_assume_init_ref(buffer) }; + Ok(handles) + } + Status::BUFFER_TOO_SMALL => Err(Error::new(status, Some(num_handles))), + _ => Err(Error::new(status, None)), + } +} + /// Returns an array of handles that support the requested protocol in a /// pool-allocated buffer. /// @@ -698,6 +744,47 @@ pub fn locate_handle_buffer(search_ty: SearchType) -> Result { }) } +/// Returns all the handles implementing a certain protocol. +/// +/// # Errors +/// +/// * [`Status::NOT_FOUND`]: no matching handles. +#[cfg(feature = "alloc")] +pub fn find_handles() -> Result> { + // Search by protocol. + let search_type = SearchType::from_proto::

(); + + // Determine how much we need to allocate. + let num_handles = match locate_handle(search_type, &mut []) { + Err(err) => { + if err.status() == Status::BUFFER_TOO_SMALL { + err.data().expect("error data is missing") + } else { + return Err(err.to_err_without_payload()); + } + } + // This should never happen: if no handles match the search then a + // `NOT_FOUND` error should be returned. + Ok(_) => panic!("locate_handle should not return success with empty buffer"), + }; + + // Allocate a large enough buffer without pointless initialization. + let mut handles = Vec::with_capacity(num_handles); + + // Perform the search. + let num_handles = locate_handle(search_type, handles.spare_capacity_mut()) + .discard_errdata()? + .len(); + + // Mark the returned number of elements as initialized. + unsafe { + handles.set_len(num_handles); + } + + // Emit output, with warnings + Ok(handles) +} + /// Find an arbitrary handle that supports a particular [`Protocol`]. Returns /// [`NOT_FOUND`] if no handles support the protocol. ///