Fix unsoundness in tray icon code

Userdata pointer must be set in the wndproc callback during window creation.
Use `Rc` and `Cell` to remove prevent creating multiple mutable references to the state and makes `TrayIcon` !Sync

src/winapi/tray_icon.rs

@@ -1,38 +1,67 @@
+use std::cell::Cell;
+use std::rc::Rc;
 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,
-    icon: HICON,
+const MSG_ID_TRAY_ICON: u32 = WM_USER;
+
+struct State {
+    msg_id: u32,
+    icon: Cell<HICON>,
+}
+
+impl State {
+    fn raw_from_hwnd(hwnd: HWND) -> *mut State {
+        let ptr = unsafe { GetWindowLongPtrA(hwnd, GWLP_USERDATA) as *mut State };
+        assert!(!ptr.is_null());
+        ptr
+    }
+
+    fn from_hwnd(hwnd: HWND) -> Rc<State> {
+        let state = Self::raw_from_hwnd(hwnd);
+        unsafe {
+            Rc::increment_strong_count(state);
+            Rc::from_raw(state)
+        }
+    }
 }
 
 pub struct TrayIcon {
     hwnd: HWND,
+    state: Rc<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);
+            // The window carried a reference to the state in its user data.
+            // It has now been destroyed, manually adjust the reference count.
+            Rc::decrement_strong_count(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 = Rc::new(State {
+            msg_id,
+            icon: Cell::new(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 +75,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,23 +87,19 @@ impl TrayIcon {
                 0,
                 0,
                 0,
-                ptr::null(),
-            );
-            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 _);
+                Rc::into_raw(state.clone()) as *const _,
+            )
+        };
+        assert_ne!(hwnd, 0);
 
-            // Create the tray icon
-            add_tray_icon(hwnd);
+        // Create the tray icon
+        add_tray_icon(hwnd, icon);
 
-            Self { hwnd }
-        }
+        Self { hwnd, state }
     }
 
     pub fn set_icon(&self, icon: HICON) {
-        userdata(self.hwnd).icon = icon;
+        self.state.icon.set(icon);
         let mut notification_data = notification_data(self.hwnd);
         notification_data.uFlags = NIF_ICON;
         notification_data.hIcon = icon;
@@ -85,40 +110,45 @@ impl TrayIcon {
 }
 
 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),
+    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);
+    }
+
+    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.get());
     }
-    0
+
+    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 userdata(hwnd: HWND) -> &'static mut UserData {
-    unsafe {
-        let userdata = GetWindowLongPtrA(hwnd, GWLP_USERDATA) as *mut UserData;
-        &mut *userdata
-    }
+    let mut notification_data: NOTIFYICONDATAA = unsafe { mem::zeroed() };
+    notification_data.cbSize = mem::size_of_val(&notification_data) as _;
+    notification_data.hWnd = hwnd;
+    notification_data
 }