Re: [PATCH] rust: cpumask: Validate CPU number in set() and clear()

[Boqun Feng <boqun.feng@gmail.com>]

On Fri, Jun 06, 2025 at 09:47:28AM +0530, Viresh Kumar wrote:
> The C `cpumask_{set|clear}_cpu()` APIs emit a warning when given an
> invalid CPU number - but only if `CONFIG_DEBUG_PER_CPU_MAPS=y` is set.
> 
> Meanwhile, `cpumask_weight()` only considers CPUs up to `nr_cpu_ids`,
> which can cause inconsistencies: a CPU number greater than `nr_cpu_ids`
> may be set in the mask, yet the weight calculation won't reflect it.
> 
> This leads to doctest failures when `nr_cpu_ids < 4`, as the test tries
> to set CPUs 2 and 3:
> 
>   rust_doctest_kernel_cpumask_rs_0.location: rust/kernel/cpumask.rs:180
>   rust_doctest_kernel_cpumask_rs_0: ASSERTION FAILED at rust/kernel/cpumask.rs:190
> 
> Fix this by validating the CPU number in the Rust `set()` and `clear()`
> methods to prevent out-of-bounds modifications.
> 

Thanks for the quick fix!

While this can fix the current problem, but it's not a good solution for
the long run. Because outside a test, we should never use an arbitrary
i32 as a cpu number (we usually get it from smp_processor_id(), or
something else). So the `< nr_cpu_ids` testing is not necessary in
normal use cases.

We should instead provide a wrapper for cpu id:

    /// # Invariants
    ///
    /// The number is always in [0..nr_cpu_ids) range.
    pub struct CpuId(i32);

and

    impl CpuId {
        /// # Safety
	/// Callers must ensure `i` is a valid cpu id (i.e. 0 <= i <
	/// nr_cpu_ids).
        pub unsafe fn from_i32_unchecked(i: i32) -> Self {
	    // INVARIANT: The function safety guarantees `i` is a valid
	    // cpu id.
	    CpuId(id);
	}

	pub fn from_i32(i: i32) -> Option<Self> {
	    if i < 0 || i >= nr_cpu_ids {
	        None
	    } else {
	        // SAFETY: `i` has just been checked as a valid cpu id.
	        Some(unsafe { Self::from_i32_unchecked(i) })
	    }
	}

	pub fn current() -> Self {
	    // SAFETY: smp_processor_id() always return valid cpu id.
	    unsafe { Self::from_i32_unchecked(smp_processor_id()) }
	}
    }

All `Cpumask` functions then take `CpuId` instead of `i32` as the
parameter. Needless to say if we were to have a cpumask_next() wrapper,
the return value will be `CpuId` (or `Option<CpuId>`), i.e. if a bit was
set in a cpumask, then it must represent a correct cpu id.

Make sense?

> Fixes: 8961b8cb3099 ("rust: cpumask: Add initial abstractions")
> Reported-by: Miguel Ojeda <ojeda@kernel.org>
> Closes: https://lore.kernel.org/all/87qzzy3ric.fsf@kernel.org/
> Reported-by: Andreas Hindborg <a.hindborg@kernel.org>
> Closes: https://lore.kernel.org/all/87qzzy3ric.fsf@kernel.org/
> Signed-off-by: Viresh Kumar <viresh.kumar@linaro.org>
> ---
>  drivers/cpufreq/rcpufreq_dt.rs |  2 +-
>  rust/kernel/cpumask.rs         | 49 +++++++++++++++++++++++-----------
>  2 files changed, 34 insertions(+), 17 deletions(-)
> 
> diff --git a/drivers/cpufreq/rcpufreq_dt.rs b/drivers/cpufreq/rcpufreq_dt.rs
> index 94ed81644fe1..f396c8f35069 100644
> --- a/drivers/cpufreq/rcpufreq_dt.rs
> +++ b/drivers/cpufreq/rcpufreq_dt.rs
> @@ -70,7 +70,7 @@ fn init(policy: &mut cpufreq::Policy) -> Result<Self::PData> {
>          let dev = unsafe { cpu::from_cpu(cpu)? };
>          let mut mask = CpumaskVar::new_zero(GFP_KERNEL)?;
>  
> -        mask.set(cpu);
> +        mask.set(cpu)?;
>  
>          let token = find_supply_names(dev, cpu)
>              .map(|names| {
> diff --git a/rust/kernel/cpumask.rs b/rust/kernel/cpumask.rs
> index c90bfac9346a..75d4ce916b4f 100644
> --- a/rust/kernel/cpumask.rs
> +++ b/rust/kernel/cpumask.rs
> @@ -37,13 +37,14 @@
>  /// use kernel::bindings;
>  /// use kernel::cpumask::Cpumask;
>  ///
> -/// fn set_clear_cpu(ptr: *mut bindings::cpumask, set_cpu: u32, clear_cpu: i32) {
> +/// fn set_clear_cpu(ptr: *mut bindings::cpumask, set_cpu: u32, clear_cpu: i32) -> Result {
>  ///     // SAFETY: The `ptr` is valid for writing and remains valid for the lifetime of the
>  ///     // returned reference.
>  ///     let mask = unsafe { Cpumask::as_mut_ref(ptr) };
>  ///
> -///     mask.set(set_cpu);
> -///     mask.clear(clear_cpu);
> +///     mask.set(set_cpu)?;
> +///     mask.clear(clear_cpu)?;
> +///     Ok(())
>  /// }
>  /// ```
>  #[repr(transparent)]
> @@ -90,9 +91,15 @@ pub fn as_raw(&self) -> *mut bindings::cpumask {
>      /// This mismatches kernel naming convention and corresponds to the C
>      /// function `__cpumask_set_cpu()`.
>      #[inline]
> -    pub fn set(&mut self, cpu: u32) {
> +    pub fn set(&mut self, cpu: u32) -> Result {
> +        // SAFETY: It is safe to read `nr_cpu_ids`.
> +        if unsafe { cpu >= bindings::nr_cpu_ids } {
> +            return Err(EINVAL);
> +        }
> +
>          // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `__cpumask_set_cpu`.
>          unsafe { bindings::__cpumask_set_cpu(cpu, self.as_raw()) };
> +        Ok(())
>      }
>  
>      /// Clear `cpu` in the cpumask.
> @@ -101,10 +108,16 @@ pub fn set(&mut self, cpu: u32) {
>      /// This mismatches kernel naming convention and corresponds to the C
>      /// function `__cpumask_clear_cpu()`.
>      #[inline]
> -    pub fn clear(&mut self, cpu: i32) {
> +    pub fn clear(&mut self, cpu: i32) -> Result {
> +        // SAFETY: It is safe to read `nr_cpu_ids`.
> +        if unsafe { cpu as u32 >= bindings::nr_cpu_ids } {

You probably want to check whether `bindings::nr_cpu_ids` can be
accessible if NR_CPUS == 1 or CONFIG_FORCE_NR_CPUS=y, because then
nr_cpu_ids is a macro definition.

Regards,
Boqun

> +            return Err(EINVAL);
> +        }
> +
>          // SAFETY: By the type invariant, `self.as_raw` is a valid argument to
>          // `__cpumask_clear_cpu`.
>          unsafe { bindings::__cpumask_clear_cpu(cpu, self.as_raw()) };
> +        Ok(())
>      }
>  
>      /// Test `cpu` in the cpumask.
> @@ -180,19 +193,23 @@ pub fn copy(&self, dstp: &mut Self) {
>  /// ```
>  /// use kernel::cpumask::CpumaskVar;
>  ///
> -/// let mut mask = CpumaskVar::new_zero(GFP_KERNEL).unwrap();
> +/// fn cpumask_test() -> Result {
> +///     let mut mask = CpumaskVar::new_zero(GFP_KERNEL).unwrap();
>  ///
> -/// assert!(mask.empty());
> -/// mask.set(2);
> -/// assert!(mask.test(2));
> -/// mask.set(3);
> -/// assert!(mask.test(3));
> -/// assert_eq!(mask.weight(), 2);
> +///     assert!(mask.empty());
> +///     mask.set(2)?;
> +///     assert!(mask.test(2));
> +///     mask.set(3)?;
> +///     assert!(mask.test(3));
> +///     assert_eq!(mask.weight(), 2);
>  ///
> -/// let mask2 = CpumaskVar::try_clone(&mask).unwrap();
> -/// assert!(mask2.test(2));
> -/// assert!(mask2.test(3));
> -/// assert_eq!(mask2.weight(), 2);
> +///     let mask2 = CpumaskVar::try_clone(&mask).unwrap();
> +///     assert!(mask2.test(2));
> +///     assert!(mask2.test(3));
> +///     assert_eq!(mask2.weight(), 2);
> +///
> +///     Ok(())
> +/// }
>  /// ```
>  pub struct CpumaskVar {
>      #[cfg(CONFIG_CPUMASK_OFFSTACK)]
> -- 
> 2.31.1.272.g89b43f80a514
>