Add support for locate_protocol

This fixes #79 and allows users to avoid allocating when they just
want to find a single implementer of a protocol.

Author: josephlr

src/table/boot.rs

@@ -103,7 +103,11 @@ pub struct BootServices {
     // Library services
     protocols_per_handle: usize,
     locate_handle_buffer: usize,
-    locate_protocol: usize,
+    locate_protocol: extern "win64" fn(
+        proto: &Guid,
+        registration: *mut c_void,
+        out_proto: &mut *mut c_void,
+    ) -> Status,
     install_multiple_protocol_interfaces: usize,
     uninstall_multiple_protocol_interfaces: usize,
 
@@ -446,6 +450,17 @@ impl BootServices {
         unsafe { (self.set_watchdog_timer)(timeout, watchdog_code, data_len, data) }.into()
     }
 
+    /// Returns a protocol implementation, if present on the system.
+    ///
+    /// The caveats of `BootServices::handle_protocol()` also apply here.
+    pub fn locate_protocol<P: Protocol>(&self) -> Result<&UnsafeCell<P>> {
+        let mut ptr = ptr::null_mut();
+        (self.locate_protocol)(&P::GUID, ptr::null_mut(), &mut ptr).into_with_val(|| {
+            let ptr = ptr as *mut P as *mut UnsafeCell<P>;
+            unsafe { &*ptr }
+        })
+    }
+
     /// Copies memory from source to destination. The buffers can overlap.
     ///
     /// This function is unsafe as it can be used to violate most safety

uefi-exts/src/boot.rs

@@ -1,4 +1,3 @@
-use core::cell::UnsafeCell;
 use uefi::proto::Protocol;
 use uefi::table::boot::{BootServices, SearchType};
 use uefi::{Handle, Result};
@@ -9,11 +8,6 @@ use crate::alloc::vec::Vec;
 pub trait BootServicesExt {
     /// Returns all the handles implementing a certain protocol.
     fn find_handles<P: Protocol>(&self) -> Result<Vec<Handle>>;
-
-    /// Returns a protocol implementation, if present on the system.
-    ///
-    /// The caveats of `BootServices::handle_protocol()` also apply here.
-    fn find_protocol<P: Protocol>(&self) -> Result<&UnsafeCell<P>>;
 }
 
 impl BootServicesExt for BootServices {
@@ -44,23 +38,4 @@ impl BootServicesExt for BootServices {
             .into_with_val(|| buffer)
             .map(|completion| completion.with_status(status2))
     }
-
-    fn find_protocol<P: Protocol>(&self) -> Result<&UnsafeCell<P>> {
-        // Retrieve all handles implementing this.
-        let (status1, handles) = self.find_handles::<P>()?.split();
-
-        // There should be at least one, otherwise find_handles would have
-        // aborted with a NOT_FOUND error.
-        let handle = *handles.first().unwrap();
-
-        // Similarly, if the search is implemented properly, trying to open
-        // the first output handle should always succeed
-        // FIXME: Consider using the EFI 1.1 LocateProtocol API instead
-        let (status2, protocol) = self.handle_protocol::<P>(handle)?.split();
-
-        // Emit output, with warnings
-        status1
-            .into_with_val(|| protocol)
-            .map(|completion| completion.with_status(status2))
-    }
 }

uefi-test-runner/src/main.rs

@@ -14,7 +14,6 @@ use core::mem;
 use uefi::prelude::*;
 use uefi::proto::console::serial::Serial;
 use uefi::table::boot::MemoryDescriptor;
-use uefi_exts::BootServicesExt;
 
 mod boot;
 mod proto;
@@ -67,7 +66,7 @@ fn check_screenshot(bt: &BootServices, name: &str) {
     if cfg!(feature = "qemu") {
         // Access the serial port (in a QEMU environment, it should always be there)
         let serial = bt
-            .find_protocol::<Serial>()
+            .locate_protocol::<Serial>()
             .expect_success("Could not find serial port");
         let serial = unsafe { &mut *serial.get() };
 

uefi-test-runner/src/proto/console/gop.rs

@@ -1,11 +1,10 @@
 use uefi::prelude::*;
 use uefi::proto::console::gop::{BltOp, BltPixel, FrameBuffer, GraphicsOutput, PixelFormat};
 use uefi::table::boot::BootServices;
-use uefi_exts::BootServicesExt;
 
 pub fn test(bt: &BootServices) {
     info!("Running graphics output protocol test");
-    if let Ok(gop) = bt.find_protocol::<GraphicsOutput>() {
+    if let Ok(gop) = bt.locate_protocol::<GraphicsOutput>() {
         let gop = gop.expect("Warnings encountered while opening GOP");
         let gop = unsafe { &mut *gop.get() };
 

uefi-test-runner/src/proto/console/pointer.rs

@@ -2,11 +2,9 @@ use uefi::prelude::*;
 use uefi::proto::console::pointer::Pointer;
 use uefi::table::boot::BootServices;
 
-use uefi_exts::BootServicesExt;
-
 pub fn test(bt: &BootServices) {
     info!("Running pointer protocol test");
-    if let Ok(pointer) = bt.find_protocol::<Pointer>() {
+    if let Ok(pointer) = bt.locate_protocol::<Pointer>() {
         let pointer = pointer.expect("Warnings encountered while opening pointer protocol");
         let pointer = unsafe { &mut *pointer.get() };
 

uefi-test-runner/src/proto/console/serial.rs

@@ -1,11 +1,10 @@
 use uefi::prelude::*;
 use uefi::proto::console::serial::{ControlBits, Serial};
 use uefi::table::boot::BootServices;
-use uefi_exts::BootServicesExt;
 
 pub fn test(bt: &BootServices) {
     info!("Running serial protocol test");
-    if let Ok(serial) = bt.find_protocol::<Serial>() {
+    if let Ok(serial) = bt.locate_protocol::<Serial>() {
         let serial = serial.expect("Warnings encountered while opening serial protocol");
         let serial = unsafe { &mut *serial.get() };
 

uefi-test-runner/src/proto/debug.rs

@@ -1,11 +1,9 @@
 use uefi::proto::debug::DebugSupport;
 use uefi::table::boot::BootServices;
 
-use uefi_exts::BootServicesExt;
-
 pub fn test(bt: &BootServices) {
     info!("Running UEFI debug connection protocol test");
-    if let Ok(debug_support) = bt.find_protocol::<DebugSupport>() {
+    if let Ok(debug_support) = bt.locate_protocol::<DebugSupport>() {
         let debug_support =
             debug_support.expect("Warnings encountered while opening debug support protocol");
         let debug_support = unsafe { &mut *debug_support.get() };