[PATCH v2 7/8] rust: pci: add typed SR-IOV PF registration data

[Zhi Wang <zhiw@nvidia.com>]

Rust PF and VF drivers bind to separate PCI devices and may reside
in different modules. A VF driver that calls PF operations needs
typed access to the data exposed by the PF driver. This data must
remain valid until VF driver removal completes.

Add `VfRegistration` to register a pinned Rust object in the PF's
driver data. Add accessors for VF drivers to borrow this object after
checking the requested Rust type. During registration teardown,
disable SR-IOV and wait for VF remove callbacks to finish before
dropping the object.

Co-developed-by: Danilo Krummrich <dakr@kernel.org>
Signed-off-by: Danilo Krummrich <dakr@kernel.org>
Signed-off-by: Zhi Wang <zhiw@nvidia.com>
---
 include/linux/pci.h      |   7 ++
 rust/kernel/pci.rs       |  38 ++++--
 rust/kernel/pci/sriov.rs | 253 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 290 insertions(+), 8 deletions(-)
 create mode 100644 rust/kernel/pci/sriov.rs

diff --git a/include/linux/pci.h b/include/linux/pci.h
index d31a8d107b1e..9b6ae544469e 100644
--- a/include/linux/pci.h
+++ b/include/linux/pci.h
@@ -551,6 +551,13 @@ struct pci_dev {
 	u16		ats_cap;	/* ATS Capability offset */
 	u8		ats_stu;	/* ATS Smallest Translation Unit */
 #endif
+#if defined(CONFIG_PCI_IOV) && defined(CONFIG_RUST)
+	/*
+	 * Private data owned by the PF's Rust driver, readable by VF drivers
+	 * through the PCI VF registration data Rust abstraction.
+	 */
+	void		*vf_registration_data_rust;
+#endif
 #ifdef CONFIG_PCI_PRI
 	u16		pri_cap;	/* PRI Capability offset */
 	u32		pri_reqs_alloc; /* Number of PRI requests allocated */
diff --git a/rust/kernel/pci.rs b/rust/kernel/pci.rs
index cb3ed2207075..831b76744a30 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -35,6 +35,8 @@
 mod id;
 mod io;
 mod irq;
+#[cfg(CONFIG_PCI_IOV)]
+pub mod sriov;
 
 pub use self::id::{
     Class,
@@ -55,6 +57,8 @@
     IrqVector,
     IrqVectorRegistration, //
 };
+#[cfg(CONFIG_PCI_IOV)]
+pub use self::sriov::VfRegistration;
 
 /// An adapter for the registration of PCI drivers.
 pub struct Adapter<T: Driver>(T);
@@ -166,7 +170,16 @@ extern "C" fn sriov_configure_callback(
         // INVARIANT: `pdev` is valid for the duration of `sriov_configure_callback()`.
         let pdev = unsafe { &*pdev.cast::<Device<device::CoreInternal<'_>>>() };
 
-        from_result(|| T::sriov_configure(pdev, nr_virtfn))
+        // SAFETY: `sriov_configure` is called only after a successful probe and before unbind, so
+        // the stored pointer has type `T::Data<'_>` and remains valid throughout this callback.
+        let data = unsafe { pdev.as_ref().drvdata_borrow::<T::Data<'_>>() };
+
+        from_result(|| {
+            if !pdev.is_physfn() {
+                return Err(ENODEV);
+            }
+            T::sriov_configure(pdev, data, nr_virtfn)
+        })
     }
 }
 
@@ -377,8 +390,13 @@ fn unbind<'bound>(dev: &'bound Device<device::Core<'_>>, this: Pin<&Self::Data<'
     ///
     /// ```
     /// # use kernel::{device::Core, pci, prelude::*};
+    /// # struct Data;
     /// #[cfg(CONFIG_PCI_IOV)]
-    /// fn sriov_configure(dev: &pci::Device<Core<'_>>, nr_virtfn: i32) -> Result<i32> {
+    /// fn sriov_configure(
+    ///     dev: &pci::Device<Core<'_>>,
+    ///     _this: Pin<&Data>,
+    ///     nr_virtfn: i32,
+    /// ) -> Result<i32> {
     ///     if nr_virtfn == 0 {
     ///         dev.disable_sriov();
     ///     } else {
@@ -388,8 +406,12 @@ fn unbind<'bound>(dev: &'bound Device<device::Core<'_>>, this: Pin<&Self::Data<'
     /// }
     /// ```
     #[cfg(CONFIG_PCI_IOV)]
-    fn sriov_configure(dev: &Device<device::Core<'_>>, nr_virtfn: i32) -> Result<i32> {
-        let _ = (dev, nr_virtfn);
+    fn sriov_configure<'bound>(
+        dev: &'bound Device<device::Core<'_>>,
+        this: Pin<&Self::Data<'bound>>,
+        nr_virtfn: i32,
+    ) -> Result<i32> {
+        let _ = (dev, this, nr_virtfn);
         build_error!(crate::error::VTABLE_DEFAULT_ERROR)
     }
 }
@@ -496,24 +518,24 @@ pub fn resource_start(&self, bar: u32) -> Result<bindings::resource_size_t> {
     }
 
     /// Returns `true` if this device is a Physical Function (PF).
+    #[cfg(CONFIG_PCI_IOV)]
     #[inline]
-    #[expect(dead_code)]
     pub(crate) fn is_physfn(&self) -> bool {
         // SAFETY: `self.as_raw` is a valid pointer to a `struct pci_dev`.
         unsafe { (*self.as_raw()).is_physfn() != 0 }
     }
 
     /// Returns `true` if this device is a Virtual Function (VF).
+    #[cfg(CONFIG_PCI_IOV)]
     #[inline]
-    #[expect(dead_code)]
-    pub(crate) fn is_virtfn(&self) -> bool {
+    pub fn is_virtfn(&self) -> bool {
         // SAFETY: `self.as_raw` is a valid pointer to a `struct pci_dev`.
         unsafe { (*self.as_raw()).is_virtfn() != 0 }
     }
 
     /// Returns the number of Virtual Functions (VF) enabled for a Physical Function (PF).
     #[cfg(CONFIG_PCI_IOV)]
-    pub(crate) fn num_vf(&self) -> i32 {
+    pub fn num_vf(&self) -> i32 {
         // SAFETY: `self.as_raw` is a valid pointer to a `struct pci_dev`.
         unsafe { bindings::pci_num_vf(self.as_raw()) }
     }
diff --git a/rust/kernel/pci/sriov.rs b/rust/kernel/pci/sriov.rs
new file mode 100644
index 000000000000..90abe878f67c
--- /dev/null
+++ b/rust/kernel/pci/sriov.rs
@@ -0,0 +1,253 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Abstractions for PCI Single Root I/O Virtualization (SR-IOV) drivers.
+
+use super::Device as PciDevice;
+use crate::{
+    bindings,
+    device, //
+    prelude::*,
+    types::{
+        CovariantForLt,
+        ForLt, //
+    },
+};
+use core::{
+    any::TypeId,
+    marker::PhantomPinned, //
+};
+
+/// Wrapper for VF registration data stored inside a [`VfRegistration`].
+///
+/// Stores a [`TypeId`] header (derived from `F`) followed by the pinned data,
+/// so that [`PciDevice::vf_registration_data_with()`] can verify the type at
+/// runtime.
+#[repr(C)]
+#[pin_data]
+struct VfRegistrationData<'a, F: ForLt + 'static> {
+    type_id: TypeId,
+    #[pin]
+    data: F::Of<'a>,
+}
+
+static_assert!(
+    core::mem::offset_of!(VfRegistrationData<'static, CovariantForLt!(())>, type_id) == 0
+);
+
+impl<'a, F: ForLt + 'static> VfRegistrationData<'a, F> {
+    /// Pin-initializer for the registration data.
+    fn new<D>(data: D) -> impl PinInit<Self, Error> + use<'a, D, F>
+    where
+        D: PinInit<F::Of<'a>, Error> + 'a,
+    {
+        try_pin_init!(Self {
+            type_id: TypeId::of::<F>(),
+            data <- data,
+        })
+    }
+}
+
+/// SR-IOV VF registration on a PF device.
+///
+/// Owns the registration data that VF drivers access via
+/// [`PciDevice::vf_registration_data_with()`] and [`PciDevice::vf_registration_data()`].
+///
+/// The Rust PCI adapter removes all VFs before invoking the PF driver's unbind callback or
+/// dropping its data. Drop of a published registration calls `pci_disable_sriov()` before
+/// clearing the pointer and letting the data fields drop.
+#[pin_data(PinnedDrop)]
+pub struct VfRegistration<'a, F: ForLt + 'static> {
+    pdev: &'a PciDevice<device::Bound>,
+    #[pin]
+    inner: VfRegistrationData<'a, F>,
+    published: bool,
+    #[pin]
+    _pin: PhantomPinned,
+}
+
+impl<'a, F: ForLt + 'static> VfRegistration<'a, F>
+where
+    for<'b> F::Of<'b>: Send + Sync,
+{
+    /// Create a new VF registration.
+    ///
+    /// Returns a pin-initializer so the registration can be embedded directly
+    /// in the PF driver's bus device private data.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the containing struct's field ordering drops
+    /// this `VfRegistration` before any resources that the registration data
+    /// borrows.
+    ///
+    /// The caller must invoke this during the PCI driver's probe and embed the result in the driver
+    /// data. On an SR-IOV PF, no VF may be enabled before probe successfully installs the complete
+    /// driver data.
+    pub unsafe fn new<'core, D>(
+        pdev: &'a PciDevice<device::Core<'core>>,
+        data: D,
+    ) -> impl PinInit<Self, Error> + use<'a, 'core, D, F>
+    where
+        D: PinInit<F::Of<'a>, Error> + 'a,
+    {
+        pin_init::pin_init_scope(move || {
+            if pdev.is_virtfn() {
+                return Err(ENODEV);
+            }
+
+            let published = pdev.is_physfn();
+            if published {
+                if pdev.num_vf() != 0 {
+                    return Err(EBUSY);
+                }
+
+                if !pdev.vf_registration_data_rust().is_null() {
+                    return Err(EBUSY);
+                }
+            }
+
+            Ok(try_pin_init!(Self {
+                pdev,
+                inner <- VfRegistrationData::new(data),
+                published,
+                _pin: PhantomPinned,
+                _: {
+                    if *published {
+                        // Store the pointer to the pinned `VfRegistrationData`
+                        // on the PCI device so VF drivers can find it.
+                        pdev.set_vf_registration_data_rust(
+                            core::ptr::from_ref(inner.as_ref().get_ref()).cast_mut().cast(),
+                        );
+                    }
+                },
+            }))
+        })
+    }
+}
+
+#[pinned_drop]
+impl<F: ForLt + 'static> PinnedDrop for VfRegistration<'_, F> {
+    fn drop(self: Pin<&mut Self>) {
+        if !self.published {
+            return;
+        }
+
+        // SAFETY: `pci_disable_sriov()` is safe to call on any `pci_dev`; it
+        // is a no-op if the device has no VFs enabled. When VFs are enabled,
+        // this blocks until all VF `remove()` callbacks complete.
+        unsafe { bindings::pci_disable_sriov(self.pdev.as_raw()) };
+
+        // After `pci_disable_sriov()` all VFs are gone, so no one can read
+        // the pointer anymore.
+        self.pdev
+            .set_vf_registration_data_rust(core::ptr::null_mut());
+
+        // The pinned `inner` field is dropped automatically after this returns.
+    }
+}
+
+// SAFETY: The inner data is `Send` (enforced by the bound), and `&PciDevice` is `Send + Sync`.
+unsafe impl<F: ForLt> Send for VfRegistration<'_, F> where for<'a> F::Of<'a>: Send {}
+
+// SAFETY: The inner data is `Send + Sync`. `VfRegistration` doesn't expose mutable access;
+// VF drivers only read the data through an immutable pinned reference.
+unsafe impl<F: ForLt> Sync for VfRegistration<'_, F> where for<'a> F::Of<'a>: Send + Sync {}
+
+impl<Ctx: device::DeviceContext> PciDevice<Ctx> {
+    /// Returns the raw `vf_registration_data_rust` pointer from this device.
+    fn vf_registration_data_rust(&self) -> *mut core::ffi::c_void {
+        // SAFETY: `self.as_raw()` is valid.
+        unsafe { (*self.as_raw()).vf_registration_data_rust }
+    }
+
+    /// Sets the `vf_registration_data_rust` pointer on this device.
+    fn set_vf_registration_data_rust(&self, ptr: *mut core::ffi::c_void) {
+        // SAFETY: `self.as_raw()` is valid. PCI probe publishes the data before enabling VFs;
+        // teardown removes all VFs before withdrawing it.
+        unsafe { (*self.as_raw()).vf_registration_data_rust = ptr };
+    }
+}
+
+impl PciDevice<device::Bound> {
+    /// Returns the PF for this VF, or [`ENODEV`] if this is not a VF.
+    fn physfn(&self) -> Result<&PciDevice> {
+        if !self.is_virtfn() {
+            return Err(ENODEV);
+        }
+
+        // SAFETY: `self.as_raw()` is valid and this VF uses the `physfn` union field.
+        let pf = unsafe { (*self.as_raw()).__bindgen_anon_1.physfn };
+        if pf.is_null() {
+            return Err(ENODEV);
+        }
+
+        // SAFETY: PCI holds a PF reference until VF removal completes. The returned borrow
+        // cannot outlive this bound VF, and `PciDevice` is a transparent wrapper of `pci_dev`.
+        Ok(unsafe { &*pf.cast() })
+    }
+
+    /// Internal helper: reads the `vf_registration_data_rust` pointer from the
+    /// PF, checks the `TypeId`, and returns a pinned reference.
+    ///
+    /// # Safety
+    ///
+    /// The returned borrow must be confined by a closure higher-ranked independently over its
+    /// borrow and data lifetimes, or `F` must be covariant in its encoded lifetime.
+    unsafe fn vf_registration_data_pinned<F: ForLt + 'static>(&self) -> Result<Pin<&F::Of<'_>>> {
+        let pf = self.physfn()?;
+
+        let ptr = pf.vf_registration_data_rust();
+        if ptr.is_null() {
+            return Err(ENOENT);
+        }
+
+        // SAFETY: The Rust PCI adapter keeps the PF data installed until VF removal completes.
+        // `ptr` points to a `VfRegistrationData` whose first field is a `TypeId`.
+        let type_id = unsafe { ptr.cast::<TypeId>().read() };
+        if type_id != TypeId::of::<F>() {
+            return Err(EINVAL);
+        }
+
+        // SAFETY: TypeId check confirms the stored type matches `F`. The data
+        // is pinned inside the PF's driver data struct. Lifetime shortening
+        // from the PF's binding scope to `'_` is layout-compatible.
+        let data_ptr = unsafe {
+            let vfrd = ptr.cast::<VfRegistrationData<'_, F>>();
+            &raw const (*vfrd).data
+        };
+
+        // SAFETY: `data` is structurally pinned inside `VfRegistrationData`.
+        Ok(unsafe { Pin::new_unchecked(&*data_ptr) })
+    }
+
+    /// Access the VF registration data through a closure with an HRTB lifetime.
+    ///
+    /// `F` is the [`ForLt`](trait@ForLt) encoding of the data type. Returns
+    /// [`ENODEV`] if this is not a VF, [`ENOENT`] if no data was registered,
+    /// or [`EINVAL`] if `F` does not match the type registered by the PF.
+    ///
+    /// The closure's borrow and the registration data's lifetime are independent, so a borrow of
+    /// the context cannot be stored in invariant registration data.
+    pub fn vf_registration_data_with<F: ForLt + 'static, R>(
+        &self,
+        f: impl for<'borrow, 'data> FnOnce(Pin<&'borrow F::Of<'data>>) -> R,
+    ) -> Result<R> {
+        // SAFETY: The higher-ranked closure prevents the borrow from escaping or being stored in
+        // invariant data by keeping its lifetime independent of the erased data lifetime.
+        let pinned = unsafe { self.vf_registration_data_pinned::<F>()? };
+        Ok(f(pinned))
+    }
+
+    /// Returns a pinned reference to the VF registration data.
+    ///
+    /// Available only when `F` implements [`CovariantForLt`](trait@crate::types::CovariantForLt),
+    /// guaranteeing that shortening the PF data lifetime is sound.
+    ///
+    /// For non-covariant types, use [`Self::vf_registration_data_with()`].
+    ///
+    /// It returns the same errors as [`Self::vf_registration_data_with()`].
+    pub fn vf_registration_data<F: CovariantForLt + 'static>(&self) -> Result<Pin<&F::Of<'_>>> {
+        // SAFETY: `CovariantForLt` permits shortening the encoded lifetime to this borrow.
+        unsafe { self.vf_registration_data_pinned::<F>() }
+    }
+}
-- 
2.53.0