Merge pull request #1269 from nicholasbishop/bishop-handle-buf

Add `boot::locate_handle_buffer`

Author: phip1611

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()
+        );
     }
 }
 

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<u8>) -> 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<HandleBuffer> {
+    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<Handle>,
+}
+
+impl Drop for HandleBuffer {
+    fn drop(&mut self) {
+        let _ = unsafe { free_pool(self.buffer.cast::<u8>()) };
+    }
+}
+
+impl Deref for HandleBuffer {
+    type Target = [Handle];
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { slice::from_raw_parts(self.buffer.as_ptr(), self.count) }
+    }
+}

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<c_void>);
+pub struct ProtocolSearchKey(pub(crate) NonNull<c_void>);