win_wrap/
hook.rs

1/*
2 * Copyright (c) 2023. The RigelA open source project team and
3 * its contributors reserve all rights.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 * http://www.apache.org/licenses/LICENSE-2.0
9 * Unless required by applicable law or agreed to in writing, software distributed under the
10 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 * See the License for the specific language governing permissions and limitations under the License.
12 */
13
14use crate::{
15    common::{
16        call_next_hook_ex, set_windows_hook_ex, unhook_windows_hook_ex, LPARAM, LRESULT,
17        WINDOWS_HOOK_ID, WPARAM,
18    },
19    message::message_loop,
20    threading::{get_current_thread_id, ThreadNotify},
21};
22use std::{
23    collections::HashMap,
24    error::Error,
25    fmt::{Debug, Display, Formatter},
26    sync::{LazyLock, Mutex, PoisonError, RwLock},
27    thread::{sleep, spawn},
28    time::{Duration, SystemTime, SystemTimeError, UNIX_EPOCH},
29};
30use windows::Win32::UI::WindowsAndMessaging::{
31    CWPRETSTRUCT, CWPSTRUCT, KBDLLHOOKSTRUCT, KBDLLHOOKSTRUCT_FLAGS, MSLLHOOKSTRUCT,
32    WH_CALLWNDPROC, WH_CALLWNDPROCRET, WH_GETMESSAGE, WH_KEYBOARD_LL, WH_MOUSE_LL,
33};
34pub use windows::Win32::UI::WindowsAndMessaging::{
35    LLKHF_ALTDOWN, LLKHF_EXTENDED, LLKHF_INJECTED, LLKHF_LOWER_IL_INJECTED, LLKHF_UP,
36};
37
38/**
39钩子类型。
40*/
41pub type HookType = WINDOWS_HOOK_ID;
42
43/**
44低级键盘钩子。
45*/
46pub const HOOK_TYPE_KEYBOARD_LL: HookType = WH_KEYBOARD_LL;
47
48/**
49低级鼠标钩子。
50*/
51pub const HOOK_TYPE_MOUSE_LL: HookType = WH_MOUSE_LL;
52
53/**
54当SendMessage()把消息交给WndProc时,在WndProc尚未执行前,系统调用CallWndProc钩子函数,钩子函数执行后才执行窗口过程。
55*/
56pub const HOOK_TYPE_CALL_WND_PROC: WINDOWS_HOOK_ID = WH_CALLWNDPROC;
57
58/**
59当SendMessage()把消息交给WndProc时,在WndProc执行完毕,系统调用CallWndProcRet钩子函数,从而可以拦截窗口过程函数的结果。
60*/
61pub const HOOK_TYPE_CALL_WND_PROC_RET: WINDOWS_HOOK_ID = WH_CALLWNDPROCRET;
62
63/**
64拦截队列消息(拦截由get_message或者post_message或者peek_message的队列消息)。
65*/
66pub const HOOK_TYPE_GET_MESSAGE: WINDOWS_HOOK_ID = WH_GETMESSAGE;
67
68/**
69低级键盘钩子信息结构。
70*/
71pub type KbdLlHookStruct = KBDLLHOOKSTRUCT;
72pub type KbdLlHookFlags = KBDLLHOOKSTRUCT_FLAGS;
73
74/**
75低级鼠标钩子信息结构。
76*/
77pub type MsLlHookStruct = MSLLHOOKSTRUCT;
78
79/**
80窗口过程函数之前的钩子信息结构。
81*/
82pub type CwpStruct = CWPSTRUCT;
83
84/**
85窗口过程函数之后的钩子信息结构。
86*/
87pub type CwpRStruct = CWPRETSTRUCT;
88
89static H_HOOK: RwLock<Option<HashMap<i32, ThreadNotify>>> = RwLock::new(None);
90static HOOK_MAP: LazyLock<Mutex<HashMap<i32, (Vec<WindowsHook>, usize)>>> =
91    LazyLock::new(Default::default);
92static HOOK_CALLBACKS: LazyLock<
93    Mutex<
94        HashMap<
95            Duration,
96            Box<
97                dyn FnMut(WindowsHook, WPARAM, LPARAM) -> Result<LRESULT, WindowsHookError>
98                    + Send
99                    + Sync,
100            >,
101        >,
102    >,
103> = LazyLock::new(Default::default);
104
105/// Windows钩子错误
106#[derive(Debug)]
107pub enum WindowsHookError {
108    Dropped,
109    Poison(String),
110    SystemTime(SystemTimeError),
111}
112
113impl Display for WindowsHookError {
114    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
115        match self {
116            Self::Dropped => write!(f, "Windows hook has been dropped."),
117            Self::Poison(e) => Display::fmt(e, f),
118            Self::SystemTime(e) => Display::fmt(e, f),
119        }
120    }
121}
122
123impl Error for WindowsHookError {}
124
125impl<T> From<PoisonError<T>> for WindowsHookError {
126    fn from(value: PoisonError<T>) -> Self {
127        Self::Poison(value.to_string())
128    }
129}
130
131impl From<SystemTimeError> for WindowsHookError {
132    fn from(value: SystemTimeError) -> Self {
133        Self::SystemTime(value)
134    }
135}
136
137/**
138Windows 的钩子。
139*/
140#[derive(Clone)]
141pub struct WindowsHook {
142    hook_type: WINDOWS_HOOK_ID,
143    id: Duration,
144    code: i32,
145}
146
147impl WindowsHook {
148    fn prepare(&mut self, code: i32) {
149        self.code = code;
150    }
151
152    fn call(self, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {
153        let mut lock = HOOK_CALLBACKS.lock()?;
154        let Some(cb) = lock.get_mut(&self.id) else {
155            return Err(WindowsHookError::Dropped);
156        };
157        let hook = self.clone();
158        cb(hook, w_param, l_param)
159    }
160
161    /// 调用钩子链中下一个函数
162    pub fn next_hook(self, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {
163        let mut lock = HOOK_MAP.lock()?;
164        let (mut vec, index) = match lock.get_mut(&self.hook_type.0) {
165            None => return Ok(call_next_hook_ex(self.code, w_param, l_param)),
166            Some(v) => {
167                let i = v.1;
168                v.1 += 1;
169                (v.0.clone(), i)
170            }
171        };
172        drop(lock);
173
174        if index >= vec.len() {
175            return Ok(call_next_hook_ex(self.code, w_param, l_param));
176        }
177
178        let hook = vec.get_mut(index);
179
180        match hook {
181            None => Ok(call_next_hook_ex(self.code, w_param, l_param)),
182            Some(x) => x.clone().call(w_param, l_param),
183        }
184    }
185
186    //noinspection SpellCheckingInspection
187    /**
188    创建一个钩子对象并开始监听事件。
189    `hook_type` 钩子类型,使用“HOOK_TYPE_”开头的常亮。
190    `cb` 用于接收事件的处理函数,需要实现三个参数:
191    ```
192    use win_wrap::{common::{LPARAM, WPARAM, LRESULT}, hook::{WindowsHookError,WindowsHook}};
193    fn cb(hook: WindowsHook, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {hook.next_hook(w_param,l_param)}
194    ```
195    也可以使用闭包
196    ```
197    use win_wrap::{common::{LPARAM, WPARAM}, hook::{WindowsHookError,WindowsHook}};
198    |hook: WindowsHook, w_param: WPARAM, l_param: LPARAM| {hook.next_hook(w_param,l_param)};
199    ```
200    处理函数运行在单独的子线程中,使用时可以通过next_hook()调用钩子链中的下一个函数,也可以直接拦截。
201    */
202    pub fn new<F>(hook_type: HookType, cb: F) -> Result<Self, WindowsHookError>
203    where
204        F: FnMut(WindowsHook, WPARAM, LPARAM) -> Result<LRESULT, WindowsHookError>
205            + Send
206            + Sync
207            + 'static,
208    {
209        let id = SystemTime::now().duration_since(UNIX_EPOCH)?;
210        let info = Self {
211            hook_type,
212            id,
213            code: 0,
214        };
215
216        let mut lock = HOOK_CALLBACKS.lock()?;
217        lock.insert(id, Box::new(cb));
218        drop(lock);
219
220        let mut lock = HOOK_MAP.lock()?;
221        match lock.get_mut(&hook_type.0) {
222            None => {
223                install(hook_type);
224                let mut vec = Vec::new();
225                vec.insert(0, info.clone());
226                lock.insert(hook_type.0, (vec, 0));
227            }
228            Some((vec, _)) => {
229                vec.insert(0, info.clone());
230            }
231        };
232        Ok(info)
233    }
234
235    /**
236    卸载钩子。
237    */
238    pub fn unhook(&self) -> Result<(), WindowsHookError> {
239        let mut lock = HOOK_CALLBACKS.lock()?;
240        lock.remove(&self.id);
241
242        let mut lock = HOOK_MAP.lock()?;
243        match lock.get(&self.hook_type.0) {
244            None => {
245                uninstall(self.hook_type);
246            }
247            Some((x, _)) => {
248                let mut x = x.clone();
249                for i in 0..x.len() {
250                    let Some(j) = x.get(i) else {
251                        continue;
252                    };
253
254                    if j == self {
255                        x.remove(i);
256                        break;
257                    }
258                }
259                if x.len() < 1 {
260                    uninstall(self.hook_type);
261                    lock.remove(&self.hook_type.0);
262                } else {
263                    lock.insert(self.hook_type.0, (x, 0));
264                }
265            }
266        }
267
268        Ok(())
269    }
270}
271
272impl PartialEq for WindowsHook {
273    fn eq(&self, other: &Self) -> bool {
274        self.id == other.id
275    }
276}
277
278impl Debug for WindowsHook {
279    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
280        write!(f, "WindowsHook({})", self.hook_type.0)
281    }
282}
283
284macro_rules! define_hook_proc {
285    ($name:tt, $r#type:tt) => {
286        unsafe extern "system" fn $name(code: i32, w_param: WPARAM, l_param: LPARAM) -> LRESULT {
287            if code < 0 {
288                return call_next_hook_ex(code, w_param, l_param);
289            }
290
291            let mut lock = match HOOK_MAP.lock() {
292                Err(_) => return call_next_hook_ex(code, w_param, l_param),
293                Ok(lock) => lock,
294            };
295
296            let mut hook = match lock.get_mut(&$r#type.0) {
297                None => {
298                    drop(lock);
299                    return call_next_hook_ex(code, w_param, l_param);
300                }
301                Some((x, i)) => {
302                    *i = 0;
303                    match x.get_mut(0) {
304                        None => {
305                            drop(lock);
306                            return call_next_hook_ex(code, w_param, l_param);
307                        }
308                        Some(x) => x.clone(),
309                    }
310                }
311            };
312            drop(lock);
313
314            hook.prepare(code);
315            if let Ok(r) = hook.next_hook(w_param, l_param) {
316                return r;
317            }
318            call_next_hook_ex(code, w_param, l_param)
319        }
320    };
321}
322define_hook_proc!(proc_keyboard_ll, HOOK_TYPE_KEYBOARD_LL);
323define_hook_proc!(proc_mouse_ll, HOOK_TYPE_MOUSE_LL);
324define_hook_proc!(proc_call_wnd_proc, HOOK_TYPE_CALL_WND_PROC);
325define_hook_proc!(proc_call_wnd_proc_ret, HOOK_TYPE_CALL_WND_PROC_RET);
326define_hook_proc!(proc_get_message, HOOK_TYPE_GET_MESSAGE);
327
328fn install(hook_type: HookType) {
329    let lock = H_HOOK.read().unwrap();
330    if let Some(x) = lock.as_ref() {
331        if x.contains_key(&hook_type.0) {
332            return;
333        }
334    }
335    drop(lock);
336
337    spawn(move || {
338        let proc = match hook_type {
339            HOOK_TYPE_KEYBOARD_LL => proc_keyboard_ll,
340            HOOK_TYPE_MOUSE_LL => proc_mouse_ll,
341            HOOK_TYPE_CALL_WND_PROC => proc_call_wnd_proc,
342            HOOK_TYPE_CALL_WND_PROC_RET => proc_call_wnd_proc_ret,
343            HOOK_TYPE_GET_MESSAGE => proc_get_message,
344            x => panic!("Unsupported hook type: {}.", x.0),
345        };
346        let mut retry = 0;
347        let h_hook = loop {
348            let h = set_windows_hook_ex(hook_type, Some(proc), None, 0);
349            if h.is_ok() {
350                break h.unwrap();
351            }
352            if retry > 5 {
353                panic!(
354                    "Can't set the windows hook ({}), and retrying it.",
355                    hook_type.0
356                );
357            }
358            retry += 1;
359            sleep(Duration::from_millis(1000));
360        };
361        let notify = ThreadNotify::new(get_current_thread_id());
362        let mut lock = H_HOOK.write().unwrap();
363        let mut map = match lock.as_ref() {
364            None => HashMap::new(),
365            Some(x) => x.clone(),
366        };
367        map.insert(hook_type.0, notify.clone());
368        *lock = Some(map);
369        drop(lock);
370        message_loop(|_| ());
371        unhook_windows_hook_ex(h_hook).unwrap_or(());
372        notify.finish();
373    });
374}
375
376fn uninstall(hook_type: HookType) {
377    let lock = H_HOOK.read().unwrap();
378    match lock.as_ref() {
379        None => {
380            return;
381        }
382        Some(x) => {
383            if let Some(x) = x.get(&hook_type.0) {
384                x.quit();
385                x.join(5000);
386            }
387        }
388    }
389}
390
391#[cfg(test)]
392mod test_hook {
393    use crate::{
394        ext::LParamExt,
395        hook::{KbdLlHookStruct, WindowsHook, WindowsHookError, HOOK_TYPE_KEYBOARD_LL},
396    };
397
398    #[test]
399    fn main() -> Result<(), WindowsHookError> {
400        let hook = WindowsHook::new(HOOK_TYPE_KEYBOARD_LL, |h, w, l| {
401            let key = l.to::<KbdLlHookStruct>();
402            println!("{}", key.vkCode);
403            h.next_hook(w, l)
404        })?;
405        std::thread::sleep(std::time::Duration::from_millis(5000));
406        hook.unhook()
407    }
408}