[PATCH v2 05/12] rust: pci: add reset callbacks

[Simon Song <xinmengs@nvidia.com>]

VFIO PCI variant drivers will use reset_prepare() and reset_done() to
forward PCI reset notifications to their VFIO registration.

Adding these callbacks to pci::Driver lets VFIO variants reuse the PCI
adapter's reset callback dispatch. A following patch removes the
duplicate vfio::pci::Adapter.

Because there is no gurantee that drvdata must be valid when reset is
issued, using callback_data_with which contains the helper to null-check
drv_data, this could occur during probe if reset is called before
set_drvdata

Suggested-by: Zhi Wang <zhiw@nvidia.com>
Signed-off-by: Simon Song <xinmengs@nvidia.com>
---
 rust/kernel/pci.rs      | 70 +++++++++++++++++++++++++++++++++++++++++
 rust/kernel/vfio/pci.rs | 12 +++++++
 2 files changed, 82 insertions(+)

diff --git a/rust/kernel/pci.rs b/rust/kernel/pci.rs
index a156927e23d..adc7e92d44e 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -109,6 +109,9 @@ unsafe fn register(
             (*pdrv.get()).remove = Some(Self::remove_callback);
             (*pdrv.get()).id_table = T::ID_TABLE.as_ptr();
             (*pdrv.get()).driver_managed_dma = T::DRIVER_MANAGED_DMA;
+            if T::HAS_RESET_PREPARE || T::HAS_RESET_DONE {
+                (*pdrv.get()).err_handler = &Self::ERROR_HANDLERS;
+            }
             #[cfg(CONFIG_PCI_IOV)]
             if T::HAS_SRIOV_ENABLE {
                 (*pdrv.get()).sriov_configure = Some(Self::sriov_configure_callback);
@@ -175,6 +178,61 @@ extern "C" fn remove_callback(pdev: *mut bindings::pci_dev) {
         // Drop private data before returning to PCI core, while its removal context is valid.
         drop(data);
     }
+
+    /// # Safety
+    ///
+    /// `dev` must belong to `T`. Any installed driver data must remain installed until this
+    /// call returns. The data lifetime is independent of the callback borrow.
+    unsafe fn callback_data_with<R>(
+        dev: *mut bindings::device,
+        f: impl for<'borrow, 'data> FnOnce(Pin<&'borrow T::Data<'data>>) -> R,
+    ) -> Option<R> {
+        // SAFETY: The callback supplies a device bound to `T`.
+        let data = unsafe { bindings::dev_get_drvdata(dev) };
+        if data.is_null() {
+            return None;
+        }
+        // SAFETY: PCI installed pinned `T::Data`; the caller excludes its removal.
+        Some(f(unsafe {
+            Pin::new_unchecked(&*data.cast::<T::Data<'_>>())
+        }))
+    }
+
+    extern "C" fn reset_prepare_callback(pdev: *mut bindings::pci_dev) {
+        // SAFETY: PCI holds the device lock throughout reset.
+        let dev = unsafe { &*pdev.cast::<Device<device::CoreInternal<'_>>>() };
+        // SAFETY: The device lock prevents removal of installed driver data.
+        let _ = unsafe {
+            Self::callback_data_with(&raw mut (*pdev).dev, |data| T::reset_prepare(dev, data))
+        };
+    }
+
+    extern "C" fn reset_done_callback(pdev: *mut bindings::pci_dev) {
+        // SAFETY: PCI holds the device lock throughout reset.
+        let dev = unsafe { &*pdev.cast::<Device<device::CoreInternal<'_>>>() };
+        // SAFETY: The device lock prevents removal of installed driver data.
+        let _ = unsafe {
+            Self::callback_data_with(&raw mut (*pdev).dev, |data| T::reset_done(dev, data))
+        };
+    }
+
+    const ERROR_HANDLERS: bindings::pci_error_handlers = bindings::pci_error_handlers {
+        error_detected: None,
+        mmio_enabled: None,
+        slot_reset: None,
+        reset_prepare: if T::HAS_RESET_PREPARE {
+            Some(Self::reset_prepare_callback)
+        } else {
+            None
+        },
+        reset_done: if T::HAS_RESET_DONE {
+            Some(Self::reset_done_callback)
+        } else {
+            None
+        },
+        resume: None,
+        cor_error_detected: None,
+    };
 }
 
 /// Declares a kernel module that exposes a single PCI driver.
@@ -375,6 +433,18 @@ fn unbind<'bound>(dev: &'bound Device<device::Core<'_>>, this: Pin<&Self::Data<'
         let _ = (dev, this);
     }
 
+    /// Called before a PCI function reset, with the device lock held.
+    fn reset_prepare<'bound>(dev: &Device<device::Core<'_>>, this: Pin<&Self::Data<'bound>>) {
+        let _ = (dev, this);
+        build_error!(crate::error::VTABLE_DEFAULT_ERROR)
+    }
+
+    /// Called after a PCI function reset, with the device lock held.
+    fn reset_done<'bound>(dev: &Device<device::Core<'_>>, this: Pin<&Self::Data<'bound>>) {
+        let _ = (dev, this);
+        build_error!(crate::error::VTABLE_DEFAULT_ERROR)
+    }
+
     /// Enables Single Root I/O Virtualization (SR-IOV) for a PF.
     ///
     /// Called when userspace writes a nonzero VF count to `sriov_numvfs`. The token carries the
diff --git a/rust/kernel/vfio/pci.rs b/rust/kernel/vfio/pci.rs
index 43c14b4e3a1..db30dbce6a8 100644
--- a/rust/kernel/vfio/pci.rs
+++ b/rust/kernel/vfio/pci.rs
@@ -1084,6 +1084,18 @@ unsafe fn unregister(&self) {
         // VGA and runtime PM callbacks. No field of the enclosing data is being destroyed.
         unsafe { bindings::vfio_pci_core_unregister_device(self.dev.core_device()) };
     }
+
+    /// Prepare the variant driver's state for a PCI reset.
+    pub fn reset_prepare(&self) {
+        // SAFETY: This owner keeps the pinned registration data alive.
+        unsafe { self.dev.registration_data_with(T::reset_prepare) };
+    }
+
+    /// Complete the variant driver's state update after a PCI reset.
+    pub fn reset_done(&self) -> Result {
+        // SAFETY: This owner keeps the pinned registration data alive.
+        unsafe { self.dev.registration_data_with(T::reset_done) }
+    }
 }
 
 impl<T: Operations> Drop for Registration<'_, T> {
-- 
2.43.7