diff --git a/uefi-test-runner/src/boot/mod.rs b/uefi-test-runner/src/boot/mod.rs index 74e817afe..e23d08118 100644 --- a/uefi-test-runner/src/boot/mod.rs +++ b/uefi-test-runner/src/boot/mod.rs @@ -5,7 +5,7 @@ use uefi::proto::device_path::media::FilePath; use uefi::proto::device_path::{DevicePath, LoadedImageDevicePath}; use uefi::table::boot::{BootServices, LoadImageSource, SearchType}; use uefi::table::{Boot, SystemTable}; -use uefi::{CString16, Identify}; +use uefi::{boot, CString16, Identify}; mod memory; mod misc; @@ -28,6 +28,12 @@ fn test_locate_handle_buffer(bt: &BootServices) { .locate_handle_buffer(SearchType::AllHandles) .expect("Failed to locate handle buffer"); assert!(!handles.is_empty(), "Could not find any handles"); + + // Compare with freestanding version. + assert_eq!( + *handles, + *boot::locate_handle_buffer(SearchType::AllHandles).unwrap() + ); } { @@ -39,6 +45,12 @@ fn test_locate_handle_buffer(bt: &BootServices) { !handles.is_empty(), "Could not find any OUTPUT protocol handles" ); + + // Compare with freestanding version. + assert_eq!( + *handles, + *boot::locate_handle_buffer(SearchType::ByProtocol(&Output::GUID)).unwrap() + ); } } diff --git a/uefi/src/boot.rs b/uefi/src/boot.rs index 6bb41bced..1fdb0e3fd 100644 --- a/uefi/src/boot.rs +++ b/uefi/src/boot.rs @@ -4,14 +4,16 @@ use crate::data_types::PhysicalAddress; use core::ffi::c_void; +use core::ops::Deref; use core::ptr::{self, NonNull}; +use core::slice; use core::sync::atomic::{AtomicPtr, Ordering}; use uefi::{table, Handle, Result, StatusExt}; #[cfg(doc)] use uefi::Status; -pub use uefi::table::boot::AllocateType; +pub use uefi::table::boot::{AllocateType, SearchType}; pub use uefi_raw::table::boot::MemoryType; /// Global image handle. This is only set by [`set_image_handle`], and it is @@ -128,3 +130,56 @@ pub unsafe fn free_pool(ptr: NonNull) -> Result { unsafe { (bt.free_pool)(ptr.as_ptr()) }.to_result() } + +/// Returns an array of handles that support the requested protocol in a +/// pool-allocated buffer. +/// +/// See [`SearchType`] for details of the available search operations. +/// +/// # Errors +/// +/// * [`Status::NOT_FOUND`]: no matching handles. +/// * [`Status::OUT_OF_RESOURCES`]: out of memory. +pub fn locate_handle_buffer(search_ty: SearchType) -> Result { + let bt = boot_services_raw_panicking(); + let bt = unsafe { bt.as_ref() }; + + 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 _, ptr::null()), + }; + + let mut num_handles: usize = 0; + let mut buffer: *mut uefi_raw::Handle = ptr::null_mut(); + unsafe { (bt.locate_handle_buffer)(ty, guid, key, &mut num_handles, &mut buffer) } + .to_result_with_val(|| HandleBuffer { + count: num_handles, + buffer: NonNull::new(buffer.cast()) + .expect("locate_handle_buffer must not return a null pointer"), + }) +} + +/// A buffer returned by [`locate_handle_buffer`] that contains an array of +/// [`Handle`]s that support the requested protocol. +#[derive(Debug, Eq, PartialEq)] +pub struct HandleBuffer { + count: usize, + buffer: NonNull, +} + +impl Drop for HandleBuffer { + fn drop(&mut self) { + let _ = unsafe { free_pool(self.buffer.cast::()) }; + } +} + +impl Deref for HandleBuffer { + type Target = [Handle]; + + fn deref(&self) -> &Self::Target { + unsafe { slice::from_raw_parts(self.buffer.as_ptr(), self.count) } + } +} diff --git a/uefi/src/table/boot.rs b/uefi/src/table/boot.rs index 18122df28..9baef842b 100644 --- a/uefi/src/table/boot.rs +++ b/uefi/src/table/boot.rs @@ -1716,4 +1716,4 @@ impl<'a> HandleBuffer<'a> { /// with [`BootServices::locate_handle`] via [`SearchType::ByRegisterNotify`]. #[derive(Debug, Clone, Copy)] #[repr(transparent)] -pub struct ProtocolSearchKey(NonNull); +pub struct ProtocolSearchKey(pub(crate) NonNull);