ib_hook/process/
mod.rs

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

ib_hook/process/
mod.rs
