Re: [PATCH v2 10/12] rust: vfio: separate common device handling from PCI

[Danilo Krummrich]

On Fri Oct 23, 2026 at 5:40 PM CEST, Simon Song wrote:
> +/// A VFIO device, independent of its bus-specific wrapper.
> +///
> +/// # Invariants
> +///
> +/// The underlying `vfio_device` is initialized and remains alive while borrowed.
> +/// Callback contexts are borrowed only for the corresponding callback and thread;
> +/// only [`Normal`] references can be reference-counted.
> +#[repr(transparent)]
> +pub struct Device<Ctx: DeviceContext = Normal> {
> +    raw: Opaque<bindings::vfio_device>,
> +    _context: PhantomData<(Ctx, NotThreadSafe)>,
> +}

I think this needs to be declared differently, see below.

> +impl<Ctx: DeviceContext> Device<Ctx> {
> +    /// # Safety
> +    ///
> +    /// `raw` must point to an initialized `vfio_device` that remains alive for
> +    /// `'a`. The caller must provide the guarantees of `Ctx` throughout the borrow.
> +    #[cfg_attr(not(CONFIG_VFIO_PCI_CORE), expect(dead_code))]
> +    unsafe fn from_raw<'a>(raw: *mut bindings::vfio_device) -> &'a Self {
> +        // SAFETY: Self is a transparent wrapper; the caller guarantees validity and context.
> +        unsafe { &*raw.cast() }
> +    }
> +
> +    fn as_raw(&self) -> *mut bindings::vfio_device {
> +        self.raw.get()
> +    }
> +}
> +
> +// SAFETY: VFIO uses the embedded class device's refcount to own the full allocation.
> +unsafe impl AlwaysRefCounted for Device {
> +    fn inc_ref(&self) {
> +        // SAFETY: A shared reference keeps the initialized VFIO device alive.
> +        unsafe { bindings::get_device(&raw mut (*self.as_raw()).device) };
> +    }
> +
> +    unsafe fn dec_ref(obj: NonNull<Self>) {
> +        // SAFETY: The caller owns the VFIO reference being released.
> +        unsafe { bindings::put_device(&raw mut (*obj.as_ref().as_raw()).device) };
> +    }
> +}
> +
> +// SAFETY: Ordinary VFIO references expose no unsynchronized callback operations.
> +unsafe impl Send for Device {}
> +// SAFETY: Shared access only changes the embedded class device's reference count.
> +unsafe impl Sync for Device {}
> +
> +/// Callback data stored after the bus-specific C device structure.
> +///
> +/// # Invariants
> +///
> +/// `reg_data` is dangling when unpublished; otherwise it borrows pinned data kept
> +/// alive until callback access has ended. `open_data` is null outside a successful
> +/// first open, otherwise owns a pinned allocation until the last close.
> +/// The bus-specific layer drains callbacks before releasing either allocation.
> +struct CallbackData<R: CovariantForLt, O: ForLt> {
> +    reg_data: UnsafeCell<NonNull<R::Of<'static>>>,
> +    open_data: UnsafeCell<*mut O::Of<'static>>,
> +}

We don't need to arbitrarily limit registration data to require covariance,
eventually we will need invariant data.

Also, I think we don't need ForLt/CovariantForLt at all. Both types are
associated types of an ops trait, so they can just be GATs, i.e.:

	type RegistrationData<'a>: Send + Sync + 'a;
	type OpenData<'a, 'data: 'a>: Send + Sync + 'a;

> +
> +#[cfg_attr(not(CONFIG_VFIO_PCI_CORE), expect(dead_code))]
> +impl<R: CovariantForLt, O: ForLt> CallbackData<R, O> {
> +    fn new() -> Self {
> +        Self {
> +            reg_data: UnsafeCell::new(NonNull::dangling()),
> +            open_data: UnsafeCell::new(core::ptr::null_mut()),
> +        }
> +    }
> +
> +    /// Access registration data without allowing its erased lifetime to escape.
> +    ///
> +    /// # Safety
> +    ///
> +    /// Registration data must be published and remain valid through this call.
> +    /// The caller must exclude concurrent changes to `reg_data`.
> +    unsafe fn registration_data_with<V>(&self, f: impl for<'a> FnOnce(&'a R::Of<'a>) -> V) -> V {
> +        // SAFETY: The caller keeps registration data alive and excludes mutation.
> +        // Covariance permits shortening its erased lifetime to this borrow.
> +        let reg_data = unsafe { (*self.reg_data.get()).cast::<R::Of<'_>>().as_ref() };
> +        f(reg_data)
> +    }
> +
> +    /// # Safety
> +    ///
> +    /// The caller must be an I/O callback while VFIO keeps the device open.
> +    /// Registration and open data must remain valid, with their pointer slots
> +    /// unchanged throughout the call.
> +    unsafe fn callback_data_with<V>(
> +        &self,
> +        f: impl for<'borrow, 'data> FnOnce(&'borrow R::Of<'data>, Pin<&'borrow O::Of<'data>>) -> V,
> +    ) -> V {
> +        // SAFETY: The caller excludes open/close and unregister while borrowing the data.
> +        unsafe {
> +            self.registration_data_with(|reg_data| {
> +                let ptr = (*self.open_data.get()).cast::<O::Of<'_>>();
> +                // The data lifetime stays independent of the callback borrow,
> +                // so borrowed OpenData fields cannot be stored back into OpenData.
> +                f(reg_data, Pin::new_unchecked(&*ptr))
> +            })
> +        }
> +    }

I don't think we necessarily need the extra closure, we only ever pass this to
callbacks of the ops trait, but it would allow us to make the accessor safe:

If we move them onto Device directly, we can leverage the type states. E.g. we
can group them with traits and only implement e.g. callback_data_with() for type
states that allow for this.

> +
> +    /// # Safety
> +    ///
> +    /// Call only during exclusive first open, with no existing open data.
> +    /// All resources borrowed for `'a` must remain valid until `close()` returns.
> +    unsafe fn open<'a, I: PinInit<O::Of<'a>, Error>>(&self, init: impl FnOnce() -> I) -> Result {
> +        // Allocate before calling the driver so allocation failure cannot follow open.
> +        let data = KBox::<O::Of<'a>>::new_uninit(GFP_KERNEL)?;
> +        let data = data.write_pin_init(init())?;
> +
> +        // SAFETY: Only the owning pointer moves; close drops the allocation in place.
> +        let raw = KBox::into_raw(unsafe { Pin::into_inner_unchecked(data) });
> +        // SAFETY: First open is exclusive; the caller keeps borrowed resources alive
> +        // through close. Lifetimes do not affect the allocation's layout.
> +        unsafe { *self.open_data.get() = raw.cast::<O::Of<'static>>() };
> +        Ok(())
> +    }
> +
> +    /// # Safety
> +    ///
> +    /// Call once after a successful open, with all callbacks drained and the next
> +    /// open excluded. Registration data and borrowed resources must remain alive.
> +    unsafe fn close(&self) {
> +        // SAFETY: Last close exclusively takes the allocation published by open.
> +        let raw = unsafe { core::mem::replace(&mut *self.open_data.get(), core::ptr::null_mut()) };
> +        // SAFETY: The allocation came from KBox::into_raw and is destroyed without moving it.
> +        unsafe { drop(KBox::from_raw(raw)) };
> +    }

This seems like a bit of a back and forth because logically we want those
methods on vfio::Device so you can take advantage of the type system
information.

But in this implementation CallbackData lives in the vfio::pci::Device, likely
for memory layout reasons.

But we can also do something like this:

	#[repr(C)]
	pub struct DeviceBase<Raw, R, O, T, Ctx = Normal> {
	    raw: Opaque<Raw>,
	    data: CallbackData<R, O>,
	    _context: PhantomData<(T, Ctx, NotThreadSafe)>,
	}

	// vfio.rs
	pub type Device<T, Ctx = Normal> = DeviceBase<
	    bindings::vfio_device,
	    T,
	    <T as Operations>::RegistrationData<'static>,
	    <T as Operations>::OpenData<'static, 'static>,
	    Ctx,
	>;
	
	// vfio/pci.rs
	pub type Device<T, Ctx = Normal> = super::DeviceBase<
	    bindings::vfio_pci_core_device,
	    T,
	    <T as Operations>::RegistrationData<'static>,
	    <T as Operations>::OpenData<'static, 'static>,
	    Ctx,
	>;
	
Of course we can't have AsRef<vfio::Device> with this, but we don't need it.

For instance, there's only one AlwaysRefCounted impl needed:

  unsafe impl<Raw: RawDevice, R, O, T> AlwaysRefCounted
      for DeviceBase<Raw, R, O, Normal>
  {
      fn inc_ref(&self) {
          let vdev = unsafe { Raw::raw_vfio_device(self.as_raw()) };

          unsafe { bindings::get_device(&raw mut (*vdev).device) };
      }

      unsafe fn dec_ref(obj: NonNull<Self>) {
          let raw = unsafe {
              Opaque::cast_into(&raw const (*obj.as_ptr()).raw)
          };

          let vdev = unsafe { Raw::raw_vfio_device(raw) };

          unsafe { bindings::put_device(&raw mut (*vdev).device) };
      }
  }

We can either rely on the fact that vfio_pci_core_device's first member is the
vfio_device, or make it explicit with something like:

	/// Maps a C device structure to its embedded VFIO device.
	///
	/// # Safety
	///
	/// Implementations must return the address of the `vfio_device` at the
	/// beginning of `Self`, preserving the input pointer's provenance.
	unsafe trait RawDevice {
	    /// # Safety
	    ///
	    /// `raw` must point to an initialized, live `Self`.
	    unsafe fn raw_vfio_device(
	        raw: *mut Self,
	    ) -> *mut bindings::vfio_device;
	}

Not that this isn't only to make use of some type system information (there's
other ways to achieve that too) but also to make vfio::Device useful for non-PCI
implementations, as they need CallbackData too.