From 6e2328448458264fd0bd872bd65d8ddde5e43db0 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Fri, 22 Jan 2021 10:42:43 -0500
Subject: [PATCH] Implement THIS_MODULE.

Closes #15
Fixes #10
---
 rust/kernel/chrdev.rs |  5 ++---
 rust/kernel/lib.rs    | 12 ++++++++++++
 rust/module.rs        |  5 +++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/chrdev.rs b/rust/kernel/chrdev.rs
index d32b9a0d584a56..52110d5c7b9a15 100644
--- a/rust/kernel/chrdev.rs
+++ b/rust/kernel/chrdev.rs
@@ -38,7 +38,7 @@ impl Builder {
         self
     }
 
-    pub fn build(self) -> KernelResult<Registration> {
+    pub fn build(self, this_module: &'static crate::ThisModule) -> KernelResult<Registration> {
         let mut dev: bindings::dev_t = 0;
         let res = unsafe {
             bindings::alloc_chrdev_region(
@@ -58,8 +58,7 @@ impl Builder {
         for (i, file_op) in self.file_ops.iter().enumerate() {
             unsafe {
                 bindings::cdev_init(&mut cdevs[i], *file_op);
-                // TODO: proper `THIS_MODULE` handling
-                cdevs[i].owner = core::ptr::null_mut();
+                cdevs[i].owner = this_module.0;
                 let rc = bindings::cdev_add(&mut cdevs[i], dev + i as bindings::dev_t, 1);
                 if rc != 0 {
                     // Clean up the ones that were allocated.
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 93bdaceaab7809..bee01ba2564680 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -44,6 +44,18 @@ pub trait KernelModule: Sized + Sync {
     fn init() -> KernelResult<Self>;
 }
 
+/// An instance equivalent to `THIS_MODULE` in C code.
+pub struct ThisModule(*mut bindings::module);
+
+// Safe because THIS_MODULE may be used from all threads within a module.
+unsafe impl Sync for ThisModule {}
+
+impl ThisModule {
+    pub const unsafe fn from_ptr(ptr: *mut bindings::module) -> ThisModule {
+        ThisModule(ptr)
+    }
+}
+
 extern "C" {
     fn rust_helper_BUG() -> !;
 }
diff --git a/rust/module.rs b/rust/module.rs
index 051a58bac8da65..1622f1c0a94608 100644
--- a/rust/module.rs
+++ b/rust/module.rs
@@ -309,6 +309,11 @@ pub fn module(ts: TokenStream) -> TokenStream {
         "
             static mut __MOD: Option<{type_}> = None;
 
+            #[cfg(MODULE)]
+            static THIS_MODULE: kernel::ThisModule = unsafe { kernel::ThisModule::from_ptr(&mut kernel::bindings::__this_module) };
+            #[cfg(not(MODULE))]
+            static THIS_MODULE: kernel::ThisModule = unsafe { kernel::ThisModule::from_ptr(core::ptr::null_mut()) };
+
             // Loadable modules need to export the `{{init,cleanup}}_module` identifiers
             #[cfg(MODULE)]
             #[no_mangle]