diff --git a/src/winapi/tray_icon.rs b/src/winapi/tray_icon.rs index e2a89f2..86448d2 100644 --- a/src/winapi/tray_icon.rs +++ b/src/winapi/tray_icon.rs @@ -1,3 +1,5 @@ +use std::cell::Cell; +use std::rc::Rc; use std::{mem, ptr}; use windows_sys::Win32::Foundation::*; @@ -5,34 +7,61 @@ 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, +} + +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 { + let state = Self::raw_from_hwnd(hwnd); + unsafe { + Rc::increment_strong_count(state); + Rc::from_raw(state) + } + } } pub struct TrayIcon { hwnd: HWND, + state: Rc, } 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); + // 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, ¬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 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(¬ification_data) as _; + notification_data.hWnd = hwnd; + notification_data }