Protocol safety improvements (#460)

* Add BootServices::get_handle_for_protocol

This is a convenience method to get any arbitrary handle that supports a
particular `Protocol`.

* Use open_protocol in shim-lock test

* Use open_protocol in multiprocessor test

* Use open_protocol in device path test

* Use open_protocol in pointer test

* Use open_protocol in graphics test

* Use open_protocol in file system test

* Use open_protocol in the serial device test

* Simplify event callback with context test

Rather than opening the Output protocol, which isn't important for this
test, just write through the context pointer and assert that the
expected data was written.

* Mark handle_protocol as unsafe

This function is already marked deprecated, mark it unsafe as well and
update the documentation to describe why.

* Deprecate BootServices::locate_protocol and mark it unsafe

This method has the same problems as `handle_protocol`; it does not mark
the handle and protocol as in use. Calls to `locate_protocol` can be
replaced by calling `get_handle_for_protocol` and `open_protocol`.

#359

Author: nicholasbishop

Diff for: 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
 

Diff for: 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
         })

Diff for: 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) {

Diff for: 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);

Diff for: 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;

Diff for: 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)

Diff for: 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");

Diff for: 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!(

Diff for: 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 {

Diff for: 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) {

Diff for: 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;

Diff for: 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);