Re: [PATCH 3/4] rust: devres: get rid of Devres' inner Arc

["Benno Lossin" <lossin@kernel.org>]

On Thu Jun 12, 2025 at 4:51 PM CEST, Danilo Krummrich wrote:
> So far Devres uses an inner memory allocation and reference count, i.e.
> an inner Arc, in order to ensure that the devres callback can't run into
> a use-after-free in case where the Devres object is dropped while the
> devres callback runs concurrently.
>
> Instead, use a completion in order to avoid a potential UAF: In
> Devres::drop(), if we detect that we can't remove the devres action
> anymore, we wait for the completion that is completed from the devres
> callback. If, in turn, we were able to successfully remove the devres
> action, we can just go ahead.
>
> This, again, allows us to get rid of the internal Arc, and instead let
> Devres consume an `impl PinInit<T, E>` in order to return an
> `impl PinInit<Devres<T>, E>`, which enables us to get away with less
> memory allocations.
>
> Additionally, having the resulting explicit synchronization in
> Devres::drop() prevents potential subtle undesired side effects of the
> devres callback dropping the final Arc reference asynchronously within
> the devres callback.
>
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>

This is really nice, good to see the extra allocations gone :)

> ---
>  drivers/gpu/nova-core/driver.rs |   7 +-
>  drivers/gpu/nova-core/gpu.rs    |   6 +-
>  rust/kernel/devres.rs           | 187 +++++++++++++++-----------------
>  rust/kernel/pci.rs              |  20 ++--
>  samples/rust/rust_driver_pci.rs |  19 ++--
>  5 files changed, 117 insertions(+), 122 deletions(-)

> @@ -86,100 +76,93 @@ struct DevresInner<T> {
>  /// # fn no_run(dev: &Device<Bound>) -> Result<(), Error> {
>  /// // SAFETY: Invalid usage for example purposes.
>  /// let iomem = unsafe { IoMem::<{ core::mem::size_of::<u32>() }>::new(0xBAAAAAAD)? };
> -/// let devres = Devres::new(dev, iomem, GFP_KERNEL)?;
> +/// let devres = KBox::pin_init(Devres::new(dev, iomem), GFP_KERNEL)?;
>  ///
>  /// let res = devres.try_access().ok_or(ENXIO)?;
>  /// res.write8(0x42, 0x0);
>  /// # Ok(())
>  /// # }
>  /// ```
> -pub struct Devres<T>(Arc<DevresInner<T>>);
> -
> -impl<T> DevresInner<T> {
> -    fn new(dev: &Device<Bound>, data: T, flags: Flags) -> Result<Arc<DevresInner<T>>> {
> -        let inner = Arc::pin_init(
> -            try_pin_init!( DevresInner {
> -                dev: dev.into(),
> -                callback: Self::devres_callback,
> -                data <- Revocable::new(data),
> -                revoke <- Completion::new(),
> -            }),
> -            flags,
> -        )?;
> -
> -        // Convert `Arc<DevresInner>` into a raw pointer and make devres own this reference until
> -        // `Self::devres_callback` is called.
> -        let data = inner.clone().into_raw();
> +#[pin_data(PinnedDrop)]
> +pub struct Devres<T> {
> +    dev: ARef<Device>,
> +    callback: unsafe extern "C" fn(*mut c_void),

Do I remember correctly that we at some point talked about adding a
comment here for why this is needed? (ie it's needed, because
`Self::callback` might return different addresses?)

> +    #[pin]
> +    data: Revocable<T>,
> +    #[pin]
> +    devm: Completion,
> +    #[pin]
> +    revoke: Completion,

Probably a good idea to add some doc comments explaining what these two
completions track.

(feel free to do these in another patch or in a follow-up)

> +}
>  
> -        // SAFETY: `devm_add_action` guarantees to call `Self::devres_callback` once `dev` is
> -        // detached.
> -        let ret =
> -            unsafe { bindings::devm_add_action(dev.as_raw(), Some(inner.callback), data as _) };
> +impl<T> Devres<T> {
> +    /// Creates a new [`Devres`] instance of the given `data`. The `data` encapsulated within the

Missing double newline after the first sentence.

> +    /// returned `Devres` instance' `data` will be revoked once the device is detached.

Maybe we should link to `Revocable` on the word `revoked`?

> +    pub fn new<'a, E>(
> +        dev: &'a Device<Bound>,
> +        data: impl PinInit<T, E> + 'a,
> +    ) -> impl PinInit<Self, Error> + 'a
> +    where
> +        T: 'a,
> +        Error: From<E>,
> +    {
> +        let callback = Self::devres_callback;

> -        Ok(Devres(inner))
> +    fn remove_action(&self) -> bool {
> +        // SAFETY:
> +        // - `self.dev` is a valid `Device`,
> +        // - the `action` and `data` pointers are the exact same ones as given to devm_add_action()
> +        //   previously,
> +        // - `self` is always valid, even if the action has been released already.
> +        (unsafe {
> +            bindings::devm_remove_action_nowarn(
> +                self.dev.as_raw(),
> +                Some(self.callback),
> +                self.as_ptr().cast_mut().cast(),
> +            )
> +        } == 0)

I don't think the parenthesis are required?

>      }
>  
>      /// Obtain `&'a T`, bypassing the [`Revocable`].

> -impl<T> Drop for Devres<T> {
> -    fn drop(&mut self) {
> +#[pinned_drop]
> +impl<T> PinnedDrop for Devres<T> {
> +    fn drop(self: Pin<&mut Self>) {
>          // SAFETY: When `drop` runs, it is guaranteed that nobody is accessing the revocable data
>          // anymore, hence it is safe not to wait for the grace period to finish.
> -        if unsafe { self.0.data.revoke_nosync() } {
> -            // We revoked `self.0.data` before the devres action did, hence try to remove it.
> -            if !DevresInner::remove_action(&self.0) {
> +        if unsafe { self.data.revoke_nosync() } {
> +            // We revoked `self.data` before the devres action did, hence try to remove it.
> +            if !self.remove_action() {
>                  // We could not remove the devres action, which means that it now runs concurrently,
> -                // hence signal that `self.0.data` has been revoked successfully.
> -                self.0.revoke.complete_all();
> +                // hence signal that `self.data` has been revoked by us successfully.
> +                self.revoke.complete_all();
> +
> +                // Wait for `Self::devres_callback` to be done using this object.
> +                self.devm.wait_for_completion();
>              }
> +        } else {
> +            // `Self::devres_callback` revokes `self.data` for us, hence wait for it to be done
> +            // using this object.
> +            self.devm.wait_for_completion();

I don't understand this change, maybe it's best to move that into a
separate commit?

---
Cheers,
Benno