From 0feab6db84b6209b935d8244b6fc94082bfa28e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20Kr=C3=B6ger?= Date: Tue, 14 May 2024 15:09:06 +0200 Subject: [PATCH] Fix unsoundness in tray icon code * Add PhantomData to TrayIcon to make it !Sync and !Send * Set userdata state pointer in the wndproc callback during window creation. --- config.toml | 32 +--------- src/main.rs | 4 +- src/winapi/tray_icon.rs | 132 ++++++++++++++++++++++++---------------- 3 files changed, 84 insertions(+), 84 deletions(-) diff --git a/config.toml b/config.toml index 591d290..f0c370e 100644 --- a/config.toml +++ b/config.toml @@ -24,37 +24,7 @@ caps_lock_layer = "shift" # a key press to a [virtual key code](https://docs.microsoft.com/en-us/windows/win32/inputdev/virtual-key-codes). # # Use the debug output to observe scan code (sc) and virtual key code (vk) information while pressing keys. -base = [ - # The `characters` target remaps multiple keys to (Unicode) characters starting - # at `scan_code`. Increments `scan_code` by one for each following character. - # Simulates a virtual key press if the character exists on the active Windows layout - # so that keyborad shortcuts (e.g. Ctrl+a) continue to work as expected. - # Injects a Unicode symbol if the character is not available on the Windows layout. - { scan_code = 0x02, characters = "1234567890[]" }, - { scan_code = 0x10, characters = "',.pyfgcrl/=\\" }, - { scan_code = 0x1E, characters = "aoeuidhtns-" }, - { scan_code = 0x2C, characters = ";qjkxbmwvz" }, - - # The `layer` target switches to another virtual layer. - # Ignores the original scan code if no `virtual_key` is specified. - { scan_code = 0x2A, layer = "shift", virtual_key = 0xA0 }, # left shift - { scan_code = 0xE036, layer = "shift", virtual_key = 0xA1 }, # right shift - { scan_code = 0x3A, layer_lock = "shift" }, # caps lock - - # The `virtual_keys` works similar to the `characters` target but takes a collection of virtual keys. - # An empty target can be used to ignore scan codes. - { scan_code = 0x021D, virtual_keys = [] }, # ignored additional scan code from `Alt Gr` - - # Emoji Keys! - # { scan_code = 0x29, characters = "🚀" }, # `~` on US layout, `^` on german layout -] - -shift = [ - { scan_code = 0x02, characters = "!@#$%^&*(){}" }, - { scan_code = 0x10, characters = "\"<>PYFGCRL?+|" }, - { scan_code = 0x1E, characters = "AOEWIDHTNS_" }, - { scan_code = 0x2C, characters = ":QJKXBMWVZ" }, -] +base = [] # A dummy layer not referenced by any layer key action. unreachable = [] diff --git a/src/main.rs b/src/main.rs index 3f3ae18..daf4005 100644 --- a/src/main.rs +++ b/src/main.rs @@ -130,7 +130,7 @@ fn main() -> Result<()> { // Arbitrary ID in the WM_APP range, used to identify which tray icon a message originates from. const WM_APP_TRAYICON: u32 = WM_APP + 873; - let tray_icon = TrayIcon::new(WM_APP_TRAYICON, icon_enabled.0); + let mut tray_icon = TrayIcon::new(WM_APP_TRAYICON, icon_enabled.0); let cmd = env::current_exe().unwrap(); let autostart = AutoStartEntry::new( @@ -139,7 +139,7 @@ fn main() -> Result<()> { ); // Enabled state can be changed by double click to the tray icon or from the context menu. - let toggle_enabled = |kbhook: &mut Option| { + let mut toggle_enabled = |kbhook: &mut Option| { if kbhook.is_some() { *kbhook = None; tray_icon.set_icon(icon_disabled.0); diff --git a/src/winapi/tray_icon.rs b/src/winapi/tray_icon.rs index e2a89f2..f8f0371 100644 --- a/src/winapi/tray_icon.rs +++ b/src/winapi/tray_icon.rs @@ -1,3 +1,4 @@ +use std::marker::PhantomData; use std::{mem, ptr}; use windows_sys::Win32::Foundation::*; @@ -5,34 +6,52 @@ use windows_sys::Win32::System::LibraryLoader::*; use windows_sys::Win32::UI::Shell::*; use windows_sys::Win32::UI::WindowsAndMessaging::*; -struct UserData { - msg_tray: u32, +const MSG_ID_TRAY_ICON: u32 = WM_USER; + +struct State { + msg_id: u32, icon: HICON, } +impl State { + fn raw_from_hwnd(hwnd: HWND) -> *mut Self { + let ptr = unsafe { GetWindowLongPtrA(hwnd, GWLP_USERDATA) as *mut Self }; + assert!(!ptr.is_null()); + ptr + } + + fn from_hwnd(hwnd: HWND) -> &'static mut Self { + unsafe { &mut *Self::raw_from_hwnd(hwnd) } + } +} + pub struct TrayIcon { hwnd: HWND, + _state: PhantomData<*mut State>, } impl Drop for TrayIcon { fn drop(&mut self) { + let state_ptr = State::raw_from_hwnd(self.hwnd); unsafe { - drop(Box::from_raw(userdata(self.hwnd))); Shell_NotifyIconA(NIM_DELETE, ¬ification_data(self.hwnd)); DestroyWindow(self.hwnd); + drop(Box::from_raw(state_ptr)); } } } impl TrayIcon { - pub fn new(msg_tray: u32, icon: HICON) -> Self { + pub fn new(msg_id: u32, icon: HICON) -> Self { assert!( - (WM_APP..WM_APP + 0x4000).contains(&msg_tray), + (WM_APP..WM_APP + 0x4000).contains(&msg_id), "message must be in the WM_APP range" ); + let state = Box::new(State { msg_id, icon }); + let class_name = "trayicon\0".as_ptr(); - unsafe { + let hwnd = unsafe { let hinstance = GetModuleHandleA(ptr::null()); // A class is unique and the `RegisterClass()` function fails when @@ -46,7 +65,7 @@ impl TrayIcon { // Never visible, only required to receive messages. // It cannot be a message-only `HWND_MESSAGE` thought because those do not // receive global messages like "TaskbarCreated". - let hwnd = CreateWindowExA( + CreateWindowExA( 0, class_name, ptr::null(), @@ -58,67 +77,78 @@ impl TrayIcon { 0, 0, 0, - ptr::null(), - ); - assert_ne!(hwnd, 0); + Box::into_raw(state) as *mut _, + ) + }; + assert_ne!(hwnd, 0); - // Set message as associated data - let userdata = Box::new(UserData { msg_tray, icon }); - SetWindowLongPtrA(hwnd, GWLP_USERDATA, Box::into_raw(userdata) as _); + // Create the tray icon + add_tray_icon(hwnd, icon); - // Create the tray icon - add_tray_icon(hwnd); - - Self { hwnd } + Self { + hwnd, + _state: PhantomData, } } - pub fn set_icon(&self, icon: HICON) { - userdata(self.hwnd).icon = icon; - let mut notification_data = notification_data(self.hwnd); - notification_data.uFlags = NIF_ICON; - notification_data.hIcon = icon; - unsafe { - Shell_NotifyIconA(NIM_MODIFY, ¬ification_data); - } + pub fn set_icon(&mut self, icon: HICON) { + update_tray_icon(self.hwnd, icon); + State::from_hwnd(self.hwnd).icon = icon; } } unsafe extern "system" fn wndproc(hwnd: HWND, msg: u32, wparam: WPARAM, lparam: LPARAM) -> LRESULT { - match msg { - WM_USER => { - // Forward the message to the main event loop. - PostMessageA(hwnd, userdata(hwnd).msg_tray, wparam, lparam); - } - msg if msg == RegisterWindowMessageA("TaskbarCreated\0".as_ptr()) => { - // Re-add the tray icon if explorer.exe has restarted. - add_tray_icon(hwnd); - } - msg => return DefWindowProcA(hwnd, msg, wparam, lparam), + println!( + "WND hwnd={:p} msg={:06}, wparam={:06}, lparam={:06}", + hwnd as *const (), msg, wparam, lparam + ); + + if msg == WM_NCCREATE { + // Attach user data to the window so it can be accessed from this + // callback function when receiving other messages. + // This must be done here because the WM_NCCREATE (which is the very + // first message of each window) and other message are dispatched to + // this callback before `CreateWindowEx()` returns. + // https://devblogs.microsoft.com/oldnewthing/20191014-00/?p=102992 + let create_params = lparam as *const CREATESTRUCTA; + SetWindowLongPtrA( + hwnd, + GWLP_USERDATA, + (*create_params).lpCreateParams as isize, + ); + return DefWindowProcA(hwnd, msg, wparam, lparam); } - 0 + + let state = State::from_hwnd(hwnd); + if msg == MSG_ID_TRAY_ICON { + // Forward the message to the main event loop. + PostMessageA(hwnd, state.msg_id, wparam, lparam); + } else if msg == RegisterWindowMessageA("TaskbarCreated\0".as_ptr()) { + // Re-add the tray icon if explorer.exe has restarted. + add_tray_icon(hwnd, state.icon); + } + + DefWindowProcA(hwnd, msg, wparam, lparam) } -fn add_tray_icon(hwnd: HWND) { +fn add_tray_icon(hwnd: HWND, icon: HICON) { let mut notification_data = notification_data(hwnd); notification_data.uFlags = NIF_MESSAGE | NIF_ICON; - notification_data.uCallbackMessage = WM_USER; - notification_data.hIcon = userdata(hwnd).icon; + notification_data.uCallbackMessage = MSG_ID_TRAY_ICON; + notification_data.hIcon = icon; unsafe { Shell_NotifyIconA(NIM_ADD, ¬ification_data) }; } -fn notification_data(hwnd: HWND) -> NOTIFYICONDATAA { - unsafe { - let mut notification_data: NOTIFYICONDATAA = mem::zeroed(); - notification_data.cbSize = mem::size_of_val(¬ification_data) as _; - notification_data.hWnd = hwnd; - notification_data - } +fn update_tray_icon(hwnd: HWND, icon: HICON) { + let mut notification_data = notification_data(hwnd); + notification_data.uFlags = NIF_ICON; + notification_data.hIcon = icon; + unsafe { Shell_NotifyIconA(NIM_MODIFY, ¬ification_data) }; } -fn userdata(hwnd: HWND) -> &'static mut UserData { - unsafe { - let userdata = GetWindowLongPtrA(hwnd, GWLP_USERDATA) as *mut UserData; - &mut *userdata - } +fn notification_data(hwnd: HWND) -> NOTIFYICONDATAA { + let mut notification_data: NOTIFYICONDATAA = unsafe { mem::zeroed() }; + notification_data.cbSize = mem::size_of_val(¬ification_data) as _; + notification_data.hWnd = hwnd; + notification_data }