Protocol safety improvements

CHANGELOG.md

@@ -10,6 +10,15 @@
   which is now marked as deprecated.
 - Implemented `core::fmt::Write` for the `Serial` protocol.
 - Added the `MemoryProtection` protocol.
+- Added `BootServices::get_handle_for_protocol`.
+
+### Changed
+
+- Marked `BootServices::handle_protocol` as `unsafe`. (This method is
+  also deprecated -- use `open_protocol` instead.)
+- Deprecated `BootServices::locate_protocol` and marked it `unsafe`. Use
+  `BootServices::get_handle_for_protocol` and
+  `BootServices::open_protocol` instead.
 
 ### Fixed
 

src/table/boot.rs

@@ -576,14 +576,21 @@ impl BootServices {
     /// protections must be implemented by user-level code, for example via a
     /// global `HashSet`.
     ///
+    /// # Safety
+    ///
+    /// This method is unsafe because the handle database is not
+    /// notified that the handle and protocol are in use; there is no
+    /// guarantee that they will remain valid for the duration of their
+    /// use. Use [`open_protocol`] instead.
+    ///
     /// [`open_protocol`]: BootServices::open_protocol
     #[deprecated(note = "it is recommended to use `open_protocol` instead")]
-    pub fn handle_protocol<P: ProtocolPointer + ?Sized>(
+    pub unsafe fn handle_protocol<P: ProtocolPointer + ?Sized>(
         &self,
         handle: Handle,
     ) -> Result<&UnsafeCell<P>> {
         let mut ptr = ptr::null_mut();
-        (self.handle_protocol)(handle, &P::GUID, &mut ptr).into_with_val(|| unsafe {
+        (self.handle_protocol)(handle, &P::GUID, &mut ptr).into_with_val(|| {
             let ptr = P::mut_ptr_from_ffi(ptr) as *const UnsafeCell<P>;
             &*ptr
         })
@@ -649,6 +656,57 @@ impl BootServices {
         }
     }
 
+    /// Find an arbitrary handle that supports a particular
+    /// [`Protocol`]. Returns [`NOT_FOUND`] if no handles support the
+    /// protocol.
+    ///
+    /// This method is a convenient wrapper around
+    /// [`BootServices::locate_handle_buffer`] for getting just one
+    /// handle. This is useful when you don't care which handle the
+    /// protocol is opened on. For example, [`DevicePathToText`] isn't
+    /// tied to a particular device, so only a single handle is expected
+    /// to exist.
+    ///
+    /// [`NOT_FOUND`]: Status::NOT_FOUND
+    /// [`DevicePathToText`]: uefi::proto::device_path::text::DevicePathToText
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use uefi::proto::device_path::text::DevicePathToText;
+    /// use uefi::table::boot::{BootServices, OpenProtocolAttributes, OpenProtocolParams};
+    /// use uefi::Handle;
+    /// # use uefi::Result;
+    ///
+    /// # fn get_fake_val<T>() -> T { todo!() }
+    /// # fn test() -> Result {
+    /// # let boot_services: &BootServices = get_fake_val();
+    /// # let image_handle: Handle = get_fake_val();
+    /// let handle = boot_services.get_handle_for_protocol::<DevicePathToText>()?;
+    /// let device_path_to_text = boot_services.open_protocol::<DevicePathToText>(
+    ///     OpenProtocolParams {
+    ///         handle,
+    ///         agent: image_handle,
+    ///         controller: None,
+    ///     },
+    ///     OpenProtocolAttributes::Exclusive,
+    /// )?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_handle_for_protocol<P: Protocol>(&self) -> Result<Handle> {
+        // Delegate to a non-generic function to potentially reduce code size.
+        self.get_handle_for_protocol_impl(&P::GUID)
+    }
+
+    fn get_handle_for_protocol_impl(&self, guid: &Guid) -> Result<Handle> {
+        self.locate_handle_buffer(SearchType::ByProtocol(guid))?
+            .handles()
+            .first()
+            .cloned()
+            .ok_or_else(|| Status::NOT_FOUND.into())
+    }
+
     /// Load an EFI image into memory and return a [`Handle`] to the image.
     ///
     /// There are two ways to load the image: by copying raw image data
@@ -961,9 +1019,18 @@ impl BootServices {
     /// Returns a protocol implementation, if present on the system.
     ///
     /// The caveats of `BootServices::handle_protocol()` also apply here.
-    pub fn locate_protocol<P: ProtocolPointer + ?Sized>(&self) -> Result<&UnsafeCell<P>> {
+    ///
+    /// # Safety
+    ///
+    /// This method is unsafe because the handle database is not
+    /// notified that the handle and protocol are in use; there is no
+    /// guarantee that they will remain valid for the duration of their
+    /// use. Use [`BootServices::get_handle_for_protocol`] and
+    /// [`BootServices::open_protocol`] instead.
+    #[deprecated(note = "it is recommended to use `open_protocol` instead")]
+    pub unsafe fn locate_protocol<P: ProtocolPointer + ?Sized>(&self) -> Result<&UnsafeCell<P>> {
         let mut ptr = ptr::null_mut();
-        (self.locate_protocol)(&P::GUID, ptr::null_mut(), &mut ptr).into_with_val(|| unsafe {
+        (self.locate_protocol)(&P::GUID, ptr::null_mut(), &mut ptr).into_with_val(|| {
             let ptr = P::mut_ptr_from_ffi(ptr) as *const UnsafeCell<P>;
             &*ptr
         })

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

@@ -1,7 +1,6 @@
 use core::ffi::c_void;
 use core::ptr::NonNull;
 
-use uefi::proto::console::text::Output;
 use uefi::table::boot::{BootServices, EventType, TimerTrigger, Tpl};
 use uefi::Event;
 
@@ -37,28 +36,35 @@ fn test_event_callback(bt: &BootServices) {
 }
 
 fn test_callback_with_ctx(bt: &BootServices) {
+    let mut data = 123u32;
+
     extern "efiapi" fn callback(_event: Event, ctx: Option<NonNull<c_void>>) {
         info!("Inside the event callback with context");
+        // Safety: this callback is run within the parent function's
+        // scope, so the context pointer is still valid.
         unsafe {
-            let ctx = &mut *(ctx.unwrap().as_ptr() as *mut Output);
-            // Clear the screen as a quick test that we successfully passed context
-            ctx.clear().expect("Failed to clear screen");
+            let ctx = ctx.unwrap().as_ptr().cast::<u32>();
+            *ctx = 456;
         }
     }
 
-    let ctx = unsafe { &mut *(bt.locate_protocol::<Output>().unwrap().get()) };
+    let ctx: *mut u32 = &mut data;
+    let ctx = NonNull::new(ctx.cast::<c_void>()).unwrap();
 
     let event = unsafe {
         bt.create_event(
             EventType::NOTIFY_WAIT,
             Tpl::CALLBACK,
             Some(callback),
-            Some(NonNull::new_unchecked(ctx as *mut _ as *mut c_void)),
+            Some(ctx),
         )
         .expect("Failed to create event with context")
     };
 
     bt.check_event(event).expect("Failed to check event");
+
+    // Check that `data` was updated inside the event callback.
+    assert_eq!(data, 456);
 }
 
 fn test_watchdog(bt: &BootServices) {

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

@@ -1,11 +1,23 @@
 use uefi::prelude::*;
 use uefi::proto::console::gop::{BltOp, BltPixel, FrameBuffer, GraphicsOutput, PixelFormat};
-use uefi::table::boot::BootServices;
+use uefi::table::boot::{BootServices, OpenProtocolAttributes, OpenProtocolParams};
 
 pub fn test(image: Handle, bt: &BootServices) {
     info!("Running graphics output protocol test");
-    if let Ok(gop) = bt.locate_protocol::<GraphicsOutput>() {
-        let gop = unsafe { &mut *gop.get() };
+    if let Ok(handle) = bt.get_handle_for_protocol::<GraphicsOutput>() {
+        let gop = &mut bt
+            .open_protocol::<GraphicsOutput>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                // For this test, don't open in exclusive mode. That
+                // would break the connection between stdout and the
+                // video console.
+                OpenProtocolAttributes::GetProtocol,
+            )
+            .expect("failed to open Graphics Output Protocol");
 
         set_graphics_mode(gop);
         fill_color(gop);

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

@@ -6,9 +6,9 @@ pub fn test(image: Handle, st: &mut SystemTable<Boot>) {
     stdout::test(st.stdout());
 
     let bt = st.boot_services();
-    serial::test(bt);
+    serial::test(image, bt);
     gop::test(image, bt);
-    pointer::test(bt);
+    pointer::test(image, bt);
 }
 
 mod gop;

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

@@ -1,10 +1,20 @@
 use uefi::proto::console::pointer::Pointer;
-use uefi::table::boot::BootServices;
+use uefi::table::boot::{BootServices, OpenProtocolAttributes, OpenProtocolParams};
+use uefi::Handle;
 
-pub fn test(bt: &BootServices) {
+pub fn test(image: Handle, bt: &BootServices) {
     info!("Running pointer protocol test");
-    if let Ok(pointer) = bt.locate_protocol::<Pointer>() {
-        let pointer = unsafe { &mut *pointer.get() };
+    if let Ok(handle) = bt.get_handle_for_protocol::<Pointer>() {
+        let mut pointer = bt
+            .open_protocol::<Pointer>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                OpenProtocolAttributes::Exclusive,
+            )
+            .expect("failed to open pointer protocol");
 
         pointer
             .reset(false)

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

@@ -1,16 +1,28 @@
 use uefi::proto::console::serial::{ControlBits, Serial};
-use uefi::table::boot::BootServices;
+use uefi::table::boot::{BootServices, OpenProtocolAttributes, OpenProtocolParams};
+use uefi::Handle;
 
-pub fn test(bt: &BootServices) {
+pub fn test(image: Handle, bt: &BootServices) {
     info!("Running serial protocol test");
-    if let Ok(serial) = bt.locate_protocol::<Serial>() {
+    if let Ok(handle) = bt.get_handle_for_protocol::<Serial>() {
+        let mut serial = bt
+            .open_protocol::<Serial>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                // For this test, don't open in exclusive mode. That
+                // would break the connection between stdout and the
+                // serial device.
+                OpenProtocolAttributes::GetProtocol,
+            )
+            .expect("failed to open serial protocol");
         // BUG: there are multiple failures in the serial tests on AArch64
         if cfg!(target_arch = "aarch64") {
             return;
         }
 
-        let serial = unsafe { &mut *serial.get() };
-
         let old_ctrl_bits = serial
             .get_control_bits()
             .expect("Failed to get device control bits");

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

@@ -29,14 +29,30 @@ pub fn test(image: Handle, bt: &BootServices) {
         .expect("Failed to open DevicePath protocol");
 
     let device_path_to_text = bt
-        .locate_protocol::<DevicePathToText>()
+        .open_protocol::<DevicePathToText>(
+            OpenProtocolParams {
+                handle: bt
+                    .get_handle_for_protocol::<DevicePathToText>()
+                    .expect("Failed to get DevicePathToText handle"),
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::Exclusive,
+        )
         .expect("Failed to open DevicePathToText protocol");
-    let device_path_to_text = unsafe { &*device_path_to_text.get() };
 
     let device_path_from_text = bt
-        .locate_protocol::<DevicePathFromText>()
+        .open_protocol::<DevicePathFromText>(
+            OpenProtocolParams {
+                handle: bt
+                    .get_handle_for_protocol::<DevicePathFromText>()
+                    .expect("Failed to get DevicePathFromText handle"),
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::Exclusive,
+        )
         .expect("Failed to open DevicePathFromText protocol");
-    let device_path_from_text = unsafe { &*device_path_from_text.get() };
 
     for path in device_path.node_iter() {
         info!(

uefi-test-runner/src/proto/media/mod.rs

@@ -27,8 +27,18 @@ fn test_file_system_info(directory: &mut Directory) {
 pub fn test(image: Handle, bt: &BootServices) {
     info!("Testing Media Access protocols");
 
-    if let Ok(sfs) = bt.locate_protocol::<SimpleFileSystem>() {
-        let sfs = unsafe { &mut *sfs.get() };
+    if let Ok(handle) = bt.get_handle_for_protocol::<SimpleFileSystem>() {
+        let mut sfs = bt
+            .open_protocol::<SimpleFileSystem>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                OpenProtocolAttributes::Exclusive,
+            )
+            .expect("failed to open SimpleFileSystem protocol");
+
         let mut directory = sfs.open_volume().unwrap();
         let mut buffer = vec![0; 128];
         loop {

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

@@ -17,7 +17,7 @@ pub fn test(image: Handle, st: &mut SystemTable<Boot>) {
     loaded_image::test(image, bt);
     media::test(image, bt);
     network::test(image, bt);
-    pi::test(bt);
+    pi::test(image, bt);
     rng::test(image, bt);
 
     #[cfg(any(
@@ -26,7 +26,7 @@ pub fn test(image: Handle, st: &mut SystemTable<Boot>) {
         target_arch = "arm",
         target_arch = "aarch64"
     ))]
-    shim::test(bt);
+    shim::test(image, bt);
 }
 
 fn find_protocol(bt: &BootServices) {

uefi-test-runner/src/proto/pi/mod.rs

@@ -1,9 +1,9 @@
 use uefi::prelude::*;
 
-pub fn test(bt: &BootServices) {
+pub fn test(image: Handle, bt: &BootServices) {
     info!("Testing Platform Initialization protocols");
 
-    mp::test(bt);
+    mp::test(image, bt);
 }
 
 mod mp;

uefi-test-runner/src/proto/pi/mp.rs

@@ -2,21 +2,30 @@ use core::ffi::c_void;
 use core::sync::atomic::{AtomicUsize, Ordering};
 use core::time::Duration;
 use uefi::proto::pi::mp::MpServices;
-use uefi::table::boot::BootServices;
-use uefi::Status;
+use uefi::table::boot::{BootServices, OpenProtocolAttributes, OpenProtocolParams};
+use uefi::{Handle, Status};
 
 /// Number of cores qemu is configured to have
 const NUM_CPUS: usize = 4;
 
-pub fn test(bt: &BootServices) {
+pub fn test(image: Handle, bt: &BootServices) {
     // These tests break CI. See #103.
     if cfg!(feature = "ci") {
         return;
     }
 
     info!("Running UEFI multi-processor services protocol test");
-    if let Ok(mp_support) = bt.locate_protocol::<MpServices>() {
-        let mp_support = unsafe { &mut *mp_support.get() };
+    if let Ok(handle) = bt.get_handle_for_protocol::<MpServices>() {
+        let mp_support = &bt
+            .open_protocol::<MpServices>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                OpenProtocolAttributes::Exclusive,
+            )
+            .expect("failed to open multi-processor services protocol");
 
         test_get_number_of_processors(mp_support);
         test_get_processor_info(mp_support);