Skip to main content

ib_hook/process/
mod.rs

1/*!
2Process utilities.
3*/
4#[cfg(feature = "sysinfo")]
5use std::path::PathBuf;
6use std::time::SystemTime;
7
8use derive_more::{Deref, Display};
9use windows::Win32::{
10    Foundation::{GetLastError, HWND, WIN32_ERROR},
11    System::Threading::{
12        GetProcessIdOfThread, GetProcessTimes, OpenProcess, OpenThread,
13        PROCESS_QUERY_LIMITED_INFORMATION, THREAD_QUERY_LIMITED_INFORMATION,
14    },
15    UI::WindowsAndMessaging::GetWindowThreadProcessId,
16};
17
18use crate::log::*;
19
20mod gui;
21pub mod module;
22
23pub use gui::*;
24
25/// Process ID.
26#[derive(Clone, Copy, PartialEq, Eq, Hash, Display, Debug, Deref)]
27#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
28pub struct Pid(pub u32);
29
30impl Pid {
31    pub fn current() -> Self {
32        Self(std::process::id())
33    }
34
35    pub fn from_tid(tid: u32) -> windows::core::Result<Self> {
36        let thread = unsafe { OpenThread(THREAD_QUERY_LIMITED_INFORMATION, false, tid) }?;
37        match unsafe { GetProcessIdOfThread(thread) } {
38            0 => Err(windows::core::Error::from_thread()),
39            pid => Ok(Pid(pid)),
40        }
41    }
42
43    fn from_hwnd_with_thread(hwnd: HWND) -> Result<(Self, u32), WIN32_ERROR> {
44        let mut pid: u32 = 0;
45        let tid = unsafe { GetWindowThreadProcessId(hwnd, Some(&mut pid)) };
46        if tid != 0 {
47            Ok((Pid(pid), tid))
48        } else {
49            Err(unsafe { GetLastError() })
50        }
51    }
52
53    /// Gets the process ID from a window handle.
54    ///
55    /// This uses [`GetWindowThreadProcessId`] to retrieve the PID associated with a window.
56    pub fn from_hwnd(hwnd: HWND) -> Result<Self, WIN32_ERROR> {
57        Self::try_from(hwnd)
58    }
59}
60
61impl TryFrom<HWND> for Pid {
62    type Error = WIN32_ERROR;
63
64    /// Gets the process ID from a window handle.
65    ///
66    /// This uses [`GetWindowThreadProcessId`] to retrieve the PID associated with a window.
67    fn try_from(hwnd: HWND) -> Result<Self, Self::Error> {
68        Self::from_hwnd_with_thread(hwnd).map(|(pid, _tid)| pid)
69    }
70}
71
72impl Pid {
73    /// Gets the start time of the process.
74    pub fn get_start_time(self) -> windows::core::Result<SystemTime> {
75        let pid = self.0;
76        let handle = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid) }?;
77
78        let mut start: nt_time::FileTime = Default::default();
79        let mut x = Default::default();
80        unsafe { GetProcessTimes(handle, &mut start as *mut _ as _, &mut x, &mut x, &mut x) }?;
81
82        Ok(start.into())
83    }
84
85    /// Gets the start time of the process.
86    ///
87    /// [`#![feature(time_systemtime_limits)]`](https://github.com/rust-lang/rust/issues/149067) is not stable yet.
88    pub fn get_start_time_or_max(self) -> SystemTime {
89        self.get_start_time()
90            .inspect_err(|e| debug!(%e, "get_start_time"))
91            .unwrap_or_else(|_| nt_time::FileTime::MAX.into())
92    }
93}
94
95#[cfg(feature = "sysinfo")]
96impl Pid {
97    pub fn with_process<R>(
98        self,
99        refresh_info: sysinfo::ProcessRefreshKind,
100        f: impl FnOnce(&sysinfo::Process) -> R,
101    ) -> Option<R> {
102        let pid = self.clone().into();
103        let mut system = sysinfo::System::new();
104        system.refresh_processes_specifics(
105            sysinfo::ProcessesToUpdate::Some(&[pid]),
106            false,
107            refresh_info,
108        );
109        system.process(pid).map(f)
110    }
111
112    pub fn image_path(self) -> Option<PathBuf> {
113        self.with_process(
114            sysinfo::ProcessRefreshKind::nothing().with_exe(sysinfo::UpdateKind::Always),
115            |p| p.exe().map(|p| p.to_owned()),
116        )
117        .flatten()
118    }
119}
120
121#[cfg(feature = "sysinfo")]
122impl From<sysinfo::Pid> for Pid {
123    fn from(pid: sysinfo::Pid) -> Self {
124        Self(pid.as_u32())
125    }
126}
127
128#[cfg(feature = "sysinfo")]
129impl Into<sysinfo::Pid> for Pid {
130    fn into(self) -> sysinfo::Pid {
131        sysinfo::Pid::from_u32(self.0)
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    use windows::Win32::UI::WindowsAndMessaging::GetDesktopWindow;
140
141    #[test]
142    fn pid_from_hwnd() {
143        // Test with a simple window - use GetDesktopWindow as it always exists
144        let desktop = unsafe { GetDesktopWindow() };
145        let pid = Pid::from_hwnd(desktop);
146
147        // Desktop window should have a valid PID (usually 0 or the session manager)
148        assert!(pid.is_ok(), "Should be able to get PID from desktop window");
149        println!("Desktop window PID: {:?}", pid.unwrap());
150    }
151
152    #[test]
153    #[cfg(feature = "sysinfo")]
154    fn pid_from_hwnd_desktop() {
155        let desktop = unsafe { GetDesktopWindow() };
156        dbg!(desktop);
157        let pid = Pid::from_hwnd(desktop).expect("Should get PID from desktop window");
158        dbg!(pid);
159
160        let exe_path = pid.image_path();
161        // Should be C:\Windows\System32\csrss.exe if permission enough
162        assert_eq!(exe_path, None);
163    }
164
165    /// Test with GetDesktopWindow and verify the process is explorer.exe
166    #[test]
167    #[cfg(feature = "sysinfo")]
168    fn pid_from_hwnd_shell() {
169        use std::path::Path;
170        use windows::Win32::UI::WindowsAndMessaging::GetShellWindow;
171
172        let desktop = unsafe { GetShellWindow() };
173        dbg!(desktop);
174        let pid = Pid::from_hwnd(desktop).expect("Should get PID from desktop window");
175        dbg!(pid);
176
177        // The desktop window is typically owned by explorer.exe
178        let exe_path = pid.image_path().unwrap();
179        assert_eq!(exe_path, Path::new(r"C:\Windows\explorer.exe"));
180    }
181
182    #[test]
183    fn get_start_time() {
184        let current_pid = std::process::id();
185        let pid = Pid(current_pid);
186
187        let start_time = pid.get_start_time().unwrap();
188        dbg!(start_time);
189        let t = start_time
190            .duration_since(SystemTime::UNIX_EPOCH)
191            .unwrap()
192            .as_secs();
193        assert!(t > 0, "{t}");
194
195        let start_time2 = pid.get_start_time().unwrap();
196        dbg!(start_time2);
197        assert_eq!(start_time, start_time2);
198    }
199}