rust: drm: device: Convert Device to AlwaysRefCounted

hoshinolina · hoshinolina · commit a5ddd9d2fc72 · 2023-05-19T22:21:10.000+09:00
Switch from being a refcount wrapper itself to a transparent wrapper
around `bindings::drm_device`. The refcounted type then becomes
ARef&lt;Device&lt;T&gt;&gt;.

Signed-off-by: Asahi Lina &lt;lina@asahilina.net&gt;
diff --git a/rust/kernel/drm/device.rs b/rust/kernel/drm/device.rs
@@ -4,73 +4,66 @@
 //!
 //! C header: [`include/linux/drm/drm_device.h`](../../../../include/linux/drm/drm_device.h)
 
-use crate::{bindings, device, drm, types::ForeignOwnable};
+use crate::{
+    bindings, device, drm,
+    types::{ARef, AlwaysRefCounted, ForeignOwnable},
+};
+use core::cell::UnsafeCell;
 use core::marker::PhantomData;
+use core::ptr::NonNull;
 
-/// Represents a reference to a DRM device. The device is reference-counted and is guaranteed to
-/// not be dropped while this object is alive.
+/// A typed DRM device with a specific driver. The device is always reference-counted.
+#[repr(transparent)]
 pub struct Device<T: drm::drv::Driver> {
-    // Type invariant: ptr must be a valid and initialized drm_device,
-    // and this value must either own a reference to it or the caller
-    // must ensure that it is never dropped if the reference is borrowed.
-    pub(super) ptr: *mut bindings::drm_device,
+    pub(super) drm: UnsafeCell<bindings::drm_device>,
     _p: PhantomData<T>,
 }
 
 impl<T: drm::drv::Driver> Device<T> {
-    // Not intended to be called externally, except via declare_drm_ioctls!()
-    #[doc(hidden)]
-    pub unsafe fn from_raw(raw: *mut bindings::drm_device) -> Device<T> {
-        Device {
-            ptr: raw,
-            _p: PhantomData,
-        }
-    }
-
-    #[allow(dead_code)]
-    pub(crate) fn raw(&self) -> *const bindings::drm_device {
-        self.ptr
+    #[allow(dead_code, clippy::mut_from_ref)]
+    pub(crate) unsafe fn raw_mut(&self) -> &mut bindings::drm_device {
+        unsafe { &mut *self.drm.get() }
     }
 
-    pub(crate) fn raw_mut(&mut self) -> *mut bindings::drm_device {
-        self.ptr
+    // Not intended to be called externally, except via declare_drm_ioctls!()
+    #[doc(hidden)]
+    pub unsafe fn borrow<'a>(raw: *const bindings::drm_device) -> &'a Self {
+        unsafe { &*(raw as *const Self) }
     }
 
     /// Returns a borrowed reference to the user data associated with this Device.
     pub fn data(&self) -> <T::Data as ForeignOwnable>::Borrowed<'_> {
-        unsafe { T::Data::borrow((*self.ptr).dev_private) }
+        // SAFETY: dev_private is guaranteed to be initialized for all
+        // Device objects exposed to users.
+        unsafe { T::Data::borrow((*self.drm.get()).dev_private) }
     }
 }
 
-impl<T: drm::drv::Driver> Drop for Device<T> {
-    fn drop(&mut self) {
-        // SAFETY: By the type invariants, we know that `self` owns a reference, so it is safe to
-        // relinquish it now.
-        unsafe { bindings::drm_dev_put(self.ptr) };
+// SAFETY: DRM device objects are always reference counted and the get/put functions
+// satisfy the requirements.
+unsafe impl<T: drm::drv::Driver> AlwaysRefCounted for Device<T> {
+    fn inc_ref(&self) {
+        unsafe { bindings::drm_dev_get(&self.drm as *const _ as *mut _) };
     }
-}
 
-impl<T: drm::drv::Driver> Clone for Device<T> {
-    fn clone(&self) -> Self {
-        // SAFETY: We get a new reference and then create a new owning object from the raw pointer
-        unsafe {
-            bindings::drm_dev_get(self.ptr);
-            Device::from_raw(self.ptr)
-        }
+    unsafe fn dec_ref(obj: NonNull<Self>) {
+        // SAFETY: The Device<T> type has the same layout as drm_device,
+        // so we can just cast.
+        unsafe { bindings::drm_dev_put(obj.as_ptr() as *mut _) };
     }
 }
 
 // SAFETY: `Device` only holds a pointer to a C device, which is safe to be used from any thread.
-unsafe impl<T: drm::drv::Driver> Send for Device<T> {}
+unsafe impl<T: drm::drv::Driver> Send for ARef<Device<T>> {}
 
 // SAFETY: `Device` only holds a pointer to a C device, references to which are safe to be used
 // from any thread.
-unsafe impl<T: drm::drv::Driver> Sync for Device<T> {}
+unsafe impl<T: drm::drv::Driver> Sync for ARef<Device<T>> {}
 
 // Make drm::Device work for dev_info!() and friends
 unsafe impl<T: drm::drv::Driver> device::RawDevice for Device<T> {
     fn raw_device(&self) -> *mut bindings::device {
-        // SAFETY: ptr must be valid per the type invariant
-        unsafe { (*self.ptr).dev }
+        // SAFETY: dev is initialized by C for all Device objects
+        unsafe { (*self.drm.get()).dev }
     }
 }
diff --git a/rust/kernel/drm/drv.rs b/rust/kernel/drm/drv.rs
@@ -12,12 +12,13 @@ use crate::{
     prelude::*,
     private::Sealed,
     str::CStr,
-    types::ForeignOwnable,
+    types::{ARef, ForeignOwnable},
     ThisModule,
 };
 use core::{
     marker::{PhantomData, PhantomPinned},
     pin::Pin,
+    ptr::NonNull,
 };
 use macros::vtable;
 
@@ -163,7 +164,7 @@ pub trait Driver {
 ///
 /// drm is always a valid pointer to an allocated drm_device
 pub struct Registration<T: Driver> {
-    drm: drm::device::Device<T>,
+    drm: ARef<drm::device::Device<T>>,
     registered: bool,
     fops: bindings::file_operations,
     vtable: Pin<Box<bindings::drm_driver>>,
@@ -253,11 +254,11 @@ impl<T: Driver> Registration<T> {
     pub fn new(parent: &dyn device::RawDevice) -> Result<Self> {
         let vtable = Pin::new(Box::try_new(Self::VTABLE)?);
         let raw_drm = unsafe { bindings::drm_dev_alloc(&*vtable, parent.raw_device()) };
-        let raw_drm = from_err_ptr(raw_drm)?;
+        let raw_drm = NonNull::new(from_err_ptr(raw_drm)? as *mut _).ok_or(ENOMEM)?;
 
         // The reference count is one, and now we take ownership of that reference as a
         // drm::device::Device.
-        let drm = unsafe { drm::device::Device::from_raw(raw_drm) };
+        let drm = unsafe { ARef::from_raw(raw_drm) };
 
         Ok(Self {
             drm,
@@ -287,10 +288,9 @@ impl<T: Driver> Registration<T> {
         // SAFETY: We never move out of `this`.
         let this = unsafe { self.get_unchecked_mut() };
         let data_pointer = <T::Data as ForeignOwnable>::into_foreign(data);
-        // SAFETY: `drm` is valid per the type invariant
-        unsafe {
-            (*this.drm.raw_mut()).dev_private = data_pointer as *mut _;
-        }
+        // SAFETY: This is the only code touching dev_private, so it is safe to upgrade to a
+        // mutable reference.
+        unsafe { this.drm.raw_mut() }.dev_private = data_pointer as *mut _;
 
         this.fops.owner = module.0;
         this.vtable.fops = &this.fops;
@@ -309,6 +309,7 @@ impl<T: Driver> Registration<T> {
 
     /// Returns a reference to the `Device` instance for this registration.
     pub fn device(&self) -> &drm::device::Device<T> {
+        // TODO: rework this, ensure this only works after registration
         &self.drm
     }
 }
@@ -329,7 +330,7 @@ impl<T: Driver> Drop for Registration<T> {
         if self.registered {
             // Get a pointer to the data stored in device before destroying it.
             // SAFETY: `drm` is valid per the type invariant
-            let data_pointer = unsafe { (*self.drm.raw_mut()).dev_private };
+            let data_pointer = unsafe { self.drm.raw_mut().dev_private };
 
             // SAFETY: Since `registered` is true, `self.drm` is both valid and registered.
             unsafe { bindings::drm_dev_unregister(self.drm.raw_mut()) };
diff --git a/rust/kernel/drm/file.rs b/rust/kernel/drm/file.rs
@@ -32,11 +32,11 @@ pub(super) unsafe extern "C" fn open_callback<T: DriverFile>(
     raw_dev: *mut bindings::drm_device,
     raw_file: *mut bindings::drm_file,
 ) -> core::ffi::c_int {
-    let drm = core::mem::ManuallyDrop::new(unsafe { drm::device::Device::from_raw(raw_dev) });
+    let drm = unsafe { drm::device::Device::borrow(raw_dev) };
     // SAFETY: This reference won't escape this function
     let file = unsafe { &mut *raw_file };
 
-    let inner = match T::open(&drm) {
+    let inner = match T::open(drm) {
         Err(e) => {
             return e.to_errno();
         }
diff --git a/rust/kernel/drm/gem/mod.rs b/rust/kernel/drm/gem/mod.rs
@@ -15,7 +15,7 @@ use crate::{
     error::{to_result, Result},
     prelude::*,
 };
-use core::{marker::PhantomPinned, mem, mem::ManuallyDrop, ops::Deref, ops::DerefMut};
+use core::{marker::PhantomPinned, mem, ops::Deref, ops::DerefMut};
 
 /// GEM object functions, which must be implemented by drivers.
 pub trait BaseDriverObject<T: BaseObject>: Sync + Send + Sized {
@@ -211,9 +211,7 @@ impl<T: IntoGEMObject> BaseObject for T {}
 #[pin_data]
 pub struct Object<T: DriverObject> {
     obj: bindings::drm_gem_object,
-    // The DRM core ensures the Device exists as long as its objects exist, so we don't need to
-    // manage the reference count here.
-    dev: ManuallyDrop<device::Device<T::Driver>>,
+    dev: *const bindings::drm_device,
     #[pin]
     inner: T,
     #[pin]
@@ -251,14 +249,12 @@ impl<T: DriverObject> Object<T> {
                 ..Default::default()
             },
             inner <- T::new(dev, size),
-            // SAFETY: The drm subsystem guarantees that the drm_device will live as long as
-            // the GEM object lives, so we can conjure a reference out of thin air.
-            dev: ManuallyDrop::new(unsafe { device::Device::from_raw(dev.ptr) }),
+            dev: dev.drm.get(),
             _p: PhantomPinned
         }))?;
 
         to_result(unsafe {
-            bindings::drm_gem_object_init(dev.raw() as *mut _, &obj.obj as *const _ as *mut _, size)
+            bindings::drm_gem_object_init(dev.raw_mut(), &obj.obj as *const _ as *mut _, size)
         })?;
 
         // SAFETY: We never move out of self
@@ -275,7 +271,9 @@ impl<T: DriverObject> Object<T> {
 
     /// Returns the `Device` that owns this GEM object.
     pub fn dev(&self) -> &device::Device<T::Driver> {
-        &self.dev
+        // SAFETY: The drm subsystem guarantees that the drm_device will live as long as
+        // the GEM object lives, so we can just borrow from the raw pointer.
+        unsafe { device::Device::borrow(self.dev) }
     }
 }
 
diff --git a/rust/kernel/drm/gem/shmem.rs b/rust/kernel/drm/gem/shmem.rs
@@ -12,7 +12,7 @@ use crate::{
 use core::{
     marker::{PhantomData, PhantomPinned},
     mem,
-    mem::{ManuallyDrop, MaybeUninit},
+    mem::MaybeUninit,
     ops::{Deref, DerefMut},
     slice,
 };
@@ -71,7 +71,7 @@ pub struct Object<T: DriverObject> {
     obj: bindings::drm_gem_shmem_object,
     // The DRM core ensures the Device exists as long as its objects exist, so we don't need to
     // manage the reference count here.
-    dev: ManuallyDrop<device::Device<T::Driver>>,
+    dev: *const bindings::drm_device,
     #[pin]
     inner: T,
 }
@@ -80,13 +80,9 @@ pub struct Object<T: DriverObject> {
 unsafe impl init::Zeroable for bindings::drm_gem_shmem_object {}
 
 unsafe extern "C" fn gem_create_object<T: DriverObject>(
-    raw_dev: *mut bindings::drm_device,
+    dev: *mut bindings::drm_device,
     size: usize,
 ) -> *mut bindings::drm_gem_object {
-    // SAFETY: GEM ensures the device lives as long as its objects live,
-    // so we can conjure up a reference from thin air and never drop it.
-    let dev = ManuallyDrop::new(unsafe { device::Device::from_raw(raw_dev) });
-
     let p = unsafe {
         bindings::krealloc(core::ptr::null(), Object::<T>::SIZE, bindings::GFP_KERNEL)
             as *mut Object<T>
@@ -98,7 +94,8 @@ unsafe extern "C" fn gem_create_object<T: DriverObject>(
 
     let init = try_pin_init!(Object {
         obj <- init::zeroed(),
-        inner <- T::new(&*dev, size),
+        // SAFETY: GEM ensures the device lives as long as its objects live
+        inner <- T::new(unsafe { device::Device::borrow(dev)}, size),
         dev,
     });
 
@@ -163,7 +160,7 @@ impl<T: DriverObject> Object<T> {
     pub fn new(dev: &device::Device<T::Driver>, size: usize) -> Result<gem::UniqueObjectRef<Self>> {
         // SAFETY: This function can be called as long as the ALLOC_OPS are set properly
         // for this driver, and the gem_create_object is called.
-        let p = unsafe { bindings::drm_gem_shmem_create(dev.raw() as *mut _, size) };
+        let p = unsafe { bindings::drm_gem_shmem_create(dev.raw_mut(), size) };
         let p = crate::container_of!(p, Object<T>, obj) as *mut _;
 
         // SAFETY: The gem_create_object callback ensures this is a valid Object<T>,
@@ -178,7 +175,9 @@ impl<T: DriverObject> Object<T> {
 
     /// Returns the `Device` that owns this GEM object.
     pub fn dev(&self) -> &device::Device<T::Driver> {
-        &self.dev
+        // SAFETY: GEM ensures that the device outlives its objects, so we can
+        // just borrow here.
+        unsafe { device::Device::borrow(self.dev) }
     }
 
     /// Creates (if necessary) and returns a scatter-gather table of DMA pages for this object.
diff --git a/rust/kernel/drm/ioctl.rs b/rust/kernel/drm/ioctl.rs
@@ -128,22 +128,20 @@ macro_rules! declare_drm_ioctls {
                                 raw_data: *mut ::core::ffi::c_void,
                                 raw_file_priv: *mut $crate::drm::ioctl::internal::drm_file,
                         ) -> core::ffi::c_int {
-                            // SAFETY: We never drop this, and the DRM core ensures the device lives
-                            // while callbacks are being called.
+                            // SAFETY: The DRM core ensures the device lives while callbacks are
+                            // being called.
                             //
                             // FIXME: Currently there is nothing enforcing that the types of the
                             // dev/file match the current driver these ioctls are being declared
                             // for, and it's not clear how to enforce this within the type system.
-                            let dev = ::core::mem::ManuallyDrop::new(unsafe {
-                                $crate::drm::device::Device::from_raw(raw_dev)
-                            });
+                            let dev = $crate::drm::device::Device::borrow(raw_dev);
                             // SAFETY: This is just the ioctl argument, which hopefully has the right type
                             // (we've done our best checking the size).
                             let data = unsafe { &mut *(raw_data as *mut $crate::uapi::$struct) };
                             // SAFETY: This is just the DRM file structure
                             let file = unsafe { $crate::drm::file::File::from_raw(raw_file_priv) };
 
-                            match $func(&*dev, data, &file) {
+                            match $func(dev, data, &file) {
                                 Err(e) => e.to_errno(),
                                 Ok(i) => i.try_into().unwrap_or(ERANGE.to_errno()),
                             }
