From 269b6bd9838ba7539017de4e238a59c072a4db0a Mon Sep 17 00:00:00 2001
From: Wedson Almeida Filho <wedsonaf@google.com>
Date: Sun, 9 Jan 2022 06:09:46 +0000
Subject: [PATCH] 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>
---
 rust/kernel/amba.rs   |  4 +++-
 rust/kernel/driver.rs | 16 +++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/rust/kernel/amba.rs b/rust/kernel/amba.rs
index 689528502ca5ae..4ebf86028c709e 100644
--- a/rust/kernel/amba.rs
+++ b/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>);
diff --git a/rust/kernel/driver.rs b/rust/kernel/driver.rs
index d72ada58138c7e..f6b02d7ec0cbbc 100644
--- a/rust/kernel/driver.rs
+++ b/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,9 +86,9 @@ 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)
     }
 
@@ -95,7 +96,11 @@ impl<T: DriverOps> Registration<T> {
     ///
     /// 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)?,
         })
     }
 }