win-wrap 0.3.4

/*
 * Copyright (c) 2023. The RigelA open source project team and
 * its contributors reserve all rights.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software distributed under the
 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */

use crate::{
    common::{
        call_next_hook_ex, set_windows_hook_ex, unhook_windows_hook_ex, LPARAM, LRESULT,
        WINDOWS_HOOK_ID, WPARAM,
    },
    message::message_loop,
    threading::{get_current_thread_id, ThreadNotify},
};
use std::{
    collections::HashMap,
    error::Error,
    fmt::{Debug, Display, Formatter},
    sync::{LazyLock, Mutex, PoisonError, RwLock},
    thread::{sleep, spawn},
    time::{Duration, SystemTime, SystemTimeError, UNIX_EPOCH},
};
use windows::Win32::UI::WindowsAndMessaging::{
    CWPRETSTRUCT, CWPSTRUCT, KBDLLHOOKSTRUCT, KBDLLHOOKSTRUCT_FLAGS, MSLLHOOKSTRUCT,
    WH_CALLWNDPROC, WH_CALLWNDPROCRET, WH_GETMESSAGE, WH_KEYBOARD_LL, WH_MOUSE_LL,
};
pub use windows::Win32::UI::WindowsAndMessaging::{
    LLKHF_ALTDOWN, LLKHF_EXTENDED, LLKHF_INJECTED, LLKHF_LOWER_IL_INJECTED, LLKHF_UP,
};

/**

钩子类型。
*/
pub type HookType = WINDOWS_HOOK_ID;

/**

低级键盘钩子。
*/
pub const HOOK_TYPE_KEYBOARD_LL: HookType = WH_KEYBOARD_LL;

/**

低级鼠标钩子。
*/
pub const HOOK_TYPE_MOUSE_LL: HookType = WH_MOUSE_LL;

/**

当SendMessage()把消息交给WndProc时,在WndProc尚未执行前,系统调用CallWndProc钩子函数,钩子函数执行后才执行窗口过程。
*/
pub const HOOK_TYPE_CALL_WND_PROC: WINDOWS_HOOK_ID = WH_CALLWNDPROC;

/**

当SendMessage()把消息交给WndProc时,在WndProc执行完毕,系统调用CallWndProcRet钩子函数,从而可以拦截窗口过程函数的结果。
*/
pub const HOOK_TYPE_CALL_WND_PROC_RET: WINDOWS_HOOK_ID = WH_CALLWNDPROCRET;

/**

拦截队列消息（拦截由get_message或者post_message或者peek_message的队列消息）。
*/
pub const HOOK_TYPE_GET_MESSAGE: WINDOWS_HOOK_ID = WH_GETMESSAGE;

/**

低级键盘钩子信息结构。
*/
pub type KbdLlHookStruct = KBDLLHOOKSTRUCT;
pub type KbdLlHookFlags = KBDLLHOOKSTRUCT_FLAGS;

/**

低级鼠标钩子信息结构。
*/
pub type MsLlHookStruct = MSLLHOOKSTRUCT;

/**

窗口过程函数之前的钩子信息结构。
*/
pub type CwpStruct = CWPSTRUCT;

/**

窗口过程函数之后的钩子信息结构。
*/
pub type CwpRStruct = CWPRETSTRUCT;

static H_HOOK: RwLock<Option<HashMap<i32, ThreadNotify>>> = RwLock::new(None);
static HOOK_MAP: LazyLock<Mutex<HashMap<i32, (Vec<WindowsHook>, usize)>>> =
    LazyLock::new(Default::default);
static HOOK_CALLBACKS: LazyLock<
    Mutex<
        HashMap<
            Duration,
            Box<
                dyn FnMut(WindowsHook, WPARAM, LPARAM) -> Result<LRESULT, WindowsHookError>
                    + Send
                    + Sync,
            >,
        >,
    >,
> = LazyLock::new(Default::default);

/// Windows钩子错误
#[derive(Debug)]
pub enum WindowsHookError {
    Dropped,
    Poison(String),
    SystemTime(SystemTimeError),
}

impl Display for WindowsHookError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Dropped => write!(f, "Windows hook has been dropped."),
            Self::Poison(e) => Display::fmt(e, f),
            Self::SystemTime(e) => Display::fmt(e, f),
        }
    }
}

impl Error for WindowsHookError {}

impl<T> From<PoisonError<T>> for WindowsHookError {
    fn from(value: PoisonError<T>) -> Self {
        Self::Poison(value.to_string())
    }
}

impl From<SystemTimeError> for WindowsHookError {
    fn from(value: SystemTimeError) -> Self {
        Self::SystemTime(value)
    }
}

/**

Windows 的钩子。
*/
#[derive(Clone)]
pub struct WindowsHook {
    hook_type: WINDOWS_HOOK_ID,
    id: Duration,
    code: i32,
}

impl WindowsHook {
    fn prepare(&mut self, code: i32) {
        self.code = code;
    }

    fn call(self, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {
        let mut lock = HOOK_CALLBACKS.lock()?;
        let Some(cb) = lock.get_mut(&self.id) else {
            return Err(WindowsHookError::Dropped);
        };
        let hook = self.clone();
        cb(hook, w_param, l_param)
    }

    /// 调用钩子链中下一个函数
    pub fn next_hook(self, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {
        let mut lock = HOOK_MAP.lock()?;
        let (mut vec, index) = match lock.get_mut(&self.hook_type.0) {
            None => return Ok(call_next_hook_ex(self.code, w_param, l_param)),
            Some(v) => {
                let i = v.1;
                v.1 += 1;
                (v.0.clone(), i)
            }
        };
        drop(lock);

        if index >= vec.len() {
            return Ok(call_next_hook_ex(self.code, w_param, l_param));
        }

        let hook = vec.get_mut(index);

        match hook {
            None => Ok(call_next_hook_ex(self.code, w_param, l_param)),
            Some(x) => x.clone().call(w_param, l_param),
        }
    }

    //noinspection SpellCheckingInspection
    /**

    创建一个钩子对象并开始监听事件。
    `hook_type` 钩子类型，使用“HOOK_TYPE_”开头的常亮。
    `cb` 用于接收事件的处理函数，需要实现三个参数：
    ```
    use win_wrap::{common::{LPARAM, WPARAM, LRESULT}, hook::{WindowsHookError,WindowsHook}};
    fn cb(hook: WindowsHook, w_param: WPARAM, l_param: LPARAM) -> Result<LRESULT, WindowsHookError> {hook.next_hook(w_param,l_param)}
    ```
    也可以使用闭包
    ```
    use win_wrap::{common::{LPARAM, WPARAM}, hook::{WindowsHookError,WindowsHook}};
    |hook: WindowsHook, w_param: WPARAM, l_param: LPARAM| {hook.next_hook(w_param,l_param)};
    ```
    处理函数运行在单独的子线程中，使用时可以通过next_hook()调用钩子链中的下一个函数，也可以直接拦截。
    */
    pub fn new<F>(hook_type: HookType, cb: F) -> Result<Self, WindowsHookError>
    where
        F: FnMut(WindowsHook, WPARAM, LPARAM) -> Result<LRESULT, WindowsHookError>
            + Send
            + Sync
            + 'static,
    {
        let id = SystemTime::now().duration_since(UNIX_EPOCH)?;
        let info = Self {
            hook_type,
            id,
            code: 0,
        };

        let mut lock = HOOK_CALLBACKS.lock()?;
        lock.insert(id, Box::new(cb));
        drop(lock);

        let mut lock = HOOK_MAP.lock()?;
        match lock.get_mut(&hook_type.0) {
            None => {
                install(hook_type);
                let mut vec = Vec::new();
                vec.insert(0, info.clone());
                lock.insert(hook_type.0, (vec, 0));
            }
            Some((vec, _)) => {
                vec.insert(0, info.clone());
            }
        };
        Ok(info)
    }

    /**

    卸载钩子。
    */
    pub fn unhook(&self) -> Result<(), WindowsHookError> {
        let mut lock = HOOK_CALLBACKS.lock()?;
        lock.remove(&self.id);

        let mut lock = HOOK_MAP.lock()?;
        match lock.get(&self.hook_type.0) {
            None => {
                uninstall(self.hook_type);
            }
            Some((x, _)) => {
                let mut x = x.clone();
                for i in 0..x.len() {
                    let Some(j) = x.get(i) else {
                        continue;
                    };

                    if j == self {
                        x.remove(i);
                        break;
                    }
                }
                if x.len() < 1 {
                    uninstall(self.hook_type);
                    lock.remove(&self.hook_type.0);
                } else {
                    lock.insert(self.hook_type.0, (x, 0));
                }
            }
        }

        Ok(())
    }
}

impl PartialEq for WindowsHook {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Debug for WindowsHook {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "WindowsHook({})", self.hook_type.0)
    }
}

macro_rules! define_hook_proc {
    ($name:tt, $r#type:tt) => {
        unsafe extern "system" fn $name(code: i32, w_param: WPARAM, l_param: LPARAM) -> LRESULT {
            if code < 0 {
                return call_next_hook_ex(code, w_param, l_param);
            }

            let mut lock = match HOOK_MAP.lock() {
                Err(_) => return call_next_hook_ex(code, w_param, l_param),
                Ok(lock) => lock,
            };

            let mut hook = match lock.get_mut(&$r#type.0) {
                None => {
                    drop(lock);
                    return call_next_hook_ex(code, w_param, l_param);
                }
                Some((x, i)) => {
                    *i = 0;
                    match x.get_mut(0) {
                        None => {
                            drop(lock);
                            return call_next_hook_ex(code, w_param, l_param);
                        }
                        Some(x) => x.clone(),
                    }
                }
            };
            drop(lock);

            hook.prepare(code);
            if let Ok(r) = hook.next_hook(w_param, l_param) {
                return r;
            }
            call_next_hook_ex(code, w_param, l_param)
        }
    };
}
define_hook_proc!(proc_keyboard_ll, HOOK_TYPE_KEYBOARD_LL);
define_hook_proc!(proc_mouse_ll, HOOK_TYPE_MOUSE_LL);
define_hook_proc!(proc_call_wnd_proc, HOOK_TYPE_CALL_WND_PROC);
define_hook_proc!(proc_call_wnd_proc_ret, HOOK_TYPE_CALL_WND_PROC_RET);
define_hook_proc!(proc_get_message, HOOK_TYPE_GET_MESSAGE);

fn install(hook_type: HookType) {
    let lock = H_HOOK.read().unwrap();
    if let Some(x) = lock.as_ref() {
        if x.contains_key(&hook_type.0) {
            return;
        }
    }
    drop(lock);

    spawn(move || {
        let proc = match hook_type {
            HOOK_TYPE_KEYBOARD_LL => proc_keyboard_ll,
            HOOK_TYPE_MOUSE_LL => proc_mouse_ll,
            HOOK_TYPE_CALL_WND_PROC => proc_call_wnd_proc,
            HOOK_TYPE_CALL_WND_PROC_RET => proc_call_wnd_proc_ret,
            HOOK_TYPE_GET_MESSAGE => proc_get_message,
            x => panic!("Unsupported hook type: {}.", x.0),
        };
        let mut retry = 0;
        let h_hook = loop {
            let h = set_windows_hook_ex(hook_type, Some(proc), None, 0);
            if h.is_ok() {
                break h.unwrap();
            }
            if retry > 5 {
                panic!(
                    "Can't set the windows hook ({}), and retrying it.",
                    hook_type.0
                );
            }
            retry += 1;
            sleep(Duration::from_millis(1000));
        };
        let notify = ThreadNotify::new(get_current_thread_id());
        let mut lock = H_HOOK.write().unwrap();
        let mut map = match lock.as_ref() {
            None => HashMap::new(),
            Some(x) => x.clone(),
        };
        map.insert(hook_type.0, notify.clone());
        *lock = Some(map);
        drop(lock);
        message_loop(|_| ());
        unhook_windows_hook_ex(h_hook).unwrap_or(());
        notify.finish();
    });
}

fn uninstall(hook_type: HookType) {
    let lock = H_HOOK.read().unwrap();
    match lock.as_ref() {
        None => {
            return;
        }
        Some(x) => {
            if let Some(x) = x.get(&hook_type.0) {
                x.quit();
                x.join(5000);
            }
        }
    }
}

#[cfg(test)]
mod test_hook {
    use crate::{
        ext::LParamExt,
        hook::{KbdLlHookStruct, WindowsHook, WindowsHookError, HOOK_TYPE_KEYBOARD_LL},
    };

    #[test]
    fn main() -> Result<(), WindowsHookError> {
        let hook = WindowsHook::new(HOOK_TYPE_KEYBOARD_LL, |h, w, l| {
            let key = l.to::<KbdLlHookStruct>();
            println!("{}", key.vkCode);
            h.next_hook(w, l)
        })?;
        std::thread::sleep(std::time::Duration::from_millis(5000));
        hook.unhook()
    }
}