Refactor tray related code

3 files changed, 96 insertions, 55 deletions

diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index b5743b7c76..a3cd2dd5f8 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -14,7 +14,6 @@ import {
   Tray,
 } from 'electron';
 import * as path from 'path';
-import { sprintf } from 'sprintf-js';
 import util from 'util';
 import config from '../config.json';
 import { closeToExpiry, hasExpired } from '../shared/account-expiry';
@@ -513,11 +512,17 @@ class ApplicationMain {
     this.tunnelStateExpectation = new Expectation(async () => {
       this.trayIconController = new TrayIconController(
         tray,
+        windowController,
         this.trayIconType(this.tunnelState, this.settings.blockWhenDisconnected),
         this.guiSettings.monochromaticIcon,
+        () => setImmediate(() => void this.connectTunnel()),
+        () => setImmediate(() => void this.reconnectTunnel()),
+        () => setImmediate(() => void this.disconnectTunnel()),
       );
       await this.trayIconController.updateTheme();
 
+      this.setTrayContextMenu();
+
       if (process.platform === 'win32') {
         nativeTheme.on('updated', async () => {
           if (this.guiSettings.monochromaticIcon) {
@@ -574,7 +579,7 @@ class ApplicationMain {
           this.setMacOsAppMenu();
           break;
         case 'linux':
-          this.tray.setContextMenu(this.createTrayContextMenu());
+          this.setTrayContextMenu();
           this.setLinuxAppMenu();
           this.windowController.window.setMenuBarVisibility(false);
           break;
@@ -716,6 +721,7 @@ class ApplicationMain {
 
       // update the tray icon to indicate that the computer is not secure anymore
       this.updateTrayIcon({ state: 'disconnected' }, false);
+      this.setTrayContextMenu();
 
       // notify renderer process
       if (this.windowController) {
@@ -886,9 +892,7 @@ class ApplicationMain {
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
 
-    if (process.platform === 'linux') {
-      this.tray?.setContextMenu(this.createTrayContextMenu());
-    }
+    this.setTrayContextMenu();
 
     this.notificationController.notifyTunnelState(
       newState,
@@ -1563,6 +1567,8 @@ class ApplicationMain {
     } else if (oldAccount && newAccount && oldAccount !== newAccount) {
       this.accountDataCache.fetch(newAccount);
     }
+
+    this.setTrayContextMenu();
   }
 
   private updateAccountData() {
@@ -1685,6 +1691,8 @@ class ApplicationMain {
       messages: messagesTranslations,
       relayLocations: relayLocationsTranslations,
     };
+
+    this.setTrayContextMenu();
   }
 
   private blockPermissionRequests() {
@@ -1882,54 +1890,6 @@ class ApplicationMain {
     Menu.setApplicationMenu(Menu.buildFromTemplate(template));
   }
 
-  private createTrayContextMenu() {
-    const template: Electron.MenuItemConstructorOptions[] = [
-      {
-        label: sprintf(messages.pgettext('tray-icon-context-menu', 'Open %(mullvadVpn)s'), {
-          mullvadVpn: 'Mullvad VPN',
-        }),
-        click: () => this.windowController?.show(),
-      },
-      { type: 'separator' },
-      {
-        label: this.getContextMenuActionButtonLabel(),
-        click: () => {
-          if (this.tunnelState.state === 'disconnected') {
-            // Workaround: gRPC calls are sometimes delayed by a few seconds and setImmediate
-            // mitigates this. https://github.com/grpc/grpc-node/issues/882
-            setImmediate(() => void this.daemonRpc.connectTunnel());
-          } else {
-            setImmediate(() => void this.daemonRpc.disconnectTunnel());
-          }
-        },
-      },
-      {
-        label: messages.gettext('Reconnect'),
-        enabled: this.tunnelState.state === 'connected' || this.tunnelState.state === 'connecting',
-        click: () => setImmediate(() => void this.daemonRpc.reconnectTunnel()),
-      },
-    ];
-
-    return Menu.buildFromTemplate(template);
-  }
-
-  private getContextMenuActionButtonLabel() {
-    switch (this.tunnelState.state) {
-      case 'disconnected':
-        return messages.gettext('Connect');
-      case 'connecting':
-        return messages.gettext('Cancel');
-      case 'connected':
-        return messages.gettext('Disconnect');
-      case 'disconnecting':
-        return '';
-      case 'error':
-        return this.tunnelState.details.blockFailure
-          ? messages.gettext('Dismiss')
-          : messages.gettext('Cancel');
-    }
-  }
-
   private addContextMenu(windowController: WindowController) {
     const menuTemplate: Electron.MenuItemConstructorOptions[] = [
       { role: 'cut' },
@@ -1991,7 +1951,11 @@ class ApplicationMain {
           // This needs to be executed on click since if it is added to the tray icon it will be
           // displayed on left click as well.
           this.tray?.on('right-click', () =>
-            this.tray?.popUpContextMenu(this.createTrayContextMenu()),
+            this.trayIconController?.popUpContextMenu(
+              this.connectedToDaemon,
+              this.settings.accountToken,
+              this.tunnelState,
+            ),
           );
           this.tray?.on('click', () => this.windowController?.show());
         } else {
@@ -2024,6 +1988,14 @@ class ApplicationMain {
     }
   }
 
+  private setTrayContextMenu() {
+    this.trayIconController?.setContextMenu(
+      this.connectedToDaemon,
+      this.settings.accountToken,
+      this.tunnelState,
+    );
+  }
+
   private installWindowsMenubarAppWindowHandlers(tray: Tray, windowController: WindowController) {
     if (!this.guiSettings.unpinnedWindow) {
       windowController.window?.on('blur', () => {
diff --git a/gui/src/main/tray-icon-controller.ts b/gui/src/main/tray-icon-controller.ts
index a16c2ad9fc..9890eba8ff 100644
--- a/gui/src/main/tray-icon-controller.ts
+++ b/gui/src/main/tray-icon-controller.ts
@@ -1,9 +1,14 @@
 import { exec as execAsync } from 'child_process';
-import { nativeImage, NativeImage, Tray } from 'electron';
+import { Menu, nativeImage, NativeImage, Tray } from 'electron';
 import path from 'path';
+import { sprintf } from 'sprintf-js';
 import { promisify } from 'util';
+import { connectEnabled, disconnectEnabled, reconnectEnabled } from '../shared/connect-helper';
+import { AccountToken, TunnelState } from '../shared/daemon-rpc-types';
+import { messages } from '../shared/gettext';
 import log from '../shared/logging';
 import KeyframeAnimation from './keyframe-animation';
+import WindowController from './window-controller';
 
 const exec = promisify(execAsync);
 
@@ -23,8 +28,12 @@ export default class TrayIconController {
 
   constructor(
     private tray: Tray,
+    private windowController: WindowController,
     private iconTypeValue: TrayIconType,
     private useMonochromaticIconValue: boolean,
+    private connect: () => void,
+    private reconnect: () => void,
+    private disconnect: () => void,
   ) {
     this.loadImages();
   }
@@ -88,6 +97,28 @@ export default class TrayIconController {
     animation.play({ end: frame });
   }
 
+  public setContextMenu(
+    connectedToDaemon: boolean,
+    accountToken: AccountToken | undefined,
+    tunnelState: TunnelState,
+  ) {
+    if (process.platform === 'linux') {
+      this.tray.setContextMenu(
+        this.createContextMenu(connectedToDaemon, accountToken, tunnelState),
+      );
+    }
+  }
+
+  public popUpContextMenu(
+    connectedToDaemon: boolean,
+    accountToken: AccountToken | undefined,
+    tunnelState: TunnelState,
+  ) {
+    this.tray.popUpContextMenu(
+      this.createContextMenu(connectedToDaemon, accountToken, tunnelState),
+    );
+  }
+
   private initAnimation() {
     const initialFrame = this.targetFrame();
     const animation = new KeyframeAnimation();
@@ -178,4 +209,40 @@ export default class TrayIconController {
         return 8;
     }
   }
+
+  private createContextMenu(
+    connectedToDaemon: boolean,
+    accountToken: AccountToken | undefined,
+    tunnelState: TunnelState,
+  ) {
+    const template: Electron.MenuItemConstructorOptions[] = [
+      {
+        label: sprintf(messages.pgettext('tray-icon-context-menu', 'Open %(mullvadVpn)s'), {
+          mullvadVpn: 'Mullvad VPN',
+        }),
+        click: () => this.windowController.show(),
+      },
+      { type: 'separator' },
+      {
+        id: 'connect',
+        label: messages.gettext('Connect'),
+        enabled: connectEnabled(connectedToDaemon, accountToken, tunnelState.state),
+        click: this.connect,
+      },
+      {
+        id: 'reconnect',
+        label: messages.gettext('Reconnect'),
+        enabled: reconnectEnabled(connectedToDaemon, accountToken, tunnelState.state),
+        click: this.reconnect,
+      },
+      {
+        id: 'disconnect',
+        label: messages.gettext('Disconnect'),
+        enabled: disconnectEnabled(connectedToDaemon, tunnelState.state),
+        click: this.disconnect,
+      },
+    ];
+
+    return Menu.buildFromTemplate(template);
+  }
 }
diff --git a/gui/src/shared/connect-helper.ts b/gui/src/shared/connect-helper.ts
index 214d22a59a..b1421cd8ef 100644
--- a/gui/src/shared/connect-helper.ts
+++ b/gui/src/shared/connect-helper.ts
@@ -8,6 +8,7 @@ export function connectEnabled(
   return (
     connectedToDaemon &&
     accountToken !== undefined &&
+    accountToken !== '' &&
     (tunnelState === 'disconnected' || tunnelState === 'disconnecting' || tunnelState === 'error')
   );
 }
@@ -20,6 +21,7 @@ export function reconnectEnabled(
   return (
     connectedToDaemon &&
     accountToken !== undefined &&
+    accountToken !== '' &&
     (tunnelState === 'connected' || tunnelState === 'connecting')
   );
 }