win32_ecoqos/utils/
process.rs

1use std::{ffi::OsString, os::windows::ffi::OsStringExt};
2
3use windows::Win32::{
4    Foundation::{CloseHandle, HANDLE},
5    System::Diagnostics::ToolHelp::{
6        CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W,
7        TH32CS_SNAPPROCESS,
8    },
9};
10
11#[derive(Debug, PartialEq, Eq)]
12#[non_exhaustive]
13/// process information from snapshot.
14pub struct Process {
15    /// win32 process id
16    pub process_id: u32,
17    /// win32 process id of it's parent
18    pub process_parent_id: u32,
19    /// win32 process name, null bytes trimmed.
20    pub process_name: OsString,
21}
22
23/// snapshot of processes
24#[derive(Debug)]
25pub struct Processes {
26    snapshot: HANDLE,
27    last_entry: Option<PROCESSENTRY32W>,
28}
29
30impl Drop for Processes {
31    fn drop(&mut self) {
32        let _ = unsafe { CloseHandle(self.snapshot) };
33    }
34}
35
36impl Processes {
37    /// try to capture a snapshot of processes.
38    pub fn try_new() -> windows_result::Result<Self> {
39        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }?;
40        Ok(Self {
41            snapshot,
42            last_entry: None,
43        })
44    }
45}
46
47impl Iterator for Processes {
48    type Item = Process;
49
50    fn next(&mut self) -> Option<Self::Item> {
51        let first = self.last_entry.is_none();
52        let mut entry = self.last_entry.take().unwrap_or(PROCESSENTRY32W {
53            dwSize: size_of::<PROCESSENTRY32W>() as u32,
54            ..Default::default()
55        });
56
57        unsafe {
58            if first {
59                Process32FirstW(self.snapshot, &mut entry as *mut _)
60            } else {
61                Process32NextW(self.snapshot, &mut entry as *mut _)
62            }
63        }
64        .ok()
65        .map(|_| entry)
66        .inspect(|entry| self.last_entry = Some(*entry))
67        .map(
68            |PROCESSENTRY32W {
69                 th32ProcessID,
70                 th32ParentProcessID,
71                 szExeFile,
72                 ..
73             }| {
74                let null = szExeFile
75                    .iter()
76                    .enumerate()
77                    .find_map(|(idx, ch)| if ch == &0 { Some(idx) } else { None })
78                    .unwrap_or(szExeFile.len());
79
80                let process_name = OsString::from_wide(&szExeFile[..null]);
81                Process {
82                    process_id: th32ProcessID,
83                    process_parent_id: th32ParentProcessID,
84                    process_name,
85                }
86            },
87        )
88    }
89}