[PATCH 08/14] rust: pci: add typed SR-IOV PF registration data

[Zhi Wang <zhiw@nvidia.com>]

A VF driver may need to invoke PF-owned functionality without gaining
access to all private data belonging to the PF driver. The PCI device
alone cannot identify the restricted data type or express how long a
borrow remains valid.

Add a Rust registration slot to `struct pci_dev` and `VfRegistration`
to publish a `ForLt`-encoded payload. The registration is initialized
and pinned inline in the PF driver data, then publishes its address as
the final step of its own initialization. It remains inactive on
conventional PCI functions and rejects VFs.

Check that no VFs or other registration exist before publication. Pinned
drop disables SR-IOV before withdrawing the pointer and dropping the
payload. Managed SR-IOV ordering additionally removes every VF before
PF unbind.

Add `vf_registration_data_with()` for higher-ranked access and
`vf_registration_data()` for covariant data. `TypeId` rejects mismatched
Rust types without exposing the PF device or its complete driver data.

Keep SR-IOV operations behind a verified `pci::sriov::Device` view and
pass a pinned reference to the PCI driver data into `sriov_configure()`.
PF and VF consumers continue to use the ordinary `pci::Driver`
abstraction.

Expose `is_virtfn()` so PF drivers can reject VFs before initializing
hardware and publishing their registration.

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      |   6 +
 rust/kernel/pci.rs       |  69 +++++----
 rust/kernel/pci/sriov.rs | 299 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 345 insertions(+), 29 deletions(-)
 create mode 100644 rust/kernel/pci/sriov.rs

diff --git a/include/linux/pci.h b/include/linux/pci.h
index bc0d36204940..1ccc7fee7495 100644
--- a/include/linux/pci.h
+++ b/include/linux/pci.h
@@ -352,6 +352,9 @@ struct rcec_ea;
  *			Such bridges are allocated additional MMIO and bus
  *			number resources to allow for hierarchy expansion.
  * @is_pciehp:		PCIe Hot-Plug Capable bridge.
+ * @vf_registration_data_rust: Rust registration data published by the PF
+ *			before enabling VFs and retained until all VFs are
+ *			removed. The PF driver must use managed_sriov.
  */
 struct pci_dev {
 	struct list_head bus_list;	/* Node in per-bus list */
@@ -551,6 +554,9 @@ 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)
+	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 f2d71daef83a..46edc8007d78 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -37,6 +37,8 @@
 mod id;
 mod io;
 mod irq;
+#[cfg(CONFIG_PCI_IOV)]
+pub mod sriov;
 
 pub use self::cap::{
     ExtCapId,
@@ -64,6 +66,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);
@@ -160,13 +164,18 @@ extern "C" fn sriov_configure_callback(
         pdev: *mut bindings::pci_dev,
         nr_virtfn: c_int,
     ) -> c_int {
-        // SAFETY: The PCI bus only ever calls the sriov_configure callback with a valid pointer to
-        // a `struct pci_dev`.
-        //
-        // INVARIANT: `pdev` is valid for the duration of `sriov_configure_callback()`.
+        // SAFETY: The PCI bus invokes this callback with a valid device bound to this driver. The
+        // `CoreInternal` context is valid for the callback's duration.
         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(|| {
+            let dev = sriov::Device::try_from_pci(pdev)?;
+            T::sriov_configure(dev, data, nr_virtfn)
+        })
     }
 }
 
@@ -355,41 +364,44 @@ fn unbind<'bound>(dev: &'bound Device<device::Core<'_>>, this: Pin<&Self::Data<'
         let _ = (dev, this);
     }
 
-    /// Single Root I/O Virtualization (SR-IOV) configure.
+    /// Configures Single Root I/O Virtualization (SR-IOV) for a Physical Function (PF).
     ///
-    /// Called when a user-space application enables or disables the SR-IOV capability for a
-    /// [`Device`] by writing the number of Virtual Functions (VF), `nr_virtfn` or zero to the
-    /// sysfs file `sriov_numvfs` for this device. Implementing this callback is optional.
+    /// The PCI core invokes this callback when userspace writes the number of Virtual Functions
+    /// (VFs), or zero, to the PF's `sriov_numvfs` sysfs file. For managed SR-IOV it is also called
+    /// with zero before [`Self::unbind`] when the PF still has enabled VFs.
     ///
-    /// Further, and unlike for a PCI driver written in C, when a PF device with enabled VFs is
-    /// unbound from its bound [`Driver`], the `sriov_configure()` callback is invoked to disable
-    /// SR-IOV before the `unbind()` callback. This guarantees that when a VF device is bound to a
-    /// driver, the underlying PF device is bound to a driver, too.
+    /// `dev` is a verified SR-IOV PF in the [`device::Core`] callback context. It can be converted
+    /// to the underlying PCI device through [`sriov::Device::as_pci`]. `this` is the private data
+    /// returned by [`Self::probe`]. Both remain valid for the duration of the callback.
     ///
-    /// Upon success, this callback must return the number of VFs that were enabled, or zero if
-    /// SR-IOV was disabled.
-    ///
-    /// See [PCI Express I/O Virtualization].
-    ///
-    /// [PCI Express I/O Virtualization]: https://docs.kernel.org/PCI/pci-iov-howto.html
+    /// Upon success, return the number of VFs that were enabled, or zero if SR-IOV was disabled.
     ///
     /// # Examples
     ///
     /// ```
     /// # use kernel::{device::Core, pci, prelude::*};
-    /// #[cfg(CONFIG_PCI_IOV)]
-    /// fn sriov_configure(dev: &pci::Device<Core<'_>>, nr_virtfn: i32) -> Result<i32> {
+    /// # struct Data;
+    /// fn sriov_configure(
+    ///     dev: &pci::sriov::Device<Core<'_>>,
+    ///     _this: Pin<&Data>,
+    ///     nr_virtfn: i32,
+    /// ) -> Result<i32> {
     ///     if nr_virtfn == 0 {
     ///         dev.disable_sriov();
     ///     } else {
     ///         dev.enable_sriov(nr_virtfn)?;
     ///     }
+    ///
     ///     Ok(nr_virtfn)
     /// }
     /// ```
     #[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 sriov::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)
     }
 }
@@ -508,24 +520,23 @@ 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)]
-    #[expect(dead_code)]
     pub(crate) 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()) }
@@ -593,7 +604,7 @@ pub fn set_master(&self) {
     /// Enable the Single Root I/O Virtualization (SR-IOV) capability for this device,
     /// where `nr_virtfn` is number of Virtual Functions (VF) to enable.
     #[cfg(CONFIG_PCI_IOV)]
-    pub fn enable_sriov(&self, nr_virtfn: i32) -> Result {
+    pub(crate) fn enable_sriov(&self, nr_virtfn: i32) -> Result {
         // SAFETY:
         // `self.as_raw` returns a valid pointer to a `struct pci_dev`.
         //
@@ -609,7 +620,7 @@ pub fn enable_sriov(&self, nr_virtfn: i32) -> Result {
 
     /// Disable the Single Root I/O Virtualization (SR-IOV) capability for this device.
     #[cfg(CONFIG_PCI_IOV)]
-    pub fn disable_sriov(&self) {
+    pub(crate) fn disable_sriov(&self) {
         // SAFETY:
         // `self.as_raw` returns a valid pointer to a `struct pci_dev`.
         //
diff --git a/rust/kernel/pci/sriov.rs b/rust/kernel/pci/sriov.rs
new file mode 100644
index 000000000000..efbe444e0733
--- /dev/null
+++ b/rust/kernel/pci/sriov.rs
@@ -0,0 +1,299 @@
+// 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,
+    num::NonZero, //
+};
+
+/// A PCI Physical Function (PF) with an SR-IOV capability.
+///
+/// This capability view is created only after the PCI abstraction verifies that the device is an
+/// SR-IOV PF. Its device context follows the same hierarchy as [`PciDevice`].
+#[repr(transparent)]
+pub struct Device<Ctx: device::DeviceContext = device::Normal>(PciDevice<Ctx>);
+
+impl<Ctx: device::DeviceContext> Device<Ctx> {
+    pub(super) fn try_from_pci(pdev: &PciDevice<Ctx>) -> Result<&Self> {
+        // SAFETY: `pdev.as_raw()` is a valid pointer to a `struct pci_dev`.
+        if unsafe { (*pdev.as_raw()).is_physfn() == 0 } {
+            return Err(ENODEV);
+        }
+
+        // CAST: `Device` is a transparent capability view of `PciDevice` with the same context.
+        // SAFETY: The check above establishes the PF invariant, and the returned reference cannot
+        // outlive `pdev`.
+        Ok(unsafe { &*core::ptr::from_ref(pdev).cast() })
+    }
+
+    /// Returns the underlying PCI device with the same device context.
+    #[inline]
+    pub fn as_pci(&self) -> &PciDevice<Ctx> {
+        &self.0
+    }
+}
+
+impl<Ctx: device::DeviceContext> AsRef<PciDevice<Ctx>> for Device<Ctx> {
+    #[inline]
+    fn as_ref(&self) -> &PciDevice<Ctx> {
+        self.as_pci()
+    }
+}
+
+impl<Ctx: device::DeviceContext> AsRef<device::Device<Ctx>> for Device<Ctx> {
+    #[inline]
+    fn as_ref(&self) -> &device::Device<Ctx> {
+        self.as_pci().as_ref()
+    }
+}
+
+impl<'a> Device<device::Core<'a>> {
+    /// Returns the total number of VFs, or [`None`] if SR-IOV is unavailable.
+    #[inline]
+    pub fn total_vfs(&self) -> Option<NonZero<u16>> {
+        self.as_pci().sriov_get_totalvfs()
+    }
+
+    /// Enables `nr_virtfn` Virtual Functions (VFs).
+    #[inline]
+    pub fn enable_sriov(&self, nr_virtfn: i32) -> Result {
+        self.as_pci().enable_sriov(nr_virtfn)
+    }
+
+    /// Disables all Virtual Functions (VFs).
+    #[inline]
+    pub fn disable_sriov(&self) {
+        self.as_pci().disable_sriov();
+    }
+}
+
+impl Device<device::Bound> {
+    /// Returns the number of currently enabled Virtual Functions (VFs).
+    #[inline]
+    pub fn num_vfs(&self) -> i32 {
+        self.as_pci().num_vf()
+    }
+}
+
+// SAFETY: `Device` is a transparent wrapper around `PciDevice`, and neither type's layout
+// depends on its device context.
+kernel::impl_device_context_deref!(unsafe { Device });
+
+#[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> {
+    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,
+        })
+    }
+}
+
+/// Typed data published by a Physical Function (PF) for its Virtual Functions (VFs).
+///
+/// The registration is initialized in place as part of the PF driver's pinned data. On an SR-IOV
+/// PF it publishes the inline payload after initialization; on a conventional PCI function it is
+/// inactive. A VF is rejected. Bound VFs access the payload through
+/// [`PciDevice::vf_registration_data()`] or [`PciDevice::vf_registration_data_with()`].
+///
+/// Managed SR-IOV removes all VFs before the PF driver is unbound. As a fallback, pinned drop also
+/// disables SR-IOV before withdrawing the payload.
+#[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,
+{
+    /// Publishes typed PF data for bound VF drivers.
+    ///
+    /// This returns a pin-initializer so the registration and payload can be embedded directly in
+    /// the PF driver's pinned data.
+    ///
+    /// Initialization returns [`ENODEV`] for a VF and [`EBUSY`] if the PF has enabled VFs or
+    /// already has a registration.
+    ///
+    /// # Safety
+    ///
+    /// 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, and the driver must use managed SR-IOV. The registration must be dropped before
+    /// anything its payload borrows and must not be forgotten. Probe must have exclusive access to
+    /// the PF registration slot. On a conventional PCI function, the registration remains
+    /// inactive.
+    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 {
+                        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: `self.pdev` is the PF on which this registration was published. The call is a
+        // no-op on the normal managed-SR-IOV teardown path, where all VFs are already disabled.
+        unsafe { bindings::pci_disable_sriov(self.pdev.as_raw()) };
+        self.pdev
+            .set_vf_registration_data_rust(core::ptr::null_mut());
+    }
+}
+
+// SAFETY: The registration and its inline payload may be released from another thread after the
+// PCI core has removed all VFs.
+unsafe impl<F: ForLt> Send for VfRegistration<'_, F> where for<'a> F::Of<'a>: Send {}
+
+// SAFETY: VF consumers receive shared references only, and the payload supports shared access.
+unsafe impl<F: ForLt> Sync for VfRegistration<'_, F> where for<'a> F::Of<'a>: Send + Sync {}
+
+impl<Ctx: device::DeviceContext> PciDevice<Ctx> {
+    fn vf_registration_data_rust(&self) -> *mut core::ffi::c_void {
+        // SAFETY: `self.as_raw()` is a valid pointer to a `struct pci_dev`.
+        unsafe { (*self.as_raw()).vf_registration_data_rust }
+    }
+
+    fn set_vf_registration_data_rust(&self, data: *mut core::ffi::c_void) {
+        // SAFETY: Publication and withdrawal are serialized by PCI probe and managed teardown.
+        unsafe { (*self.as_raw()).vf_registration_data_rust = data };
+    }
+}
+
+impl PciDevice<device::Bound> {
+    /// # 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<'_>>> {
+        if !self.is_virtfn() {
+            return Err(ENODEV);
+        }
+
+        // SAFETY: A VF's `physfn` pointer remains valid for the VF's lifetime. Managed SR-IOV also
+        // keeps the PF driver bound until this VF is unbound.
+        let pf_dev = unsafe { (*self.as_raw()).__bindgen_anon_1.physfn };
+        if pf_dev.is_null() {
+            return Err(ENODEV);
+        }
+
+        // SAFETY: The PF cannot withdraw the pointer until managed teardown has removed this VF.
+        let ptr = unsafe { (*pf_dev).vf_registration_data_rust };
+        if ptr.is_null() {
+            return Err(ENOENT);
+        }
+
+        // SAFETY: The published pointer addresses 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: The type check identifies `F`; lifetime parameters do not affect layout, and the
+        // inline data remains pinned for this VF borrow.
+        let data = unsafe {
+            let registration = ptr.cast::<VfRegistrationData<'_, F>>();
+            &raw const (*registration).data
+        };
+
+        // SAFETY: `data` is structurally pinned in the PF driver's pinned registration.
+        Ok(unsafe { Pin::new_unchecked(&*data) })
+    }
+
+    /// Accesses typed data published by this VF's PF through a closure.
+    ///
+    /// Returns [`ENODEV`] if this device is not a VF, [`ENOENT`] if its PF has not published
+    /// data, or [`EINVAL`] if the registered type does not match `F`.
+    ///
+    /// 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 data = unsafe { self.vf_registration_data_pinned::<F>()? };
+        Ok(f(data))
+    }
+
+    /// Returns typed data published by this VF's PF.
+    ///
+    /// This direct accessor is available only when the encoded data is covariant in its lifetime.
+    /// Use [`Self::vf_registration_data_with()`] for invariant data.
+    ///
+    /// 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>() }
+    }
+}