[PATCH v2] rust: block: set GenDisk block_device_operations.owner to THIS_MODULE

GenDiskBuilder left block_device_operations.owner NULL. Pass the driver's
ThisModule into GenDiskBuilder::build(), heap-allocate the operations
table, and keep it alive until the gendisk is released via free_disk.
Update rnull as the in-tree caller.

v2:
- Free fops in free_disk instead of GenDisk::drop to fix use-after-free
  when the device stays open after removal. (Sashiko)
- Install the cleanup guard before fops allocation to avoid leaking gendisk
  on -ENOMEM. (Sashiko)
- Link to v1: https://lore.kernel.org/all/[email protected]/

Signed-off-by: Adarsh Das <[email protected]>
---
 drivers/block/rnull/configfs.rs  |  1 +
 drivers/block/rnull/rnull.rs     |  3 +-
 rust/kernel/block/mq.rs          |  9 ++--
 rust/kernel/block/mq/gen_disk.rs | 82 ++++++++++++++++++++++----------
 4 files changed, 66 insertions(+), 29 deletions(-)

diff --git a/drivers/block/rnull/configfs.rs b/drivers/block/rnull/configfs.rs
index 7c2eb5c0b722..bba30d590f68 100644
--- a/drivers/block/rnull/configfs.rs
+++ b/drivers/block/rnull/configfs.rs
@@ -147,6 +147,7 @@ fn store(this: &DeviceConfig, page: &[u8]) -> Result {
 
         if !guard.powered && power_op {
             guard.disk = Some(NullBlkDevice::new(
+                &THIS_MODULE,
                 &guard.name,
                 guard.block_size,
                 guard.rotational,
diff --git a/drivers/block/rnull/rnull.rs b/drivers/block/rnull/rnull.rs
index 0ca8715febe8..4265a133cbf0 100644
--- a/drivers/block/rnull/rnull.rs
+++ b/drivers/block/rnull/rnull.rs
@@ -46,6 +46,7 @@ fn init(_module: &'static ThisModule) -> impl PinInit<Self, Error> {
 
 impl NullBlkDevice {
     fn new(
+        this_module: &'static ThisModule,
         name: &CStr,
         block_size: u32,
         rotational: bool,
@@ -61,7 +62,7 @@ fn new(
             .logical_block_size(block_size)?
             .physical_block_size(block_size)?
             .rotational(rotational)
-            .build(fmt!("{}", name.to_str()?), tagset, queue_data)
+            .build(this_module, fmt!("{}", name.to_str()?), tagset, queue_data)
     }
 }
 
diff --git a/rust/kernel/block/mq.rs b/rust/kernel/block/mq.rs
index 1fd0d54dd549..33561e0f67af 100644
--- a/rust/kernel/block/mq.rs
+++ b/rust/kernel/block/mq.rs
@@ -8,8 +8,8 @@
 //! - Implement [`Operations`] for a type `T`.
 //! - Create a [`TagSet<T>`].
 //! - Create a [`GenDisk<T>`], via the [`GenDiskBuilder`].
-//! - Add the disk to the system by calling [`GenDiskBuilder::build`] passing in
-//!   the `TagSet` reference.
+//! - Add the disk to the system by calling [`GenDiskBuilder::build`], passing in
+//!   the driver's [`ThisModule`], the disk name, the `TagSet`, and queue data.
 //!
 //! The types available in this module that have direct C counterparts are:
 //!
@@ -86,9 +86,12 @@
 //!
 //! let tagset: Arc<TagSet<MyBlkDevice>> =
 //!     Arc::pin_init(TagSet::new(1, 256, 1), flags::GFP_KERNEL)?;
+//! # // SAFETY: Dummy `ThisModule` for doctest compilation only.
+//! # static THIS_MODULE: ThisModule =
+//! #     unsafe { ThisModule::from_ptr(core::ptr::null_mut()) };
 //! let mut disk = gen_disk::GenDiskBuilder::new()
 //!     .capacity_sectors(4096)
-//!     .build(fmt!("myblk"), tagset, ())?;
+//!     .build(&THIS_MODULE, fmt!("myblk"), tagset, ())?;
 //!
 //! # Ok::<(), kernel::error::Error>(())
 //! ```
diff --git a/rust/kernel/block/mq/gen_disk.rs b/rust/kernel/block/mq/gen_disk.rs
index fc97dd873974..a51027e8c1c1 100644
--- a/rust/kernel/block/mq/gen_disk.rs
+++ b/rust/kernel/block/mq/gen_disk.rs
@@ -17,6 +17,23 @@
     types::{ForeignOwnable, ScopeGuard},
 };
 
+/// # Safety
+///
+/// `disk` must be valid.
+unsafe extern "C" fn free_fops(disk: *mut bindings::gendisk) {
+    // SAFETY: `disk` is valid.
+    let fops = unsafe { (*disk).fops };
+    if fops.is_null() {
+        return;
+    }
+
+    // SAFETY: `disk` is valid; `fops` came from `KBox::into_raw` in `build`.
+    unsafe {
+        (*disk).fops = core::ptr::null_mut();
+        drop(KBox::from_raw(fops.cast_mut()));
+    }
+}
+
 /// A builder for [`GenDisk`].
 ///
 /// Use this struct to configure and add new [`GenDisk`] to the VFS.
@@ -95,8 +112,12 @@ pub fn capacity_sectors(mut self, capacity: u64) -> Self {
     }
 
     /// Build a new `GenDisk` and add it to the VFS.
+    ///
+    /// `this_module` must be the [`ThisModule`] for the kernel module registering
+    /// the disk.
     pub fn build<T: Operations>(
         self,
+        this_module: &'static ThisModule,
         name: fmt::Arguments<'_>,
         tagset: Arc<TagSet<T>>,
         queue_data: T::QueueData,
@@ -125,32 +146,16 @@ pub fn build<T: Operations>(
             )
         })?;
 
-        const TABLE: bindings::block_device_operations = bindings::block_device_operations {
-            submit_bio: None,
-            open: None,
-            release: None,
-            ioctl: None,
-            compat_ioctl: None,
-            check_events: None,
-            unlock_native_capacity: None,
-            getgeo: None,
-            set_read_only: None,
-            swap_slot_free_notify: None,
-            report_zones: None,
-            devnode: None,
-            alternative_gpt_sector: None,
-            get_unique_id: None,
-            // TODO: Set to `THIS_MODULE`.
-            owner: core::ptr::null_mut(),
-            pr_ops: core::ptr::null_mut(),
-            free_disk: None,
-            poll_bio: None,
-        };
-
-        // SAFETY: `gendisk` is a valid pointer as we initialized it above
-        unsafe { (*gendisk).fops = &TABLE };
-
         let cleanup_failure = ScopeGuard::new_with_data((gendisk, data), |(gendisk, data)| {
+            // SAFETY: `gendisk` came from `__blk_mq_alloc_disk()` above and
+            // has not been added to the VFS on this cleanup path.
+            let fops = unsafe { (*gendisk).fops };
+            if !fops.is_null() {
+                // SAFETY: `gendisk` came from `__blk_mq_alloc_disk()` above.
+                unsafe { (*gendisk).fops = core::ptr::null_mut() };
+                // SAFETY: `fops` came from `KBox::into_raw` below on this path.
+                drop(unsafe { KBox::from_raw(fops.cast_mut()) });
+            }
             // SAFETY: `gendisk` came from `__blk_mq_alloc_disk()` above and
             // has not been added to the VFS on this cleanup path.
             unsafe { bindings::put_disk(gendisk) };
@@ -159,6 +164,33 @@ pub fn build<T: Operations>(
             drop(unsafe { T::QueueData::from_foreign(data) });
         });
 
+        let fops = KBox::new(
+            bindings::block_device_operations {
+                submit_bio: None,
+                open: None,
+                release: None,
+                ioctl: None,
+                compat_ioctl: None,
+                check_events: None,
+                unlock_native_capacity: None,
+                getgeo: None,
+                set_read_only: None,
+                swap_slot_free_notify: None,
+                report_zones: None,
+                devnode: None,
+                alternative_gpt_sector: None,
+                get_unique_id: None,
+                owner: this_module.as_ptr(),
+                pr_ops: core::ptr::null_mut(),
+                free_disk: Some(free_fops),
+                poll_bio: None,
+            },
+            GFP_KERNEL,
+        )?;
+
+        // SAFETY: `gendisk` is a valid pointer as we initialized it above.
+        unsafe { (*gendisk).fops = KBox::into_raw(fops).cast() };
+
         // The failure guard now owns both pieces of cleanup; the early guard
         // must not run on this path anymore.
         recover_data.dismiss();
-- 
2.55.0