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

["Benno Lossin" <lossin@kernel.org>]

On Fri Jun 6, 2025 at 6:37 AM CEST, Boqun Feng wrote:
> 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?

Just to make sure, the `nr_cpu_ids` stays constant, right?

>> @@ -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.

Just define a helper function?

---
Cheers,
Benno