From 9e90dede6c134e34f6aad6af01002fa1f4dbc7ed Mon Sep 17 00:00:00 2001
From: Anton Bulakh <self@necauqua.dev>
Date: Mon, 8 Nov 2021 21:13:51 +0200
Subject: [PATCH] Implement the ConnectController/DisconnectController methods

---
 CHANGELOG.md      |  2 ++
 src/table/boot.rs | 51 +++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d784835ca..9c2381cde 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,8 @@
   `LoadedImage::load_options_as_cstr16`.
 - Added `Align::offset_up_to_alignment`, `Align::round_up_to_alignment`,
   and `Align::align_buf`.
+- Added `BootServices::connect_controller` and
+  `BootServices::disconnect_controller`.
 
 ### Changed
 
diff --git a/src/table/boot.rs b/src/table/boot.rs
index 04689cab4..f9e207fc4 100644
--- a/src/table/boot.rs
+++ b/src/table/boot.rs
@@ -119,8 +119,17 @@ pub struct BootServices {
     ) -> Status,
 
     // Driver support services
-    connect_controller: usize,
-    disconnect_controller: usize,
+    connect_controller: unsafe extern "efiapi" fn(
+        controller: Handle,
+        driver_image: Option<Handle>,
+        remaining_device_path: *const DevicePath,
+        recursive: bool,
+    ) -> Status,
+    disconnect_controller: unsafe extern "efiapi" fn(
+        controller: Handle,
+        driver_image: Option<Handle>,
+        child: Option<Handle>,
+    ) -> Status,
 
     // Protocol open / close services
     open_protocol: extern "efiapi" fn(
@@ -694,6 +703,44 @@ impl BootServices {
         unsafe { (self.set_watchdog_timer)(timeout, watchdog_code, data_len, data) }.into()
     }
 
+    /// Connect one or more drivers to a controller.
+    ///
+    /// Usually one disconnects and then reconnects certain drivers
+    /// to make them rescan some state that changed, e.g. reconnecting
+    /// a `BlockIO` handle after your app changed the partitions somehow.
+    pub fn connect_controller(
+        &self,
+        controller: Handle,
+        driver_image: Option<Handle>,
+        remaining_device_path: Option<&DevicePath>,
+        recursive: bool,
+    ) -> Result {
+        unsafe {
+            (self.connect_controller)(
+                controller,
+                driver_image,
+                remaining_device_path
+                    .map(|dp| dp as _)
+                    .unwrap_or(ptr::null()),
+                recursive,
+            )
+        }
+        .into_with_err(|_| ())
+    }
+
+    /// Disconnect one or more drivers from a controller.
+    ///
+    /// See [`connect_controller`][Self::connect_controller].
+    pub fn disconnect_controller(
+        &self,
+        controller: Handle,
+        driver_image: Option<Handle>,
+        child: Option<Handle>,
+    ) -> Result {
+        unsafe { (self.disconnect_controller)(controller, driver_image, child) }
+            .into_with_err(|_| ())
+    }
+
     /// Open a protocol interface for a handle.
     ///
     /// This function attempts to get the protocol implementation of a