uefi: add boot::[un_]install_multiple_protocol_interface()

Author: phip1611

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

@@ -1,14 +1,19 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
+use alloc::boxed::Box;
+use alloc::vec::Vec;
 use core::ffi::c_void;
 use core::ptr::{self, NonNull};
 
 use uefi::boot::{
     EventType, OpenProtocolAttributes, OpenProtocolParams, SearchType, TimerTrigger, Tpl,
 };
 use uefi::mem::memory_map::MemoryType;
-use uefi::proto::unsafe_protocol;
-use uefi::{Event, Guid, Identify, boot, guid, system};
+use uefi::proto::device_path::build::DevicePathBuilder;
+use uefi::proto::device_path::{DevicePath, build};
+use uefi::proto::{ProtocolPointer, unsafe_protocol};
+use uefi::{Event, Guid, Identify, ResultExt, boot, cstr16, guid, system};
+use uefi_raw::Status;
 
 pub fn test() {
     test_tpl();
@@ -25,6 +30,8 @@ pub fn test() {
     test_install_protocol_interface();
     test_reinstall_protocol_interface();
     test_uninstall_protocol_interface();
+    test_install_multiple_protocol_interface();
+    test_uninstall_multiple_protocol_interface();
     test_install_configuration_table();
     info!("Testing crc32...");
     test_calculate_crc32();
@@ -210,13 +217,130 @@ fn test_uninstall_protocol_interface() {
             &mut *sp
         };
 
-        boot::uninstall_protocol_interface(handle, &TestProtocol::GUID, interface_ptr.cast())
+        boot::uninstall_protocol_interface::<TestProtocol>(handle, interface_ptr.cast())
             .expect("Failed to uninstall protocol interface");
 
         boot::free_pool(NonNull::new(interface_ptr.cast()).unwrap()).unwrap();
     }
 }
 
+fn test_install_multiple_protocol_interface() {
+    info!("Installing multiple test protocols");
+
+    let alloc: *mut TestProtocol =
+        boot::allocate_pool(MemoryType::BOOT_SERVICES_DATA, size_of::<TestProtocol>())
+            .unwrap()
+            .cast()
+            .as_ptr();
+    unsafe { alloc.write(TestProtocol { data: 123 }) };
+
+    let dvp = {
+        let mut vec = Vec::new();
+        DevicePathBuilder::with_vec(&mut vec)
+            .push(&build::media::FilePath {
+                path_name: cstr16!("foobar"),
+            })
+            .unwrap()
+            .finalize()
+            .unwrap()
+            .to_boxed()
+    };
+    // Memory must stay valid as long as handle with interfaces lives:
+    // => so we leak the memory but will free it in the uninstall hook again.
+    let dvp = Box::leak(dvp);
+
+    let handle = unsafe {
+        boot::install_multiple_protocol_interface(
+            None,
+            &[
+                (&TestProtocol::GUID, alloc.cast()),
+                (&DevicePath::GUID, dvp.as_ffi_ptr().cast()),
+            ],
+        )
+        .expect("Failed to install protocol interface")
+    };
+
+    // Test we indeed installed the protocols.
+    {
+        assert_eq!(
+            boot::test_protocol::<DevicePath>(OpenProtocolParams {
+                handle,
+                agent: boot::image_handle(),
+                controller: None,
+            }),
+            Ok(true)
+        );
+    }
+
+    // Test that installing the device path protocol multiple times results in
+    // EFI_ALREADY_STARTED
+    {
+        let res = unsafe {
+            boot::install_multiple_protocol_interface(
+                Some(handle),
+                &[(&DevicePath::GUID, dvp.as_ffi_ptr().cast())],
+            )
+        };
+        assert_eq!(res.status(), Status::ALREADY_STARTED);
+    }
+}
+
+fn test_uninstall_multiple_protocol_interface() {
+    info!("Uninstalling multiple test protocols");
+
+    let handles = boot::locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
+        .expect("Failed to find protocol after it was installed");
+    let handle = *handles.first().unwrap();
+
+    let interface_test_protocol: *mut TestProtocol = unsafe {
+        let mut sp = boot::open_protocol::<TestProtocol>(
+            OpenProtocolParams {
+                handle,
+                agent: boot::image_handle(),
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .unwrap();
+        assert_eq!(sp.data, 123);
+        &mut *sp
+    };
+
+    let interface_dvp: *mut DevicePath = unsafe {
+        let mut sp = boot::open_protocol::<DevicePath>(
+            OpenProtocolParams {
+                handle,
+                agent: boot::image_handle(),
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .unwrap();
+        &mut *sp
+    };
+
+    unsafe {
+        boot::uninstall_multiple_protocol_interface(
+            handle,
+            &[
+                (&TestProtocol::GUID, interface_test_protocol.cast()),
+                (&DevicePath::GUID, interface_dvp.cast()),
+            ],
+        )
+    }
+    .expect("should uninstall multiple protocols");
+
+    let dvp = unsafe {
+        DevicePath::mut_ptr_from_ffi(interface_dvp.cast())
+            .as_mut()
+            .expect("should be valid device path")
+    };
+
+    // Reconstruct the Rust box to ensure that the object is properly freed in
+    // the Rust global allocator
+    let _ = unsafe { Box::from_raw(dvp) };
+}
+
 fn test_install_configuration_table() {
     // Get the current number of entries.
     let initial_table_count = system::with_config_table(|t| t.len());

uefi/CHANGELOG.md

@@ -10,6 +10,9 @@
 - Added `boot::test_protocol_by_guid()`
 - Added `boot::register_protocol_notify_by_guid()`
 - Added `boot::[re_,un_]install_protocol_interface_by_guid()` functions.
+- Added `boot::[un_]install_multiple_protocol_interface`. Currently, this
+  replicates the functionality of the EDK2 implementation rather than using it
+  due to Rusts limited support for variadic arguments.
 
 ## Changed
 - Changed ordering of `proto::pci::PciIoAddress` to (bus -> dev -> fun -> reg -> ext_reg).

uefi/src/boot.rs

@@ -37,17 +37,20 @@ use crate::proto::{BootPolicy, Protocol, ProtocolPointer};
 use crate::runtime::{self, ResetType};
 use crate::table::Revision;
 use crate::util::opt_nonnull_to_ptr;
-use crate::{Char16, Error, Event, Guid, Handle, Result, Status, StatusExt, table};
+use crate::{
+    Char16, Error, Event, Guid, Handle, Identify, Result, ResultExt, Status, StatusExt, table,
+};
+#[cfg(feature = "alloc")]
+use alloc::vec::Vec;
 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::time::Duration;
 use core::{mem, slice};
+use log::error;
 use uefi_raw::table::boot::{AllocateType as RawAllocateType, InterfaceType, TimerDelay};
-#[cfg(feature = "alloc")]
-use {alloc::vec::Vec, uefi::ResultExt};
 
 /// Global image handle. This is only set by [`set_image_handle`], and it is
 /// only read by [`image_handle`].
@@ -713,7 +716,12 @@ pub fn disconnect_controller(
     .to_result_with_err(|_| ())
 }
 
-/// Installs a protocol interface on a device handle.
+/// Installs a protocol interface on a device handle. If no handle is
+/// specified, a new handle will be allocated and returned.
+///
+/// It is recommended to use [`install_multiple_protocol_interface`] when you
+/// plan to install multiple protocols, as it performs more error checking
+/// and cleanup under the hood.
 ///
 /// When a protocol interface is installed, firmware will call all functions
 /// that have registered to wait for that interface to be installed.
@@ -736,7 +744,7 @@ pub fn disconnect_controller(
 pub unsafe fn install_protocol_interface<P: ProtocolPointer + ?Sized>(
     handle: Option<Handle>,
     interface: *const c_void,
-) -> Result<Handle> {
+) -> Result<Handle /* new if input was None */> {
     unsafe { install_protocol_interface_by_guid(handle, &P::GUID, interface) }
 }
 
@@ -750,7 +758,7 @@ pub unsafe fn install_protocol_interface_by_guid(
     handle: Option<Handle>,
     protocol: &Guid,
     interface: *const c_void,
-) -> Result<Handle> {
+) -> Result<Handle /* new if Input was None */> {
     let bt = boot_services_raw_panicking();
     let bt = unsafe { bt.as_ref() };
 
@@ -852,6 +860,168 @@ pub unsafe fn uninstall_protocol_interface_by_guid(
     unsafe { (bt.uninstall_protocol_interface)(handle.as_ptr(), protocol, interface).to_result() }
 }
 
+/// Installs multiple protocol interfaces for a given handle at once, and
+/// reverts all operations if a single operation fails. If no handle is
+/// specified, a new handle will be allocated and returned.
+///
+/// When a protocol interface is installed, firmware will call all functions
+/// that have registered to wait for that interface to be installed.
+///
+/// As Rust is not having proper C variadic support, this function emulates the
+/// behavior of the `CoreInstallMultipleProtocolInterfaces` function from
+/// `edk2`. Effectively, the behavior is the same, but it doesn't use the
+/// corresponding boot service under the hood.
+///
+/// # Arguments
+///
+/// - `handle`: Either `None` to allocate a new handle or an existing handle.
+/// - `pairs`: Pairs of the [`Guid`] of the [`Protocol`] to install and the
+///   protocol implementation. The memory backing the implementation
+///   **must live as long as the handle!**. Callers need to ensure a matching
+///   lifetime!
+///
+/// # Safety
+///
+/// The caller is responsible for ensuring that they pass a valid `Guid` for `protocol`.
+///
+/// # Errors
+///
+/// * [`Status::OUT_OF_RESOURCES`]: failed to allocate a new handle.
+/// * [`Status::INVALID_PARAMETER`]: this protocol is already installed on the handle.
+/// * [`Status::ALREADY_STARTED`]: A Device Path Protocol instance was passed in that is already present in the handle database.
+pub unsafe fn install_multiple_protocol_interface(
+    mut handle: Option<Handle>,
+    pairs: &[(&Guid, *const c_void)],
+) -> Result<Handle /* new if input was None */> {
+    // TODO once Rust has sensible variadic argument support, we should
+    // fallback to the correct boot service.
+
+    // Taken from edk2 source.
+    const TPL_NOTIFY: Tpl = Tpl(16);
+    let tpl = TPL_NOTIFY;
+    // SAFETY: We do not want our loop to be interrupted.
+    let _old_tpl = unsafe { raise_tpl(tpl) };
+
+    // Variables that are updated in the loop.
+    let mut installed_count = 0;
+    let mut status = Status::SUCCESS;
+
+    // try to install all interfaces and update `handle` if it is `None`
+    for (guid, interface) in pairs {
+        // prevent multiple installations of the device path protocol on the
+        // same handle:
+        if let Some(handle) = handle {
+            if *guid == &DevicePath::GUID
+                && test_protocol_by_guid(
+                    guid,
+                    OpenProtocolParams {
+                        handle,
+                        agent: image_handle(),
+                        controller: None,
+                    },
+                )?
+            {
+                status = Status::ALREADY_STARTED;
+                break;
+            }
+        }
+
+        let result = unsafe { install_protocol_interface_by_guid(handle, guid, *interface) };
+
+        match (result, handle, installed_count) {
+            (Ok(new_handle), None, 0) => {
+                handle = Some(new_handle);
+            }
+            (Ok(_handle), _, _) => {}
+            (Err(err), _, _) => {
+                error!("Failed to install protocol interface: {err}");
+                // next, we need to uninstall for all succeeded iterations
+                status = err.status();
+                break;
+            }
+        }
+
+        installed_count += 1;
+    }
+
+    if !status.is_success() {
+        // try to uninstall all that were just successfully installed
+        for (guid, interface) in pairs.iter().take(installed_count) {
+            let res =
+                unsafe { uninstall_protocol_interface_by_guid(handle.unwrap(), guid, *interface) };
+            if let Err(e) = res {
+                let handle_addr = &raw const *handle.as_ref().unwrap();
+                // We don't fail here, as this would break the contract of the
+                // function.
+                error!(
+                    "Failed to uninstall interface after failed multiple install attempt: handle={handle_addr:?}, guid={}, interface={:?}, error={e}",
+                    guid, interface
+                );
+            }
+        }
+
+        Err(status.into())
+    } else {
+        Ok(handle.unwrap())
+    }
+}
+
+/// Removes one or more protocol interfaces into the boot services environment.
+///
+/// If any errors are generated while the protocol interfaces are being
+/// uninstalled, then the protocols uninstalled prior to the error will be
+/// reinstalled.
+///
+/// # Safety
+///
+/// The caller is responsible for ensuring that there are no references to a protocol interface
+/// that has been removed. Some protocols may not be able to be removed as there is no information
+/// available regarding the references. This includes Console I/O, Block I/O, Disk I/o, and handles
+/// to device protocols.
+///
+/// # Errors
+///
+/// * [`Status::NOT_FOUND`]: the interface was not found on the handle.
+/// * [`Status::ACCESS_DENIED`]: the interface is still in use and cannot be uninstalled.
+pub unsafe fn uninstall_multiple_protocol_interface(
+    handle: Handle,
+    pairs: &[(&Guid, *const c_void)],
+) -> Result<()> {
+    let mut uninstalled_count = 0;
+    let mut status = Status::SUCCESS;
+
+    // try to install all interfaces and update `handle` if it is `None`
+    for (guid, interface) in pairs {
+        let result = unsafe { uninstall_protocol_interface_by_guid(handle, guid, *interface) };
+
+        if result.is_err() {
+            // next, we need to install for all succeeded iterations
+            status = result.status();
+            break;
+        }
+
+        uninstalled_count += 1;
+    }
+
+    if !status.is_success() {
+        // try to uninstall all failed ones
+        for (guid, interface) in pairs.iter().take(uninstalled_count) {
+            let res = unsafe { install_protocol_interface_by_guid(Some(handle), guid, *interface) };
+            if let Err(e) = res {
+                let handle_addr = &raw const handle;
+                // We don't fail here, as this would break the contract of the
+                // function.
+                error!(
+                    "Failed to install interface after failed multiple uninstall attempt: handle={handle_addr:?}, guid={}, interface={:?}, error={e}",
+                    guid, interface
+                );
+            }
+        }
+    }
+
+    Ok(())
+}
+
 /// Registers `event` to be signaled whenever a protocol interface is registered for
 /// `protocol` by [`install_protocol_interface`] or [`reinstall_protocol_interface`].
 ///