[PATCH 1/1] rust: pwm: track the owner module when registering `Chip`

[Priya Bala Govindasamy]

`UnregisteredChip::register` calls the C function `__pwmchip_add`
to register a PWM chip with the PWM core.
`__pwmchip_add` documents that its `owner` argument is supposed to
be `THIS_MODULE`, but `UnregisteredChip::register` passes
`core::ptr::null_mut()`. This is acceptable for built-in drivers.

However, for loadable drivers, registering the chip with a null owner
prevents the PWM core from taking a reference to the owner module when
a PWM is requested or its character device is opened.
This means that the owner module can be unloaded while references to the
chip remain. This can lead to use-after-free and
kernel panic.

Fix this by marking PwmOps as a `#[vtable]` trait so each implementation
exposes its actual OwnerModule, then pass that to `__pwmchip_add()`.

Fixes: d8046cd50879 ("rust: pwm: Add complete abstraction layer")
Reported-by: Dylan Zueck <dzueck at uci.edu>
Assisted-by: LLM
Signed-off-by: Priya Bala Govindasamy <pgovind2 at uci.edu>
---
 drivers/pwm/pwm_th1520.rs |  1 +
 rust/kernel/pwm.rs        | 11 +++++++++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/drivers/pwm/pwm_th1520.rs b/drivers/pwm/pwm_th1520.rs
index 728993998446..6a1f4219d6fa 100644
--- a/drivers/pwm/pwm_th1520.rs
+++ b/drivers/pwm/pwm_th1520.rs
@@ -94,6 +94,7 @@ struct Th1520PwmDriverData {
     clk: Clk,
 }
 
+#[vtable]
 impl pwm::PwmOps for Th1520PwmDriverData {
     type WfHw = Th1520WfHw;
 
diff --git a/rust/kernel/pwm.rs b/rust/kernel/pwm.rs
index 8aa47304bec3..ae7b449d67e4 100644
--- a/rust/kernel/pwm.rs
+++ b/rust/kernel/pwm.rs
@@ -12,6 +12,7 @@
     device::{self, Bound},
     devres,
     error::{self, to_result},
+    module::this_module,
     prelude::*,
     sync::aref::{ARef, AlwaysRefCounted},
     types::Opaque, //
@@ -175,6 +176,7 @@ pub struct RoundedWaveform<WfHw> {
 }
 
 /// Trait defining the operations for a PWM driver.
+#[vtable]
 pub trait PwmOps: 'static + Send + Sync + Sized {
     /// The driver-specific hardware representation of a waveform.
     ///
@@ -673,9 +675,14 @@ impl<T: PwmOps> UnregisteredChip<'_, T> {
     pub fn register(self) -> Result<ARef<Chip<T>>> {
         let c_chip_ptr = self.chip.as_raw();
 
-        // SAFETY: `c_chip_ptr` points to a valid chip with its ops initialized.
+        let owner = this_module::<T::OwnerModule>().as_ptr();
+
+        // SAFETY:
+        // - `c_chip_ptr` points to a valid chip with its ops initialized.
+        // - `owner` is the module implementing `PwmOps`. It is null for a built-in
+        //   implementation or a valid module pointer for a loadable implementation.
         // `__pwmchip_add` is the C function to register the chip with the PWM core.
-        to_result(unsafe { bindings::__pwmchip_add(c_chip_ptr, core::ptr::null_mut()) })?;
+        to_result(unsafe { bindings::__pwmchip_add(c_chip_ptr, owner) })?;
 
         let registration = Registration {
             chip: ARef::clone(&self.chip),
-- 
2.34.1