rust: device: introduce Device::drvdata()

In C dev_get_drvdata() has specific requirements under which it is valid
to access the returned pointer. That is, drivers have to ensure that

(1) for the duration the returned pointer is accessed the driver is
bound and remains to be bound to the corresponding device,

(2) the returned void * is treated according to the driver's private
data type, i.e. according to what has been passed to
dev_set_drvdata().

In Rust, (1) can be ensured by simply requiring the Bound device
context, i.e. provide the drvdata() method for Device<Bound> only.

For (2) we would usually make the device type generic over the driver
type, e.g. Device<T: Driver>, where <T as Driver>::Data is the type of
the driver's private data.

However, a device does not have a driver type known at compile time and
may be bound to multiple drivers throughout its lifetime.

Hence, in order to be able to provide a safe accessor for the driver's
device private data, we have to do the type check on runtime.

This is achieved by letting a driver assert the expected type, which is
then compared to a type hash stored in struct device_private when
dev_set_drvdata() is called.

Example:

// `dev` is a `&Device<Bound>`.
let data = dev.drvdata::<SampleDriver>()?;

There are two aspects to note:

(1) Technically, the same check could be achieved by comparing the
struct device_driver pointer of struct device with the struct
device_driver pointer of the driver struct (e.g. struct
pci_driver).

However, this would - in addition the pointer comparison - require
to tie back the private driver data type to the struct
device_driver pointer of the driver struct to prove correctness.

Besides that, accessing the driver struct (stored in the module
structure) isn't trivial and would result into horrible code and
API ergonomics.

(2) Having a direct accessor to the driver's private data is not
commonly required (at least in Rust): Bus callback methods already
provide access to the driver's device private data through a &self
argument, while other driver entry points such as IRQs,
workqueues, timers, IOCTLs, etc. have their own private data with
separate ownership and lifetime.

In other words, a driver's device private data is only relevant
for driver model contexts (such a file private is only relevant
for file contexts).

Having that said, the motivation for accessing the driver's device
private data with Device<Bound>::drvdata() are interactions between
drivers. For instance, when an auxiliary driver calls back into its
parent, the parent has to be capable to derive its private data from the
corresponding device (i.e. the parent of the auxiliary device).

Reviewed-by: Alice Ryhl <[e-mail]>
Reviewed-by: Greg Kroah-Hartman <[e-mail]>
[ * Remove unnecessary `const _: ()` block,
* rename type_id_{store,match}() to {set,match}_type_id(),
* assert size_of::<bindings::driver_type>() >= size_of::<TypeId>(),
* add missing check in case Device::drvdata() is called from probe().

- Danilo ]
Signed-off-by: Danilo Krummrich <[e-mail]>

drivers/base/base.h

···
 };
 #define to_driver(obj) container_of(obj, struct driver_private, kobj)
 
+#ifdef CONFIG_RUST
+/**
+ * struct driver_type - Representation of a Rust driver type.
+ */
+struct driver_type {
+	/**
+	 * @id: Representation of core::any::TypeId.
+	 */
+	u8 id[16];
+} __packed;
+#endif
+
 /**
  * struct device_private - structure to hold the private to the driver core portions of the device structure.
  *
···
  * @async_driver - pointer to device driver awaiting probe via async_probe
  * @device - pointer back to the struct device that this structure is
  * associated with.
+ * @driver_type - The type of the bound Rust driver.
  * @dead - This device is currently either in the process of or has been
  *	removed from the system. Any asynchronous events scheduled for this
  *	device should exit without taking any action.
···
 	const struct device_driver *async_driver;
 	char *deferred_probe_reason;
 	struct device *device;
+#ifdef CONFIG_RUST
+	struct driver_type driver_type;
+#endif
 	u8 dead:1;
 };
 #define to_device_private_parent(obj)	\

rust/bindings/bindings_helper.h

···
 #include <linux/xarray.h>
 #include <trace/events/rust_sample.h>
 
+/*
+ * The driver-core Rust code needs to know about some C driver-core private
+ * structures.
+ */
+#include <../../drivers/base/base.h>
+
 #if defined(CONFIG_DRM_PANIC_SCREEN_QR_CODE)
 // Used by `#[export]` in `drivers/gpu/drm/drm_panic_qr.rs`.
 #include <drm/drm_panic.h>

rust/kernel/device.rs

···
     sync::aref::ARef,
     types::{ForeignOwnable, Opaque},
 };
-use core::{marker::PhantomData, ptr};
+use core::{any::TypeId, marker::PhantomData, ptr};
 
 #[cfg(CONFIG_PRINTK)]
 use crate::c_str;
 
 pub mod property;
+
+// Assert that we can `read()` / `write()` a `TypeId` instance from / into `struct driver_type`.
+static_assert!(core::mem::size_of::<bindings::driver_type>() >= core::mem::size_of::<TypeId>());
 
 /// The core representation of a device in the kernel's driver model.
 ///
···
 }
 
 impl Device<CoreInternal> {
+    fn set_type_id<T: 'static>(&self) {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        let private = unsafe { (*self.as_raw()).p };
+
+        // SAFETY: For a bound device (implied by the `CoreInternal` device context), `private` is
+        // guaranteed to be a valid pointer to a `struct device_private`.
+        let driver_type = unsafe { &raw mut (*private).driver_type };
+
+        // SAFETY: `driver_type` is valid for (unaligned) writes of a `TypeId`.
+        unsafe {
+            driver_type
+                .cast::<TypeId>()
+                .write_unaligned(TypeId::of::<T>())
+        };
+    }
+
     /// Store a pointer to the bound driver's private data.
     pub fn set_drvdata<T: 'static>(&self, data: impl PinInit<T, Error>) -> Result {
         let data = KBox::pin_init(data, GFP_KERNEL)?;
 
         // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
         unsafe { bindings::dev_set_drvdata(self.as_raw(), data.into_foreign().cast()) };
+        self.set_type_id::<T>();
 
         Ok(())
     }
···
     pub unsafe fn drvdata_obtain<T: 'static>(&self) -> Pin<KBox<T>> {
         // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
         let ptr = unsafe { bindings::dev_get_drvdata(self.as_raw()) };
+
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        unsafe { bindings::dev_set_drvdata(self.as_raw(), core::ptr::null_mut()) };
 
         // SAFETY:
         // - By the safety requirements of this function, `ptr` comes from a previous call to
···
     ///   [`Device::drvdata_obtain`].
     /// - The type `T` must match the type of the `ForeignOwnable` previously stored by
     ///   [`Device::set_drvdata`].
-    pub unsafe fn drvdata_borrow<T: ForeignOwnable>(&self) -> T::Borrowed<'_> {
+    pub unsafe fn drvdata_borrow<T: 'static>(&self) -> Pin<&T> {
+        // SAFETY: `drvdata_unchecked()` has the exact same safety requirements as the ones
+        // required by this method.
+        unsafe { self.drvdata_unchecked() }
+    }
+}
+
+impl Device<Bound> {
+    /// Borrow the driver's private data bound to this [`Device`].
+    ///
+    /// # Safety
+    ///
+    /// - Must only be called after a preceding call to [`Device::set_drvdata`] and before
+    ///   [`Device::drvdata_obtain`].
+    /// - The type `T` must match the type of the `ForeignOwnable` previously stored by
+    ///   [`Device::set_drvdata`].
+    unsafe fn drvdata_unchecked<T: 'static>(&self) -> Pin<&T> {
         // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
         let ptr = unsafe { bindings::dev_get_drvdata(self.as_raw()) };
 
···
         //   `into_foreign()`.
         // - `dev_get_drvdata()` guarantees to return the same pointer given to `dev_set_drvdata()`
         //   in `into_foreign()`.
-        unsafe { T::borrow(ptr.cast()) }
+        unsafe { Pin::<KBox<T>>::borrow(ptr.cast()) }
+    }
+
+    fn match_type_id<T: 'static>(&self) -> Result {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        let private = unsafe { (*self.as_raw()).p };
+
+        // SAFETY: For a bound device, `private` is guaranteed to be a valid pointer to a
+        // `struct device_private`.
+        let driver_type = unsafe { &raw mut (*private).driver_type };
+
+        // SAFETY:
+        // - `driver_type` is valid for (unaligned) reads of a `TypeId`.
+        // - A bound device guarantees that `driver_type` contains a valid `TypeId` value.
+        let type_id = unsafe { driver_type.cast::<TypeId>().read_unaligned() };
+
+        if type_id != TypeId::of::<T>() {
+            return Err(EINVAL);
+        }
+
+        Ok(())
+    }
+
+    /// Access a driver's private data.
+    ///
+    /// Returns a pinned reference to the driver's private data or [`EINVAL`] if it doesn't match
+    /// the asserted type `T`.
+    pub fn drvdata<T: 'static>(&self) -> Result<Pin<&T>> {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        if unsafe { bindings::dev_get_drvdata(self.as_raw()) }.is_null() {
+            return Err(ENOENT);
+        }
+
+        self.match_type_id::<T>()?;
+
+        // SAFETY:
+        // - The above check of `dev_get_drvdata()` guarantees that we are called after
+        //   `set_drvdata()` and before `drvdata_obtain()`.
+        // - We've just checked that the type of the driver's private data is in fact `T`.
+        Ok(unsafe { self.drvdata_unchecked() })
     }
 }