From ea08e7f864a616488627490796b8a6e3cc724cde Mon Sep 17 00:00:00 2001 From: Guillaume Gomez Date: Tue, 21 Nov 2023 23:00:32 +0100 Subject: [PATCH] Correctly handle `Vec::reserve` --- src/unix/users.rs | 3 + src/windows/cpu.rs | 23 +++++- src/windows/system.rs | 182 +++++++++++++++++++++--------------------- src/windows/utils.rs | 2 + 4 files changed, 114 insertions(+), 96 deletions(-) diff --git a/src/unix/users.rs b/src/unix/users.rs index db5f6bbe4..41f83c1c8 100644 --- a/src/unix/users.rs +++ b/src/unix/users.rs @@ -64,6 +64,9 @@ pub(crate) unsafe fn get_group_name( { // If there was not enough memory, we give it more. if last_errno == libc::ERANGE as _ { + // Needs to be updated for `Vec::reserve` to actually add additional capacity. + // In here it's "fine" since we never read from `buffer`. + buffer.set_len(buffer.capacity()); buffer.reserve(2048); continue; } diff --git a/src/windows/cpu.rs b/src/windows/cpu.rs index 642483988..ce4f6145b 100644 --- a/src/windows/cpu.rs +++ b/src/windows/cpu.rs @@ -501,19 +501,25 @@ pub(crate) fn get_frequencies(nb_cpus: usize) -> Vec { } pub(crate) fn get_physical_core_count() -> Option { - // we cannot use the number of cpus here to pre calculate the buf size - // GetLogicalCpuInformationEx with RelationProcessorCore passed to it not only returns - // the logical cores but also numa nodes + // We cannot use the number of cpus here to pre calculate the buf size. + // `GetLogicalCpuInformationEx` with `RelationProcessorCore` passed to it not only returns + // the logical cores but also numa nodes. // // GetLogicalCpuInformationEx: https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getlogicalprocessorinformationex let mut needed_size = 0; unsafe { + // This function call will always return an error as it only returns "success" when it + // has written at least one item in the buffer (which it cannot do here). let _err = GetLogicalProcessorInformationEx(RelationAll, None, &mut needed_size); let mut buf: Vec = Vec::with_capacity(needed_size as _); loop { + // Needs to be updated for `Vec::reserve` to actually add additional capacity if + // `GetLogicalProcessorInformationEx` fails because the buffer isn't big enough. + buf.set_len(needed_size as _); + if GetLogicalProcessorInformationEx( RelationAll, Some(buf.as_mut_ptr().cast()), @@ -540,6 +546,17 @@ pub(crate) fn get_physical_core_count() -> Option { } else { 1 }; + needed_size = match needed_size.checked_add(reserve as _) { + Some(new_size) => new_size, + None => { + sysinfo_debug!( + "get_physical_core_count: buffer size is too big ({} + {})", + needed_size, + reserve, + ); + return None; + } + }; buf.reserve(reserve); } diff --git a/src/windows/system.rs b/src/windows/system.rs index 685e75272..79c58ac80 100644 --- a/src/windows/system.rs +++ b/src/windows/system.rs @@ -201,14 +201,14 @@ impl SystemInner { let mut buffer_size = 512 * 1024; let mut process_information: Vec = Vec::with_capacity(buffer_size); - loop { - let mut cb_needed = 0; - // reserve(n) ensures the Vec has capacity for n elements on top of len - // so we should reserve buffer_size - len. len will always be zero at this point - // this is a no-op on the first call as buffer_size == capacity - process_information.reserve(buffer_size); - - unsafe { + unsafe { + loop { + let mut cb_needed = 0; + // reserve(n) ensures the Vec has capacity for n elements on top of len + // so we should reserve buffer_size - len. len will always be zero at this point + // this is a no-op on the first call as buffer_size == capacity + process_information.reserve(buffer_size); + let ntstatus = NtQuerySystemInformation( SystemProcessInformation, process_information.as_mut_ptr() as *mut _, @@ -238,105 +238,101 @@ impl SystemInner { } } } - } - // If we reach this point NtQuerySystemInformation succeeded - // and the buffer contents are initialized - unsafe { + // If we reach this point NtQuerySystemInformation succeeded + // and the buffer contents are initialized process_information.set_len(buffer_size); - } - // Parse the data block to get process information - let mut process_ids = Vec::with_capacity(500); - let mut process_information_offset = 0; - loop { - let p = unsafe { - process_information + // Parse the data block to get process information + let mut process_ids = Vec::with_capacity(500); + let mut process_information_offset = 0; + loop { + let p = process_information .as_ptr() .offset(process_information_offset) - as *const SYSTEM_PROCESS_INFORMATION - }; - - process_ids.push(Wrap(p)); + as *const SYSTEM_PROCESS_INFORMATION; - // read_unaligned is necessary to avoid - // misaligned pointer dereference: address must be a multiple of 0x8 but is 0x... - // under x86_64 wine (and possibly other systems) - let pi = unsafe { ptr::read_unaligned(p) }; + process_ids.push(Wrap(p)); - if pi.NextEntryOffset == 0 { - break; - } + // read_unaligned is necessary to avoid + // misaligned pointer dereference: address must be a multiple of 0x8 but is 0x... + // under x86_64 wine (and possibly other systems) + let pi = ptr::read_unaligned(p); - process_information_offset += pi.NextEntryOffset as isize; - } - let process_list = Wrap(UnsafeCell::new(&mut self.process_list)); - let nb_cpus = if refresh_kind.cpu() { - self.cpus.len() as u64 - } else { - 0 - }; + if pi.NextEntryOffset == 0 { + break; + } - let now = get_now(); + process_information_offset += pi.NextEntryOffset as isize; + } + let process_list = Wrap(UnsafeCell::new(&mut self.process_list)); + let nb_cpus = if refresh_kind.cpu() { + self.cpus.len() as u64 + } else { + 0 + }; - #[cfg(feature = "multithread")] - use rayon::iter::ParallelIterator; - - // TODO: instead of using parallel iterator only here, would be better to be - // able to run it over `process_information` directly! - let processes = into_iter(process_ids) - .filter_map(|pi| { - // as above, read_unaligned is necessary - let pi = unsafe { ptr::read_unaligned(pi.0) }; - let pid = Pid(pi.UniqueProcessId as _); - if let Some(proc_) = unsafe { (*process_list.0.get()).get_mut(&pid) } { - let proc_ = &mut proc_.inner; - if proc_ - .get_start_time() - .map(|start| start == proc_.start_time()) - .unwrap_or(true) - { - if refresh_kind.memory() { - proc_.memory = pi.WorkingSetSize as _; - proc_.virtual_memory = pi.VirtualSize as _; + let now = get_now(); + + #[cfg(feature = "multithread")] + use rayon::iter::ParallelIterator; + + // TODO: instead of using parallel iterator only here, would be better to be + // able to run it over `process_information` directly! + let processes = into_iter(process_ids) + .filter_map(|pi| { + // as above, read_unaligned is necessary + let pi = ptr::read_unaligned(pi.0); + let pid = Pid(pi.UniqueProcessId as _); + if let Some(proc_) = (*process_list.0.get()).get_mut(&pid) { + let proc_ = &mut proc_.inner; + if proc_ + .get_start_time() + .map(|start| start == proc_.start_time()) + .unwrap_or(true) + { + if refresh_kind.memory() { + proc_.memory = pi.WorkingSetSize as _; + proc_.virtual_memory = pi.VirtualSize as _; + } + proc_.update(refresh_kind, nb_cpus, now); + return None; } - proc_.update(refresh_kind, nb_cpus, now); - return None; + // If the PID owner changed, we need to recompute the whole process. + sysinfo_debug!("owner changed for PID {}", proc_.pid()); } - // If the PID owner changed, we need to recompute the whole process. - sysinfo_debug!("owner changed for PID {}", proc_.pid()); - } - let name = get_process_name(&pi, pid); - let (memory, virtual_memory) = if refresh_kind.memory() { - (pi.WorkingSetSize as _, pi.VirtualSize as _) - } else { - (0, 0) - }; - let mut p = ProcessInner::new_full( - pid, - if pi.InheritedFromUniqueProcessId as usize != 0 { - Some(Pid(pi.InheritedFromUniqueProcessId as _)) + let name = get_process_name(&pi, pid); + let (memory, virtual_memory) = if refresh_kind.memory() { + (pi.WorkingSetSize as _, pi.VirtualSize as _) } else { - None - }, - memory, - virtual_memory, - name, - now, - refresh_kind, - ); - p.update(refresh_kind.without_memory(), nb_cpus, now); - Some(Process { inner: p }) - }) - .collect::>(); - for p in processes.into_iter() { - self.process_list.insert(p.pid(), p); + (0, 0) + }; + let mut p = ProcessInner::new_full( + pid, + if pi.InheritedFromUniqueProcessId as usize != 0 { + Some(Pid(pi.InheritedFromUniqueProcessId as _)) + } else { + None + }, + memory, + virtual_memory, + name, + now, + refresh_kind, + ); + p.update(refresh_kind.without_memory(), nb_cpus, now); + Some(Process { inner: p }) + }) + .collect::>(); + for p in processes.into_iter() { + self.process_list.insert(p.pid(), p); + } + self.process_list.retain(|_, v| { + let x = v.inner.updated; + v.inner.updated = false; + x + }); } - self.process_list.retain(|_, v| { - let x = v.inner.updated; - v.inner.updated = false; - x - }); } pub(crate) fn processes(&self) -> &HashMap { diff --git a/src/windows/utils.rs b/src/windows/utils.rs index 25545b9aa..56deb796d 100644 --- a/src/windows/utils.rs +++ b/src/windows/utils.rs @@ -98,6 +98,8 @@ pub(crate) fn get_reg_string_value(hkey: HKEY, path: &str, field_name: &str) -> match new_key.get_value(&c_field_name, &mut buf, &mut buf_len) { Ok(()) => break, Err(err) if err.code() == Foundation::ERROR_MORE_DATA.to_hresult() => { + // Needs to be updated for `Vec::reserve` to actually add additional capacity. + buf.set_len(buf.capacity()); buf.reserve(buf_len as _); } _ => return None,