From 889dd0f0e7723a6263f01dbcbf3d882037f0b724 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 3 Nov 2026 07:12:42 +0000 Subject: [PATCH] Reject an empty processor performance query A successful NtQuerySystemInformation result with a zero returned length is not a sample. Return an error and drop the default-zero counters instead of treating them as idle time. Co-authored-by: jameslincn --- src/system/cpu_sampler.rs | 173 ++++++++++++++++++++++++++++++++++---- 1 file changed, 158 insertions(+), 15 deletions(-) diff --git a/src/system/cpu_sampler.rs b/src/system/cpu_sampler.rs index 0e6e724..40def28 100644 --- a/src/system/cpu_sampler.rs +++ b/src/system/cpu_sampler.rs @@ -61,23 +61,69 @@ pub(crate) fn query_processor_performance( ) }; - if status >= 0 { - if returned != 0 { - let returned = returned as usize; - if !returned.is_multiple_of(item_size) || returned > byte_len as usize { - return Err(STATUS_INVALID_BUFFER_SIZE); - } - if returned / item_size != count { - return Err(STATUS_INFO_LENGTH_MISMATCH); - } - } - return Ok(()); + match take_processor_performance_query(output, status, returned, byte_len, count) { + Ok(ProcessorPerformanceQuery::Complete) => return Ok(()), + Ok(ProcessorPerformanceQuery::GrowTo(next_count)) => count = next_count, + Err(error) => return Err(error), } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProcessorPerformanceQuery { + Complete, + GrowTo(usize), +} - if status != STATUS_INFO_LENGTH_MISMATCH || returned as usize <= byte_len as usize { - return Err(status); +/// Classifies one processor-performance query. +/// +/// A success status with a zero returned length is an error. A nonzero length is +/// accepted only when it is an exact multiple of the struct and matches `count`. +fn decide_processor_performance_query( + status: i32, + returned: u32, + buffer_bytes: u32, + count: usize, +) -> Result { + let item_size = size_of::(); + if status >= 0 { + if returned == 0 { + return Err(STATUS_INVALID_BUFFER_SIZE); + } + let returned_len = returned as usize; + let buffer_len = buffer_bytes as usize; + if !returned_len.is_multiple_of(item_size) || returned_len > buffer_len { + return Err(STATUS_INVALID_BUFFER_SIZE); } - count = (returned as usize).div_ceil(item_size); + if returned_len / item_size != count { + return Err(STATUS_INFO_LENGTH_MISMATCH); + } + return Ok(ProcessorPerformanceQuery::Complete); + } + + if status != STATUS_INFO_LENGTH_MISMATCH || returned as usize <= buffer_bytes as usize { + return Err(status); + } + Ok(ProcessorPerformanceQuery::GrowTo( + (returned as usize).div_ceil(item_size), + )) +} + +fn take_processor_performance_query( + output: &mut Vec, + status: i32, + returned: u32, + buffer_bytes: u32, + count: usize, +) -> Result { + match decide_processor_performance_query(status, returned, buffer_bytes, count) { + // A successful query that wrote nothing leaves the caller's default-zero + // rows in place. Drop them so they cannot be read as idle counters. + Err(error) if status >= 0 && returned == 0 => { + output.clear(); + Err(error) + } + decision => decision, } } @@ -100,8 +146,13 @@ pub(crate) fn checked_summed_processor_times( #[cfg(test)] mod tests { + use std::mem::size_of; + use super::{ - ProcessorPerformance, checked_summed_processor_times, query_processor_performance, + ProcessorPerformance, ProcessorPerformanceQuery, STATUS_INFO_LENGTH_MISMATCH, + STATUS_INVALID_BUFFER_SIZE, checked_summed_processor_times, + decide_processor_performance_query, query_processor_performance, + take_processor_performance_query, }; use windows_sys::Win32::System::Threading::{ALL_PROCESSOR_GROUPS, GetActiveProcessorCount}; @@ -155,4 +206,96 @@ mod tests { Some((15, 70, 27)) ); } + + #[test] + fn success_with_zero_returned_length_is_an_error() { + let bytes = size_of::() as u32; + assert!(decide_processor_performance_query(0, 0, bytes, 1).is_err()); + assert!(decide_processor_performance_query(0, 0, bytes, 0).is_err()); + } + + #[test] + fn success_with_exact_multiple_is_ok() { + let item_size = size_of::(); + let count = 3; + let bytes = u32::try_from(item_size * count).unwrap(); + assert_eq!( + decide_processor_performance_query(0, bytes, bytes, count), + Ok(ProcessorPerformanceQuery::Complete) + ); + } + + #[test] + fn non_multiple_returned_length_is_an_error() { + let bytes = size_of::() as u32; + assert_eq!( + decide_processor_performance_query(0, bytes - 1, bytes, 1), + Err(STATUS_INVALID_BUFFER_SIZE) + ); + } + + #[test] + fn short_buffer_keeps_invalid_buffer_size() { + let one = size_of::() as u32; + assert_eq!( + decide_processor_performance_query(0, one * 2, one, 2), + Err(STATUS_INVALID_BUFFER_SIZE) + ); + } + + #[test] + fn count_mismatch_keeps_length_mismatch() { + let one = size_of::() as u32; + let bytes = one * 2; + assert_eq!( + decide_processor_performance_query(0, bytes, bytes, 1), + Err(STATUS_INFO_LENGTH_MISMATCH) + ); + } + + #[test] + fn length_mismatch_grows_when_the_kernel_reports_a_larger_buffer() { + let one = size_of::() as u32; + assert_eq!( + decide_processor_performance_query(STATUS_INFO_LENGTH_MISMATCH, one * 3, one, 1), + Ok(ProcessorPerformanceQuery::GrowTo(3)) + ); + } + + #[test] + fn length_mismatch_without_a_larger_return_keeps_the_status() { + let bytes = size_of::() as u32; + assert_eq!( + decide_processor_performance_query(STATUS_INFO_LENGTH_MISMATCH, bytes, bytes, 1), + Err(STATUS_INFO_LENGTH_MISMATCH) + ); + } + + #[test] + fn zero_returned_length_drops_default_processor_times() { + let mut processors = vec![ProcessorPerformance::default(); 2]; + let bytes = size_of::() as u32; + assert!(take_processor_performance_query(&mut processors, 0, 0, bytes, 2).is_err()); + assert!(processors.is_empty()); + } + + #[test] + fn failed_zero_return_keeps_existing_rows() { + let mut processors = vec![ProcessorPerformance { + idle_time: 4, + ..ProcessorPerformance::default() + }]; + let bytes = size_of::() as u32; + assert_eq!( + take_processor_performance_query( + &mut processors, + STATUS_INFO_LENGTH_MISMATCH, + 0, + bytes, + 1 + ), + Err(STATUS_INFO_LENGTH_MISMATCH) + ); + assert_eq!(processors[0].idle_time, 4); + } }