rust: add ThisModule to the arguments of driver registration

So that the kernel is aware of any loadable modules hosting drivers and
can therefore increment/decrement refcounts as appropriate to prevent
unload when the driver is in use.

This should be transparent to all drivers registered with macros. Those
registered with direct calls will need to pass the extra argument.

Signed-off-by: Wedson Almeida Filho <wedsonaf@google.com>

Author: wedsonaf

rust/kernel/amba.rs

@@ -6,7 +6,7 @@
 
 use crate::{
     bindings, c_types, device, driver, error::from_kernel_result, io_mem::Resource, power,
-    str::CStr, to_result, types::PointerWrapper, Error, Result,
+    str::CStr, to_result, types::PointerWrapper, Error, Result, ThisModule,
 };
 use core::{marker::PhantomData, ops::Deref};
 
@@ -72,12 +72,14 @@ where
     unsafe fn register(
         reg: *mut bindings::amba_driver,
         name: &'static CStr,
+        module: &'static ThisModule,
         id_table: *const bindings::amba_id,
     ) -> Result {
         // SAFETY: By the safety requirements of this function (defined in the trait defintion),
         // `reg` is non-null and valid.
         let amba = unsafe { &mut *reg };
         amba.drv.name = name.as_char_ptr();
+        amba.drv.owner = module.0;
         amba.id_table = id_table;
         amba.probe = Some(probe_callback::<T>);
         amba.remove = Some(remove_callback::<T>);

rust/kernel/driver.rs

@@ -44,6 +44,7 @@ pub trait DriverOps {
     unsafe fn register(
         reg: *mut Self::RegType,
         name: &'static CStr,
+        module: &'static ThisModule,
         id_table: *const Self::RawIdType,
     ) -> Result;
 
@@ -85,17 +86,21 @@ impl<T: DriverOps> Registration<T> {
     /// Allocates a pinned registration object and registers it.
     ///
     /// Returns a pinned heap-allocated representation of the registration.
-    pub fn new_pinned(name: &'static CStr) -> Result<Pin<Box<Self>>> {
+    pub fn new_pinned(name: &'static CStr, module: &'static ThisModule) -> Result<Pin<Box<Self>>> {
         let mut reg = Pin::from(Box::try_new(Self::new())?);
-        reg.as_mut().register(name)?;
+        reg.as_mut().register(name, module)?;
         Ok(reg)
     }
 
     /// Registers a driver with its subsystem.
     ///
     /// It must be pinned because the memory block that represents the registration is potentially
     /// self-referential.
-    pub fn register(self: Pin<&mut Self>, name: &'static CStr) -> Result {
+    pub fn register(
+        self: Pin<&mut Self>,
+        name: &'static CStr,
+        module: &'static ThisModule,
+    ) -> Result {
         // SAFETY: We never move out of `this`.
         let this = unsafe { self.get_unchecked_mut() };
         if this.is_registered {
@@ -114,6 +119,7 @@ impl<T: DriverOps> Registration<T> {
             T::register(
                 this.concrete_reg.get(),
                 name,
+                module,
                 &this.id_table[0] as *const _ as *const _,
             )
         }?;
@@ -174,9 +180,9 @@ pub struct Module<T: DriverOps> {
 }
 
 impl<T: DriverOps> KernelModule for Module<T> {
-    fn init(name: &'static CStr, _module: &'static ThisModule) -> Result<Self> {
+    fn init(name: &'static CStr, module: &'static ThisModule) -> Result<Self> {
         Ok(Self {
-            _driver: Registration::new_pinned(name)?,
+            _driver: Registration::new_pinned(name, module)?,
         })
     }
 }