rust: init: add `PinnedDrop` trait and macros

The `PinnedDrop` trait that facilitates destruction of pinned types.
It has to be implemented via the `#[pinned_drop]` macro, since the
`drop` function should not be called by normal code, only by other
destructors. It also only works on structs that are annotated with
`#[pin_data(PinnedDrop)]`.
Co-developed-by: Gary Guo <gary@garyguo.net>
Signed-off-by: Gary Guo <gary@garyguo.net>
Signed-off-by: Benno Lossin <benno.lossin@proton.me>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Andreas Hindborg <a.hindborg@samsung.com>
Link: https://lore.kernel.org/r/20230408122429.1103522-10-y86-dev@protonmail.comSigned-off-by: Miguel Ojeda <ojeda@kernel.org>

rust/kernel/init.rs

@@ -104,6 +104,78 @@
 //! }
 //! ```
 //!
+//! ## Manual creation of an initializer
+//!
+//! Often when working with primitives the previous approaches are not sufficient. That is where
+//! [`pin_init_from_closure()`] comes in. This `unsafe` function allows you to create a
+//! [`impl PinInit<T, E>`] directly from a closure. Of course you have to ensure that the closure
+//! actually does the initialization in the correct way. Here are the things to look out for
+//! (we are calling the parameter to the closure `slot`):
+//! - when the closure returns `Ok(())`, then it has completed the initialization successfully, so
+//!   `slot` now contains a valid bit pattern for the type `T`,
+//! - when the closure returns `Err(e)`, then the caller may deallocate the memory at `slot`, so
+//!   you need to take care to clean up anything if your initialization fails mid-way,
+//! - you may assume that `slot` will stay pinned even after the closure returns until `drop` of
+//!   `slot` gets called.
+//!
+//! ```rust
+//! use kernel::{prelude::*, init};
+//! use core::{ptr::addr_of_mut, marker::PhantomPinned, pin::Pin};
+//! # mod bindings {
+//! #     pub struct foo;
+//! #     pub unsafe fn init_foo(_ptr: *mut foo) {}
+//! #     pub unsafe fn destroy_foo(_ptr: *mut foo) {}
+//! #     pub unsafe fn enable_foo(_ptr: *mut foo, _flags: u32) -> i32 { 0 }
+//! # }
+//! /// # Invariants
+//! ///
+//! /// `foo` is always initialized
+//! #[pin_data(PinnedDrop)]
+//! pub struct RawFoo {
+//!     #[pin]
+//!     foo: Opaque<bindings::foo>,
+//!     #[pin]
+//!     _p: PhantomPinned,
+//! }
+//!
+//! impl RawFoo {
+//!     pub fn new(flags: u32) -> impl PinInit<Self, Error> {
+//!         // SAFETY:
+//!         // - when the closure returns `Ok(())`, then it has successfully initialized and
+//!         //   enabled `foo`,
+//!         // - when it returns `Err(e)`, then it has cleaned up before
+//!         unsafe {
+//!             init::pin_init_from_closure(move |slot: *mut Self| {
+//!                 // `slot` contains uninit memory, avoid creating a reference.
+//!                 let foo = addr_of_mut!((*slot).foo);
+//!
+//!                 // Initialize the `foo`
+//!                 bindings::init_foo(Opaque::raw_get(foo));
+//!
+//!                 // Try to enable it.
+//!                 let err = bindings::enable_foo(Opaque::raw_get(foo), flags);
+//!                 if err != 0 {
+//!                     // Enabling has failed, first clean up the foo and then return the error.
+//!                     bindings::destroy_foo(Opaque::raw_get(foo));
+//!                     return Err(Error::from_kernel_errno(err));
+//!                 }
+//!
+//!                 // All fields of `RawFoo` have been initialized, since `_p` is a ZST.
+//!                 Ok(())
+//!             })
+//!         }
+//!     }
+//! }
+//!
+//! #[pinned_drop]
+//! impl PinnedDrop for RawFoo {
+//!     fn drop(self: Pin<&mut Self>) {
+//!         // SAFETY: Since `foo` is initialized, destroying is safe.
+//!         unsafe { bindings::destroy_foo(self.foo.get()) };
+//!     }
+//! }
+//! ```
+//!
 //! [`sync`]: kernel::sync
 //! [pinning]: https://doc.rust-lang.org/std/pin/index.html
 //! [structurally pinned fields]:
@@ -1084,3 +1156,42 @@ impl<T> InPlaceInit<T> for UniqueArc<T> {
         Ok(unsafe { this.assume_init() })
     }
 }
+
+/// Trait facilitating pinned destruction.
+///
+/// Use [`pinned_drop`] to implement this trait safely:
+///
+/// ```rust
+/// # use kernel::sync::Mutex;
+/// use kernel::macros::pinned_drop;
+/// use core::pin::Pin;
+/// #[pin_data(PinnedDrop)]
+/// struct Foo {
+///     #[pin]
+///     mtx: Mutex<usize>,
+/// }
+///
+/// #[pinned_drop]
+/// impl PinnedDrop for Foo {
+///     fn drop(self: Pin<&mut Self>) {
+///         pr_info!("Foo is being dropped!");
+///     }
+/// }
+/// ```
+///
+/// # Safety
+///
+/// This trait must be implemented via the [`pinned_drop`] proc-macro attribute on the impl.
+///
+/// [`pinned_drop`]: kernel::macros::pinned_drop
+pub unsafe trait PinnedDrop: __internal::HasPinData {
+    /// Executes the pinned destructor of this type.
+    ///
+    /// While this function is marked safe, it is actually unsafe to call it manually. For this
+    /// reason it takes an additional parameter. This type can only be constructed by `unsafe` code
+    /// and thus prevents this function from being called where it should not.
+    ///
+    /// This extra parameter will be generated by the `#[pinned_drop]` proc-macro attribute
+    /// automatically.
+    fn drop(self: Pin<&mut Self>, only_call_from_drop: __internal::OnlyCallFromDrop);
+}

rust/kernel/init/__internal.rs

@@ -161,3 +161,18 @@ impl<T: ?Sized> Drop for DropGuard<T> {
         }
     }
 }
+
+/// Token used by `PinnedDrop` to prevent calling the function without creating this unsafely
+/// created struct. This is needed, because the `drop` function is safe, but should not be called
+/// manually.
+pub struct OnlyCallFromDrop(());
+
+impl OnlyCallFromDrop {
+    /// # Safety
+    ///
+    /// This function should only be called from the [`Drop::drop`] function and only be used to
+    /// delegate the destruction to the pinned destructor [`PinnedDrop::drop`] of the same type.
+    pub unsafe fn new() -> Self {
+        Self(())
+    }
+}

rust/macros/lib.rs

@@ -8,6 +8,7 @@ mod concat_idents;
 mod helpers;
 mod module;
 mod pin_data;
+mod pinned_drop;
 mod vtable;
 
 use proc_macro::TokenStream;
@@ -180,6 +181,10 @@ pub fn concat_idents(ts: TokenStream) -> TokenStream {
 /// This macro enables the use of the [`pin_init!`] macro. When pin-initializing a `struct`,
 /// then `#[pin]` directs the type of initializer that is required.
 ///
+/// If your `struct` implements `Drop`, then you need to add `PinnedDrop` as arguments to this
+/// macro, and change your `Drop` implementation to `PinnedDrop` annotated with
+/// `#[`[`macro@pinned_drop`]`]`, since dropping pinned values requires extra care.
+///
 /// # Examples
 ///
 /// ```rust,ignore
@@ -191,9 +196,53 @@ pub fn concat_idents(ts: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
+/// ```rust,ignore
+/// #[pin_data(PinnedDrop)]
+/// struct DriverData {
+///     #[pin]
+///     queue: Mutex<Vec<Command>>,
+///     buf: Box<[u8; 1024 * 1024]>,
+///     raw_info: *mut Info,
+/// }
+///
+/// #[pinned_drop]
+/// impl PinnedDrop for DriverData {
+///     fn drop(self: Pin<&mut Self>) {
+///         unsafe { bindings::destroy_info(self.raw_info) };
+///     }
+/// }
+/// ```
+///
 /// [`pin_init!`]: ../kernel/macro.pin_init.html
 //  ^ cannot use direct link, since `kernel` is not a dependency of `macros`.
 #[proc_macro_attribute]
 pub fn pin_data(inner: TokenStream, item: TokenStream) -> TokenStream {
     pin_data::pin_data(inner, item)
 }
+
+/// Used to implement `PinnedDrop` safely.
+///
+/// Only works on structs that are annotated via `#[`[`macro@pin_data`]`]`.
+///
+/// # Examples
+///
+/// ```rust,ignore
+/// #[pin_data(PinnedDrop)]
+/// struct DriverData {
+///     #[pin]
+///     queue: Mutex<Vec<Command>>,
+///     buf: Box<[u8; 1024 * 1024]>,
+///     raw_info: *mut Info,
+/// }
+///
+/// #[pinned_drop]
+/// impl PinnedDrop for DriverData {
+///     fn drop(self: Pin<&mut Self>) {
+///         unsafe { bindings::destroy_info(self.raw_info) };
+///     }
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn pinned_drop(args: TokenStream, input: TokenStream) -> TokenStream {
+    pinned_drop::pinned_drop(args, input)
+}

rust/macros/pinned_drop.rs

+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+use proc_macro::{TokenStream, TokenTree};
+
+pub(crate) fn pinned_drop(_args: TokenStream, input: TokenStream) -> TokenStream {
+    let mut toks = input.into_iter().collect::<Vec<_>>();
+    assert!(!toks.is_empty());
+    // Ensure that we have an `impl` item.
+    assert!(matches!(&toks[0], TokenTree::Ident(i) if i.to_string() == "impl"));
+    // Ensure that we are implementing `PinnedDrop`.
+    let mut nesting: usize = 0;
+    let mut pinned_drop_idx = None;
+    for (i, tt) in toks.iter().enumerate() {
+        match tt {
+            TokenTree::Punct(p) if p.as_char() == '<' => {
+                nesting += 1;
+            }
+            TokenTree::Punct(p) if p.as_char() == '>' => {
+                nesting = nesting.checked_sub(1).unwrap();
+                continue;
+            }
+            _ => {}
+        }
+        if i >= 1 && nesting == 0 {
+            // Found the end of the generics, this should be `PinnedDrop`.
+            assert!(
+                matches!(tt, TokenTree::Ident(i) if i.to_string() == "PinnedDrop"),
+                "expected 'PinnedDrop', found: '{:?}'",
+                tt
+            );
+            pinned_drop_idx = Some(i);
+            break;
+        }
+    }
+    let idx = pinned_drop_idx
+        .unwrap_or_else(|| panic!("Expected an `impl` block implementing `PinnedDrop`."));
+    // Fully qualify the `PinnedDrop`, as to avoid any tampering.
+    toks.splice(idx..idx, quote!(::kernel::init::));
+    // Take the `{}` body and call the declarative macro.
+    if let Some(TokenTree::Group(last)) = toks.pop() {
+        let last = last.stream();
+        quote!(::kernel::__pinned_drop! {
+            @impl_sig(#(#toks)*),
+            @impl_body(#last),
+        })
+    } else {
+        TokenStream::from_iter(toks)
+    }
+}