From 80701815c7fb6cbb2070e6a9e63e87e6abfe1054 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20Kr=C3=B6ger?= Date: Wed, 12 Jun 2024 17:17:06 +0200 Subject: [PATCH] Properly handle hook reentrancy --- src/winapi/keyboard.rs | 133 ++++++++++++++++++++++++----------------- 1 file changed, 79 insertions(+), 54 deletions(-) diff --git a/src/winapi/keyboard.rs b/src/winapi/keyboard.rs index 3a4480e..47bffae 100644 --- a/src/winapi/keyboard.rs +++ b/src/winapi/keyboard.rs @@ -1,6 +1,6 @@ //! Safe abstraction over the low-level windows keyboard hook API. -use std::cell::Cell; +use std::cell::{Cell, RefCell}; use std::fmt::Display; use std::{mem, ptr}; @@ -16,6 +16,7 @@ const HOOK_EXECUTING: usize = 1; thread_local! { /// Stores a type erased pointer to the hook closure. static HOOK: Cell = const { Cell::new(HOOK_INVALID) }; + static QUEUED_INPUTS: RefCell> = const { RefCell::new(Vec::new()) }; } /// Wrapper for the low-level keyboard hook API. @@ -139,32 +140,36 @@ where let key_event = KeyEvent::from_hook_lparam(hook_lparam); let injected = hook_lparam.flags & LLKHF_INJECTED != 0; - // `SendInput()` internally calls the hook function. Filter out injected events - // to prevent recursion and potential stack overflows if our remapping logic - // sent the injected event. + // `SendInput()` internally calls the hook function. Filter out injected + // events to prevent recursion and potential stack overflows if our + // remapping logic has sent the injected event. if injected { return CallNextHookEx(0, code, wparam, lparam); } - let mut handled = false; - - // As long as we prevent recursion by dropping injected events, windows - // should not be calling the hook again while it is already executing. - // That means that the pointer in TLS storage should always be valid. + // There are two conditions for which the hook can be re-entered: + // 1. `SendInput()` creates new injected input events. + // We queue input events and send them after the hook closure has + // returned to prevent a second mutable borrow to the closure. + // 2. `CallNextHookEx()` when more than one unhandled input event is queued. + // Not exposed to the user and such must not be called within the closure. + // + // How to trigger 2.: + // The classic CMD window has a "Quick Edit Mode" option which is enabled + // by default. Windows stops to read from stdout and stderr when the user + // selects characters in the CMD window. + // Any write to stdout (e.g. a call to `println!()`) blocks while + // "Quick Edit Mode" is active. The key event which exits the "Quick Edit + // Mode" triggers the hook a second time. + + // Replace the pointer to the closure with a marker, so that `send_key()` + // can detect if it was called from within the hook. let hook_ptr = HOOK.replace(HOOK_EXECUTING) as *mut F; - if let Some(hook) = unsafe { hook_ptr.as_mut() } { - handled = hook(key_event); - HOOK.set(hook_ptr as usize); - } else { - // There is one special case with classical CMD windows: - // The "Quick Edit Mode" option which is enabled by default. - // Windows stops to read from stdout and stderr when the user - // selects characters in the cmd window. - // Any write to stdout (e.g. a call to `println!()`) blocks while - // "Quick Edit Mode" is active. The nasty part is that the key event - // which exits the "Quick Edit Mode" triggers the hook a second time - // before the blocked write to stdout can return. - } + let hook = unsafe { hook_ptr.as_mut().unwrap() }; + let handled = hook(key_event); + HOOK.set(hook_ptr as usize); + + send_queued_inputs(); if handled { -1 @@ -175,49 +180,69 @@ where /// Sends a virtual key event. pub fn send_key(key: KeyEvent) { - unsafe { - let mut inputs: [INPUT; 2] = mem::zeroed(); - - let n_inputs = match key.key { + QUEUED_INPUTS.with_borrow_mut(|queued_inputs| { + match key.key { KeyType::VirtualKey(vk) => { - inputs[0].r#type = INPUT_KEYBOARD; - inputs[0].Anonymous.ki = KEYBDINPUT { - wVk: vk.into(), - wScan: key.scan_code, - dwFlags: if key.up { KEYEVENTF_KEYUP } else { 0 }, - time: key.time, - dwExtraInfo: 0, - }; - 1 + queued_inputs.push(INPUT { + r#type: INPUT_KEYBOARD, + Anonymous: INPUT_0 { + ki: KEYBDINPUT { + wVk: vk.into(), + wScan: key.scan_code, + dwFlags: if key.up { KEYEVENTF_KEYUP } else { 0 }, + time: key.time, + dwExtraInfo: 0, + }, + }, + }); } KeyType::Unicode(c) => { // Sends a unicode character, knows as `VK_PACKET`. // Interestingly this is faster than sending a regular virtual key event. - inputs - .iter_mut() - .zip(c.to_utf16()) - .map(|(input, c)| { - input.r#type = INPUT_KEYBOARD; - input.Anonymous.ki = KEYBDINPUT { - wVk: 0, - wScan: c, - dwFlags: KEYEVENTF_UNICODE | if key.up { KEYEVENTF_KEYUP } else { 0 }, - time: key.time, - dwExtraInfo: 0, - }; - }) - .count() + for c in c.to_utf16() { + queued_inputs.push(INPUT { + r#type: INPUT_KEYBOARD, + Anonymous: INPUT_0 { + ki: KEYBDINPUT { + wVk: 0, + wScan: c, + dwFlags: KEYEVENTF_UNICODE | if key.up { KEYEVENTF_KEYUP } else { 0 }, + time: key.time, + dwExtraInfo: 0, + }, + }, + }); + } } }; + }); - SendInput( - n_inputs as _, - inputs.as_mut_ptr(), - mem::size_of::() as _, - ); + if HOOK.get() != HOOK_EXECUTING { + // Send inputs only when not called from within the hook process. + send_queued_inputs(); + } else { + // Events will be sent when leaving the hook to prevent re-entrancy edge cases. } } +fn send_queued_inputs() { + let mut queued_inputs = QUEUED_INPUTS.with_borrow_mut(|qi| std::mem::replace(qi, Vec::new())); + + if !queued_inputs.is_empty() { + unsafe { + SendInput( + queued_inputs.len() as u32, + queued_inputs.as_ptr(), + mem::size_of::() as _, + ) + }; + queued_inputs.clear(); + } + + // Re-use the previous allocation + QUEUED_INPUTS.with_borrow_mut(move |qi| std::mem::replace(qi, queued_inputs)); +} + /// Returns a virtual key code if the requested character can be typed with a /// single key press/release. pub fn get_virtual_key(c: char) -> Option {