From faccb0c5c673c3b0efd6419402546fe9adf02496 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/Makefile         |  2 +-
 rust/kernel/chrdev.rs |  5 ++---
 rust/kernel/lib.rs    | 12 ++++++++++++
 rust/module.rs        |  5 +++++
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/rust/Makefile b/rust/Makefile
index f5af10260ffb5e..5c4d60d4f80172 100644
--- a/rust/Makefile
+++ b/rust/Makefile
@@ -38,7 +38,7 @@ quiet_cmd_bindgen = BINDGEN $@
       cmd_bindgen = \
 	$(BINDGEN) $< $(addprefix --opaque-type , $(bindgen_opaque_types)) \
 		--use-core --with-derive-default --ctypes-prefix c_types \
-		--size_t-is-usize -o $@ -- $(bindgen_c_flags)
+		--size_t-is-usize -o $@ -- $(bindgen_c_flags) -DMODULE
 
 $(objtree)/rust/bindings_generated.rs: $(srctree)/rust/kernel/bindings_helper.h FORCE
 	$(call if_changed_dep,bindgen)
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..f1206fef190d8c 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);
+
+// SAFETY: `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..81813904b936cc 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(&kernel::bindings::__this_module as *const _ as *mut _) }};
+            #[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]