diff --git a/Cargo.toml b/Cargo.toml index 7194e10..31f1d92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ block2 = "0.6" windows = { version = "0.58", features = [ "Win32_Foundation", "Win32_UI_WindowsAndMessaging", + "Win32_System_RemoteDesktop", ] } [target.'cfg(target_os = "linux")'.dependencies] diff --git a/src/platform/windows/listener.rs b/src/platform/windows/listener.rs index d32b339..867b86b 100644 --- a/src/platform/windows/listener.rs +++ b/src/platform/windows/listener.rs @@ -5,13 +5,19 @@ use std::sync::mpsc::{self, Sender}; use std::sync::{Arc, Mutex}; use std::thread::{self, JoinHandle}; -use windows::Win32::Foundation::{LPARAM, LRESULT, WPARAM}; +use std::ptr; + +use windows::Win32::Foundation::{HWND, LPARAM, LRESULT, WPARAM}; +use windows::Win32::System::RemoteDesktop::{ + WTSRegisterSessionNotification, WTSUnRegisterSessionNotification, +}; use windows::Win32::UI::WindowsAndMessaging::{ - CallNextHookEx, DispatchMessageW, MsgWaitForMultipleObjects, PeekMessageW, SetWindowsHookExW, - TranslateMessage, UnhookWindowsHookEx, KBDLLHOOKSTRUCT, LLKHF_EXTENDED, MSG, MSLLHOOKSTRUCT, - PM_REMOVE, QS_ALLINPUT, WH_KEYBOARD_LL, WH_MOUSE_LL, WM_KEYDOWN, WM_KEYUP, WM_LBUTTONDOWN, - WM_LBUTTONUP, WM_MBUTTONDOWN, WM_MBUTTONUP, WM_QUIT, WM_RBUTTONDOWN, WM_RBUTTONUP, - WM_SYSKEYDOWN, WM_SYSKEYUP, WM_XBUTTONDOWN, WM_XBUTTONUP, + CallNextHookEx, CreateWindowExW, DestroyWindow, DispatchMessageW, MsgWaitForMultipleObjects, + PeekMessageW, RegisterClassW, SetWindowsHookExW, TranslateMessage, UnhookWindowsHookEx, + HWND_MESSAGE, KBDLLHOOKSTRUCT, LLKHF_EXTENDED, MSG, MSLLHOOKSTRUCT, PM_REMOVE, QS_ALLINPUT, + WH_KEYBOARD_LL, WH_MOUSE_LL, WINDOW_EX_STYLE, WINDOW_STYLE, WM_KEYDOWN, WM_KEYUP, + WM_LBUTTONDOWN, WM_LBUTTONUP, WM_MBUTTONDOWN, WM_MBUTTONUP, WM_QUIT, WM_RBUTTONDOWN, + WM_RBUTTONUP, WM_SYSKEYDOWN, WM_SYSKEYUP, WM_XBUTTONDOWN, WM_XBUTTONUP, WNDCLASSW, }; use crate::error::Result; @@ -22,6 +28,11 @@ use super::keycode::{vk_to_key, vk_to_modifier}; const HOOK_LOOP_TIMEOUT_MS: u32 = 10; +// WTS constants not in the windows crate's generated bindings +const NOTIFY_FOR_THIS_SESSION: u32 = 0; +const WM_WTSSESSION_CHANGE: u32 = 0x02B1; +const WTS_SESSION_UNLOCK: usize = 0x8; + /// Thread-local state for the keyboard hook callback. /// /// Windows low-level hooks require a callback function with a specific signature, @@ -36,18 +47,38 @@ thread_local! { static HOOK_CONTEXT: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } -/// Drain all pending thread messages and return `true` if WM_QUIT was received. -fn drain_thread_messages(msg: &mut MSG) -> bool { +/// Result of draining the thread message queue. +enum DrainResult { + /// Normal -- all messages processed, loop continues. + Continue, + /// WM_QUIT received -- exit the message loop. + Quit, + /// Windows session was unlocked -- hooks need re-installation. + SessionUnlock, +} + +/// Drain all pending thread messages. +fn drain_thread_messages(msg: &mut MSG) -> DrainResult { + let mut session_unlocked = false; unsafe { while PeekMessageW(msg, None, 0, 0, PM_REMOVE).as_bool() { if msg.message == WM_QUIT { - return true; + return DrainResult::Quit; + } + if msg.message == WM_WTSSESSION_CHANGE + && msg.wParam == WPARAM(WTS_SESSION_UNLOCK) + { + session_unlocked = true; } let _ = TranslateMessage(msg); DispatchMessageW(msg); } } - false + if session_unlocked { + DrainResult::SessionUnlock + } else { + DrainResult::Continue + } } /// Wait for new input/messages or until timeout expires. @@ -57,6 +88,74 @@ fn wait_for_message_or_timeout(timeout_ms: u32) { } } +/// Create a message-only window and register for session change notifications. +/// Returns the window handle, or `None` if setup fails (non-fatal -- hooks still work, +/// they just won't auto-recover after session lock/unlock). +unsafe fn create_session_notification_window() -> Option { + let class_name: Vec = "HandyKeysSessionWatcher\0".encode_utf16().collect(); + let wnd_class = WNDCLASSW { + lpfnWndProc: Some(windows::Win32::UI::WindowsAndMessaging::DefWindowProcW), + lpszClassName: windows::core::PCWSTR(class_name.as_ptr()), + ..Default::default() + }; + RegisterClassW(&wnd_class); + + let hwnd = CreateWindowExW( + WINDOW_EX_STYLE::default(), + windows::core::PCWSTR(class_name.as_ptr()), + windows::core::PCWSTR(ptr::null()), + WINDOW_STYLE::default(), + 0, + 0, + 0, + 0, + Some(HWND_MESSAGE), + None, + None, + None, + ); + + let hwnd = match hwnd { + Ok(h) => h, + Err(_) => return None, + }; + + if WTSRegisterSessionNotification(hwnd, NOTIFY_FOR_THIS_SESSION).is_err() { + let _ = DestroyWindow(hwnd); + return None; + } + + Some(hwnd) +} + +/// Clean up the session notification window. +unsafe fn destroy_session_notification_window(hwnd: HWND) { + let _ = WTSUnRegisterSessionNotification(hwnd); + let _ = DestroyWindow(hwnd); +} + +/// Re-install low-level hooks after session unlock invalidated them. +/// Returns the new hook handles, or `None` on failure. +unsafe fn reinstall_hooks( + old_kb: windows::Win32::UI::WindowsAndMessaging::HHOOK, + old_mouse: windows::Win32::UI::WindowsAndMessaging::HHOOK, +) -> Option<( + windows::Win32::UI::WindowsAndMessaging::HHOOK, + windows::Win32::UI::WindowsAndMessaging::HHOOK, +)> { + let _ = UnhookWindowsHookEx(old_kb); + let _ = UnhookWindowsHookEx(old_mouse); + + let kb = SetWindowsHookExW(WH_KEYBOARD_LL, Some(keyboard_hook_proc), None, 0).ok()?; + match SetWindowsHookExW(WH_MOUSE_LL, Some(mouse_hook_proc), None, 0) { + Ok(mouse) => Some((kb, mouse)), + Err(_) => { + let _ = UnhookWindowsHookEx(kb); + None + } + } +} + /// Internal listener state returned to KeyboardListener pub(crate) struct WindowsListenerState { pub event_receiver: mpsc::Receiver, @@ -86,7 +185,7 @@ pub(crate) fn spawn(blocking_hotkeys: Option) -> Result h, Err(e) => { eprintln!("Failed to install keyboard hook: {:?}", e); @@ -97,11 +196,10 @@ pub(crate) fn spawn(blocking_hotkeys: Option) -> Result h, Err(e) => { eprintln!("Failed to install mouse hook: {:?}", e); - // Clean up keyboard hook before returning unsafe { let _ = UnhookWindowsHookEx(kb_hook); } @@ -109,6 +207,10 @@ pub(crate) fn spawn(blocking_hotkeys: Option) -> Result) -> Result break, + DrainResult::SessionUnlock => unsafe { + if let Some((new_kb, new_mouse)) = reinstall_hooks(kb_hook, mouse_hook) { + kb_hook = new_kb; + mouse_hook = new_mouse; + eprintln!("handy-keys: hooks re-installed after session unlock"); + } else { + eprintln!("handy-keys: failed to re-install hooks after session unlock"); + break; + } + }, + DrainResult::Continue => {} } - // Wait for messages or timeout — unlike thread::sleep, this returns + // Wait for messages or timeout -- unlike thread::sleep, this returns // immediately when a message arrives, so hook callbacks are never delayed. wait_for_message_or_timeout(HOOK_LOOP_TIMEOUT_MS); } + // Clean up session notification window + if let Some(hwnd) = session_hwnd { + unsafe { + destroy_session_notification_window(hwnd); + } + } + // Clean up the hooks unsafe { let _ = UnhookWindowsHookEx(kb_hook); @@ -349,7 +469,7 @@ mod tests { PostQuitMessage(0); } let mut msg = MSG::default(); - assert!(drain_thread_messages(&mut msg)); + assert!(matches!(drain_thread_messages(&mut msg), DrainResult::Quit)); clear_message_queue(); } }