Merge pull request #1321 from nicholasbishop/bishop-locate-handle

boot: Add freestanding locate_handle and find_handles

Author: phip1611

uefi-test-runner/src/boot/mod.rs

@@ -15,12 +15,12 @@ pub fn test(st: &SystemTable<Boot>) {
     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::<Output>().unwrap();
+        assert_eq!(*handles, handles_vec);
     }
 }
 

uefi/src/boot.rs

@@ -4,18 +4,23 @@
 
 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};
 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<P: ProtocolPointer + ?Sized>(
     }
 }
 
+/// 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<Handle>],
+) -> Result<&'buf [Handle], Option<usize>> {
+    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::<Handle>();
+    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::<Handle>();
+
+    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<HandleBuffer> {
         })
 }
 
+/// Returns all the handles implementing a certain protocol.
+///
+/// # Errors
+///
+/// * [`Status::NOT_FOUND`]: no matching handles.
+#[cfg(feature = "alloc")]
+pub fn find_handles<P: ProtocolPointer + ?Sized>() -> Result<Vec<Handle>> {
+    // Search by protocol.
+    let search_type = SearchType::from_proto::<P>();
+
+    // 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.
 ///