Skip to content

Commit

Permalink
Replace NtQuerySystemInformation with CreateToolhelp32Snapshot, `…
Browse files Browse the repository at this point in the history 
…Process32FirstW` and `Process32NextW` (#1451)
  • Loading branch information
@provrb
provrb authored Jan 17, 2025
1 parent 8a7dc24 commit 6749d8a
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 164 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ disk = [
]
system = [
"windows/Win32_Foundation",
"windows/Win32_System_Diagnostics_ToolHelp",
"windows/Wdk_System_SystemInformation",
"windows/Wdk_System_SystemServices",
"windows/Wdk_System_Threading",
Expand Down
50 changes: 40 additions & 10 deletions src/windows/process.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,11 @@ use std::os::windows::process::{CommandExt, ExitStatusExt};
use std::path::{Path, PathBuf};
use std::process::{self, ExitStatus};
use std::ptr::null_mut;
use std::str;
use std::str::{self, FromStr};
use std::sync::{Arc, OnceLock};
use std::time::Instant;

use libc::c_void;
use ntapi::ntexapi::SYSTEM_PROCESS_INFORMATION;
use ntapi::ntrtl::RTL_USER_PROCESS_PARAMETERS;
use ntapi::ntwow64::{PEB32, RTL_USER_PROCESS_PARAMETERS32};
use windows::core::PCWSTR;
Expand All @@ -35,10 +34,13 @@ use windows::Win32::Foundation::{
};
use windows::Win32::Security::{GetTokenInformation, TokenUser, TOKEN_QUERY, TOKEN_USER};
use windows::Win32::System::Diagnostics::Debug::ReadProcessMemory;
use windows::Win32::System::Diagnostics::ToolHelp::PROCESSENTRY32W;
use windows::Win32::System::Memory::{
GetProcessHeap, HeapAlloc, HeapFree, VirtualQueryEx, HEAP_ZERO_MEMORY, MEMORY_BASIC_INFORMATION,
};
use windows::Win32::System::ProcessStatus::GetModuleFileNameExW;
use windows::Win32::System::ProcessStatus::{
GetModuleFileNameExW, GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS_EX,
};
use windows::Win32::System::RemoteDesktop::ProcessIdToSessionId;
use windows::Win32::System::SystemInformation::OSVERSIONINFOEXW;
use windows::Win32::System::Threading::{
Expand Down Expand Up @@ -281,7 +283,6 @@ impl ProcessInner {
nb_cpus: u64,
now: u64,
refresh_parent: bool,
pi: &SYSTEM_PROCESS_INFORMATION,
) {
if refresh_kind.cpu() {
compute_cpu_usage(self, nb_cpus);
Expand All @@ -290,8 +291,21 @@ impl ProcessInner {
update_disk_usage(self);
}
if refresh_kind.memory() {
self.memory = pi.WorkingSetSize as _;
self.virtual_memory = pi.VirtualSize as _;
let mut mem_info = PROCESS_MEMORY_COUNTERS_EX::default();
if let Some(handle) = self.get_handle() {
if let Err(_error) = unsafe {
GetProcessMemoryInfo(
handle,
&mut mem_info as *mut _ as *mut _,
std::mem::size_of_val::<PROCESS_MEMORY_COUNTERS_EX>(&mem_info) as _,
)
} {
sysinfo_debug!("GetProcessMemoryInfo failed: {_error:?}");
} else {
self.memory = mem_info.WorkingSetSize as _;
self.virtual_memory = mem_info.PrivateUsage as _;
}
}
}
unsafe {
get_process_user_id(self, refresh_kind);
Expand All @@ -309,12 +323,28 @@ impl ProcessInner {
self.updated = true;
}

pub(crate) fn get_handle(&self) -> Option<HANDLE> {
self.handle.as_ref().map(|h| ***h)
pub(crate) fn from_process_entry(entry: &PROCESSENTRY32W, now: u64) -> Self {
let pid = Pid::from_u32(entry.th32ProcessID);
let name = match OsString::from_str(
String::from_utf16_lossy(&entry.szExeFile).trim_end_matches('\0'),
) {
Ok(name) => name,
Err(_) => format!("<no name> Process {pid}").into(),
};
let ppid = {
if entry.th32ParentProcessID == 0 {
// no parent pid
None
} else {
Some(Pid::from_u32(entry.th32ParentProcessID))
}
};

Self::new(pid, ppid, now, name)
}

pub(crate) fn get_start_time(&self) -> Option<u64> {
self.handle.as_ref().map(|handle| get_start_time(***handle))
pub(crate) fn get_handle(&self) -> Option<HANDLE> {
self.handle.as_ref().map(|h| ***h)
}

pub(crate) fn kill_with(&self, signal: Signal) -> Option<bool> {
Expand Down
224 changes: 70 additions & 154 deletions src/windows/system.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,17 @@ use crate::{
use crate::sys::cpu::*;
use crate::{Process, ProcessInner};

use crate::utils::into_iter;

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::ffi::{OsStr, OsString};
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::ptr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::os::windows::ffi::OsStrExt;
use std::time::{Duration, SystemTime};

use ntapi::ntexapi::SYSTEM_PROCESS_INFORMATION;
use windows::core::{PCWSTR, PWSTR};
use windows::Wdk::System::SystemInformation::{NtQuerySystemInformation, SystemProcessInformation};
use windows::Win32::Foundation::{self, HANDLE, STATUS_INFO_LENGTH_MISMATCH, STILL_ACTIVE};
use windows::Win32::Foundation::{self, HANDLE, STILL_ACTIVE};
use windows::Win32::System::Diagnostics::ToolHelp::{
CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
};
use windows::Win32::System::ProcessStatus::{K32GetPerformanceInfo, PERFORMANCE_INFORMATION};
use windows::Win32::System::Registry::{
RegCloseKey, RegOpenKeyExW, RegQueryValueExW, HKEY, HKEY_LOCAL_MACHINE, KEY_READ, REG_NONE,
Expand Down Expand Up @@ -56,13 +52,6 @@ impl SystemInner {
}
}

// Useful for parallel iterations.
struct Wrap<T>(T);

#[allow(clippy::non_send_fields_in_send_ty)]
unsafe impl<T> Send for Wrap<T> {}
unsafe impl<T> Sync for Wrap<T> {}

/// Calculates system boot time in seconds with improved precision.
/// Uses nanoseconds throughout to avoid rounding errors in uptime calculation,
/// converting to seconds only at the end for stable results. Result is capped
Expand Down Expand Up @@ -232,132 +221,83 @@ impl SystemInner {
}
};

// Windows 10 notebook requires at least 512KiB of memory to make it in one go
let mut buffer_size = 512 * 1024;
let mut process_information: Vec<u8> = Vec::with_capacity(buffer_size);
let now = get_now();

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);

match NtQuerySystemInformation(
SystemProcessInformation,
process_information.as_mut_ptr() as *mut _,
buffer_size as _,
&mut cb_needed,
)
.ok()
{
Ok(()) => break,
Err(err) if err.code() == STATUS_INFO_LENGTH_MISMATCH.to_hresult() => {
// GetNewBufferSize
if cb_needed == 0 {
buffer_size *= 2;
continue;
}
// allocating a few more kilo bytes just in case there are some new process
// kicked in since new call to NtQuerySystemInformation
buffer_size = (cb_needed + (1024 * 10)) as usize;
continue;
}
Err(_err) => {
sysinfo_debug!(
"Couldn't get process infos: NtQuerySystemInformation returned {}",
_err,
);
return 0;
}
}
let nb_cpus = if refresh_kind.cpu() {
self.cpus.len() as u64
} else {
0
};

// Use the amazing and cool CreateToolhelp32Snapshot function.
// Take a snapshot of all running processes. Match the result to an error
let snapshot = match unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) } {
Ok(handle) => handle,
Err(_err) => {
sysinfo_debug!(
"Error capturing process snapshot: CreateToolhelp32Snapshot returned {}",
_err
);
return 0;
}
};

// If we reach this point NtQuerySystemInformation succeeded
// and the buffer contents are initialized
process_information.set_len(buffer_size);
// https://learn.microsoft.com/en-us/windows/win32/api/tlhelp32/ns-tlhelp32-processentry32w
// Microsoft documentation states that for PROCESSENTRY32W, before calling Process32FirstW,
// the 'dwSize' field MUST be set to the size of the PROCESSENTRY32W. Otherwise, Process32FirstW fails.
let mut process_entry = PROCESSENTRY32W {
dwSize: size_of::<PROCESSENTRY32W>() as u32,
..Default::default()
};

let nb_updated = AtomicUsize::new(0);
let mut num_procs = 0; // keep track of the number of updated processes
let process_list = &mut self.process_list;

// 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 the first process
unsafe {
if let Err(_error) = Process32FirstW(snapshot, &mut process_entry) {
sysinfo_debug!("Process32FirstW has failed: {_error:?}");
return 0;
}
}

// 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);
// Iterate over processes in the snapshot.
// Use Process32NextW to process the next PROCESSENTRY32W in the snapshot
loop {
let proc_id = Pid::from_u32(process_entry.th32ProcessID);

if filter_callback(Pid(pi.UniqueProcessId as _), filter_array) {
process_ids.push(Wrap(p));
}
if filter_callback(proc_id, filter_array) {
// exists already
if let Some(p) = process_list.get_mut(&proc_id) {
// Update with the most recent information
let p = &mut p.inner;
p.update(refresh_kind, nb_cpus, now, false);

// Update parent process
let parent = if process_entry.th32ParentProcessID == 0 {
None
} else {
Some(Pid::from_u32(process_entry.th32ParentProcessID))
};

if pi.NextEntryOffset == 0 {
break;
p.parent = parent;
} else {
// Make a new 'ProcessInner' using the Windows PROCESSENTRY32W struct.
let mut p = ProcessInner::from_process_entry(&process_entry, now);
p.update(refresh_kind, nb_cpus, now, false);
process_list.insert(proc_id, Process { inner: p });
}

process_information_offset += pi.NextEntryOffset as isize;
num_procs += 1;
}
let process_list = Wrap(UnsafeCell::new(&mut self.process_list));
let nb_cpus = if refresh_kind.cpu() {
self.cpus.len() as u64
} else {
0
};

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| {
nb_updated.fetch_add(1, Ordering::Relaxed);
// as above, read_unaligned is necessary
let pi = ptr::read_unaligned(pi.0);
let pid = Pid(pi.UniqueProcessId as _);
let ppid: usize = pi.InheritedFromUniqueProcessId as _;
let parent = if ppid != 0 {
Some(Pid(pi.InheritedFromUniqueProcessId as _))
} else {
None
};
// Not sure why we need to make this
let process_list: &Wrap<UnsafeCell<&mut HashMap<Pid, Process>>> = &process_list;
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)
{
proc_.update(refresh_kind, nb_cpus, now, false, &pi);
// Update the parent in case it changed.
proc_.parent = parent;
return None;
}
// If the PID owner changed, we need to recompute the whole process.
sysinfo_debug!("owner changed for PID {}", pid);
}
let name = get_process_name(&pi, pid);
let mut p = ProcessInner::new(pid, parent, now, name);
p.update(refresh_kind, nb_cpus, now, false, &pi);
Some(Process { inner: p })
})
.collect::<Vec<_>>();
for p in processes.into_iter() {
self.process_list.insert(p.pid(), p);

// nothing else to process
if unsafe { Process32NextW(snapshot, &mut process_entry).is_err() } {
break;
}
nb_updated.into_inner()
}

num_procs
}

pub(crate) fn processes(&self) -> &HashMap<Pid, Process> {
Expand Down Expand Up @@ -514,30 +454,6 @@ pub(crate) fn is_proc_running(handle: HANDLE) -> bool {
&& exit_code == STILL_ACTIVE.0 as u32
}

#[allow(clippy::size_of_in_element_count)]
//^ needed for "name.Length as usize / std::mem::size_of::<u16>()"
pub(crate) fn get_process_name(process: &SYSTEM_PROCESS_INFORMATION, process_id: Pid) -> OsString {
let name = &process.ImageName;
if name.Buffer.is_null() {
match process_id.0 {
0 => "Idle".to_owned(),
4 => "System".to_owned(),
_ => format!("<no name> Process {process_id}"),
}
.into()
} else {
unsafe {
let slice = std::slice::from_raw_parts(
name.Buffer,
// The length is in bytes, not the length of string
name.Length as usize / std::mem::size_of::<u16>(),
);

OsString::from_wide(slice)
}
}
}

fn get_dns_hostname() -> Option<String> {
let mut buffer_size = 0;
// Running this first to get the buffer size since the DNS name can be longer than MAX_COMPUTERNAME_LENGTH
Expand Down

0 comments on commit 6749d8a

Please sign in to comment.