From f8e53ebaab1ea1ec948aed15d32991dca46a15ae Mon Sep 17 00:00:00 2001
From: Frank Richter <frank.richter@gmail.com>
Date: Sun, 12 May 2024 18:42:35 +0200
Subject: [PATCH] Handle TaskbarCreated message.

Fixes icon disappearing upon Explorer restart (#9) and
more gracefully handles "early" starts without an available
task bar.
---
 HDRTray.cpp    | 20 +++++++++++++++++++-
 NotifyIcon.cpp |  7 +++++++
 NotifyIcon.hpp |  2 ++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/HDRTray.cpp b/HDRTray.cpp
index eea1adb..e603627 100644
--- a/HDRTray.cpp
+++ b/HDRTray.cpp
@@ -161,6 +161,7 @@ BOOL InitInstance(HINSTANCE hInstance, int nCmdShow)
 }
 
 static std::unique_ptr<NotifyIcon> notify_icon;
+static UINT msg_TaskbarCreated;
 
 //
 //  FUNCTION: WndProc(HWND, UINT, WPARAM, LPARAM)
@@ -177,9 +178,13 @@ LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
     switch (message)
     {
     case WM_CREATE:
+        msg_TaskbarCreated = RegisterWindowMessage(L"TaskbarCreated");
         notify_icon.reset(new NotifyIcon(hWnd));
         if (!notify_icon->Add())
-            return -1;
+        {
+            // Set up a timer, this is the amount of time we wait for TaskbarCreated
+            SetTimer(hWnd, 1, 30000, nullptr);
+        }
         break;
     case WM_COMMAND:
         {
@@ -216,7 +221,20 @@ LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
         break;
     case NotifyIcon::MESSAGE:
         return notify_icon->HandleMessage(hWnd, wParam, lParam);
+    case WM_TIMER:
+        KillTimer(hWnd, 1);
+        // No TaskbarCreated was received, exit
+        if (!notify_icon->WasAdded())
+            DestroyWindow(hWnd);
+        break;
     default:
+        if (message == msg_TaskbarCreated)
+        {
+            // Taskbar was created (Explorer restart, DPI change), so re-create notify icon
+            notify_icon->Remove();
+            if (!notify_icon->Add())
+                DestroyWindow(hWnd);
+        }
         return DefWindowProc(hWnd, message, wParam, lParam);
     }
     return 0;
diff --git a/NotifyIcon.cpp b/NotifyIcon.cpp
index afe33d8..64fce36 100644
--- a/NotifyIcon.cpp
+++ b/NotifyIcon.cpp
@@ -89,6 +89,11 @@ NotifyIcon::~NotifyIcon()
     DestroyMenu(popup_menu);
 }
 
+bool NotifyIcon::WasAdded() const
+{
+    return added;
+}
+
 bool NotifyIcon::Add()
 {
     FetchHDRStatus();
@@ -106,6 +111,7 @@ bool NotifyIcon::Add()
     wrap_Shell_NotifyIconW(NIM_SETVERSION, &notify_setversion);
 
     UpdateIcon();
+    added = true;
     return true;
 }
 
@@ -113,6 +119,7 @@ void NotifyIcon::Remove()
 {
     auto notify_delete = notify_template;
     wrap_Shell_NotifyIconW(NIM_DELETE, &notify_delete);
+    added = false;
 }
 
 void NotifyIcon::UpdateHDRStatus()
diff --git a/NotifyIcon.hpp b/NotifyIcon.hpp
index a44bd78..17e57c7 100644
--- a/NotifyIcon.hpp
+++ b/NotifyIcon.hpp
@@ -26,6 +26,7 @@
 
 class NotifyIcon
 {
+    bool added = false;
     NOTIFYICONDATAW notify_template;
 
     enum { iconsetDarkMode = 0, iconsetLightMode, numIconsets };
@@ -44,6 +45,7 @@ class NotifyIcon
     NotifyIcon(HWND hwnd);
     ~NotifyIcon();
 
+    bool WasAdded() const;
     bool Add();
     void Remove();