1  // Copyright (c) Microsoft Corporation
  2  // The Microsoft Corporation licenses this file to you under the MIT license.
  3  // See the LICENSE file in the project root for more information.
  4  
  5  using System.Diagnostics.CodeAnalysis;
  6  using System.Runtime.InteropServices;
  7  using CommunityToolkit.Mvvm.Messaging;
  8  using Microsoft.CmdPal.UI.Messages;
  9  using Microsoft.CmdPal.UI.ViewModels;
 10  using Microsoft.CmdPal.UI.ViewModels.Messages;
 11  using Microsoft.UI.Xaml;
 12  using Windows.Win32;
 13  using Windows.Win32.Foundation;
 14  using Windows.Win32.UI.Shell;
 15  using Windows.Win32.UI.WindowsAndMessaging;
 16  using WinRT.Interop;
 17  using RS_ = Microsoft.CmdPal.UI.Helpers.ResourceLoaderInstance;
 18  
 19  namespace Microsoft.CmdPal.UI.Helpers;
 20  
 21  [SuppressMessage("StyleCop.CSharp.NamingRules", "SA1310:Field names should not contain underscore", Justification = "Stylistically, window messages are WM_*")]
 22  [SuppressMessage("StyleCop.CSharp.NamingRules", "SA1306:Field names should begin with lower-case letter", Justification = "Stylistically, window messages are WM_*")]
 23  internal sealed partial class TrayIconService
 24  {
 25      private const uint MY_NOTIFY_ID = 1000;
 26      private const uint WM_TRAY_ICON = PInvoke.WM_USER + 1;
 27  
 28      private readonly SettingsModel _settingsModel;
 29      private readonly uint WM_TASKBAR_RESTART;
 30  
 31      private Window? _window;
 32      private HWND _hwnd;
 33      private WNDPROC? _originalWndProc;
 34      private WNDPROC? _trayWndProc;
 35      private NOTIFYICONDATAW? _trayIconData;
 36      private DestroyIconSafeHandle? _largeIcon;
 37      private DestroyMenuSafeHandle? _popupMenu;
 38  
 39      public TrayIconService(SettingsModel settingsModel)
 40      {
 41          _settingsModel = settingsModel;
 42  
 43          // TaskbarCreated is the message that's broadcast when explorer.exe
 44          // restarts. We need to know when that happens to be able to bring our
 45          // notification area icon back
 46          WM_TASKBAR_RESTART = PInvoke.RegisterWindowMessage("TaskbarCreated");
 47      }
 48  
 49      public void SetupTrayIcon(bool? showSystemTrayIcon = null)
 50      {
 51          if (showSystemTrayIcon ?? _settingsModel.ShowSystemTrayIcon)
 52          {
 53              if (_window is null)
 54              {
 55                  _window = new Window();
 56                  _hwnd = new HWND(WindowNative.GetWindowHandle(_window));
 57  
 58                  // LOAD BEARING: If you don't stick the pointer to HotKeyPrc into a
 59                  // member (and instead like, use a local), then the pointer we marshal
 60                  // into the WindowLongPtr will be useless after we leave this function,
 61                  // and our **WindProc will explode**.
 62                  _trayWndProc = WindowProc;
 63                  var hotKeyPrcPointer = Marshal.GetFunctionPointerForDelegate(_trayWndProc);
 64                  _originalWndProc = Marshal.GetDelegateForFunctionPointer<WNDPROC>(PInvoke.SetWindowLongPtr(_hwnd, WINDOW_LONG_PTR_INDEX.GWL_WNDPROC, hotKeyPrcPointer));
 65              }
 66  
 67              if (_trayIconData is null)
 68              {
 69                  // We need to stash this handle, so it doesn't clean itself up. If
 70                  // explorer restarts, we'll come back through here, and we don't
 71                  // really need to re-load the icon in that case. We can just use
 72                  // the handle from the first time.
 73                  _largeIcon = GetAppIconHandle();
 74                  _trayIconData = new NOTIFYICONDATAW()
 75                  {
 76                      cbSize = (uint)Marshal.SizeOf<NOTIFYICONDATAW>(),
 77                      hWnd = _hwnd,
 78                      uID = MY_NOTIFY_ID,
 79                      uFlags = NOTIFY_ICON_DATA_FLAGS.NIF_MESSAGE | NOTIFY_ICON_DATA_FLAGS.NIF_ICON | NOTIFY_ICON_DATA_FLAGS.NIF_TIP,
 80                      uCallbackMessage = WM_TRAY_ICON,
 81                      hIcon = (HICON)_largeIcon.DangerousGetHandle(),
 82                      szTip = RS_.GetString("AppStoreName"),
 83                  };
 84              }
 85  
 86              var d = (NOTIFYICONDATAW)_trayIconData;
 87  
 88              // Add the notification icon
 89              PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_ADD, in d);
 90  
 91              if (_popupMenu is null)
 92              {
 93                  _popupMenu = PInvoke.CreatePopupMenu_SafeHandle();
 94                  PInvoke.InsertMenu(_popupMenu, 0, MENU_ITEM_FLAGS.MF_BYPOSITION | MENU_ITEM_FLAGS.MF_STRING, PInvoke.WM_USER + 1, RS_.GetString("TrayMenu_Settings"));
 95                  PInvoke.InsertMenu(_popupMenu, 1, MENU_ITEM_FLAGS.MF_BYPOSITION | MENU_ITEM_FLAGS.MF_STRING, PInvoke.WM_USER + 2, RS_.GetString("TrayMenu_Close"));
 96              }
 97          }
 98          else
 99          {
100              Destroy();
101          }
102      }
103  
104      public void Destroy()
105      {
106          if (_trayIconData is not null)
107          {
108              var d = (NOTIFYICONDATAW)_trayIconData;
109              if (PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_DELETE, in d))
110              {
111                  _trayIconData = null;
112              }
113          }
114  
115          if (_popupMenu is not null)
116          {
117              _popupMenu.Close();
118              _popupMenu = null;
119          }
120  
121          if (_largeIcon is not null)
122          {
123              _largeIcon.Close();
124              _largeIcon = null;
125          }
126  
127          if (_window is not null)
128          {
129              _window.Close();
130              _window = null;
131              _hwnd = HWND.Null;
132          }
133      }
134  
135      private DestroyIconSafeHandle GetAppIconHandle()
136      {
137          var exePath = Path.Combine(AppContext.BaseDirectory, "Microsoft.CmdPal.UI.exe");
138          DestroyIconSafeHandle largeIcon;
139          PInvoke.ExtractIconEx(exePath, 0, out largeIcon, out _, 1);
140          return largeIcon;
141      }
142  
143      private LRESULT WindowProc(
144          HWND hwnd,
145          uint uMsg,
146          WPARAM wParam,
147          LPARAM lParam)
148      {
149          switch (uMsg)
150          {
151              case PInvoke.WM_COMMAND:
152                  {
153                      if (wParam == PInvoke.WM_USER + 1)
154                      {
155                          WeakReferenceMessenger.Default.Send(new OpenSettingsMessage());
156                      }
157                      else if (wParam == PInvoke.WM_USER + 2)
158                      {
159                          WeakReferenceMessenger.Default.Send<QuitMessage>();
160                      }
161                  }
162  
163                  break;
164  
165              // Shell_NotifyIcon can fail when we invoke it during the time explorer.exe isn't present/ready to handle it.
166              // We'll also never receive WM_TASKBAR_RESTART message if the first call to Shell_NotifyIcon failed, so we use
167              // WM_WINDOWPOSCHANGING which is always received on explorer startup sequence.
168              case PInvoke.WM_WINDOWPOSCHANGING:
169                  {
170                      if (_trayIconData is null)
171                      {
172                          SetupTrayIcon();
173                      }
174                  }
175  
176                  break;
177              default:
178                  // WM_TASKBAR_RESTART isn't a compile-time constant, so we can't
179                  // use it in a case label
180                  if (uMsg == WM_TASKBAR_RESTART)
181                  {
182                      // Handle the case where explorer.exe restarts.
183                      // Even if we created it before, do it again
184                      SetupTrayIcon();
185                  }
186                  else if (uMsg == WM_TRAY_ICON)
187                  {
188                      switch ((uint)lParam.Value)
189                      {
190                          case PInvoke.WM_RBUTTONUP:
191                              {
192                                  if (_popupMenu is not null)
193                                  {
194                                      PInvoke.GetCursorPos(out var cursorPos);
195                                      PInvoke.SetForegroundWindow(_hwnd);
196                                      PInvoke.TrackPopupMenuEx(_popupMenu, (uint)TRACK_POPUP_MENU_FLAGS.TPM_LEFTALIGN | (uint)TRACK_POPUP_MENU_FLAGS.TPM_BOTTOMALIGN, cursorPos.X, cursorPos.Y, _hwnd, null);
197                                  }
198                              }
199  
200                              break;
201                          case PInvoke.WM_LBUTTONUP:
202                          case PInvoke.WM_LBUTTONDBLCLK:
203                              WeakReferenceMessenger.Default.Send<HotkeySummonMessage>(new(string.Empty, HWND.Null));
204                              break;
205                      }
206                  }
207  
208                  break;
209          }
210  
211          return PInvoke.CallWindowProc(_originalWndProc, hwnd, uMsg, wParam, lParam);
212      }
213  }