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

@@ -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 = []

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<KeyboardHook>| {
+    let mut toggle_enabled = |kbhook: &mut Option<KeyboardHook>| {
         if kbhook.is_some() {
             *kbhook = None;
             tray_icon.set_icon(icon_disabled.0);

src/winapi/tray_icon.rs

@@ -1,38 +1,57 @@
+use std::marker::PhantomData;
 use std::{mem, ptr};
 
 use windows_sys::Win32::Foundation::*;
 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, &notification_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, &notification_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, &notification_data) };
 }
 
-fn notification_data(hwnd: HWND) -> NOTIFYICONDATAA {
-    unsafe {
-        let mut notification_data: NOTIFYICONDATAA = mem::zeroed();
-        notification_data.cbSize = mem::size_of_val(&notification_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, &notification_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(&notification_data) as _;
+    notification_data.hWnd = hwnd;
+    notification_data
 }