Update to new device state rpc calls

3 files changed, 92 insertions, 102 deletions

diff --git a/gui/src/main/daemon-rpc.ts b/gui/src/main/daemon-rpc.ts
index 84e832cd5f..c40e40b864 100644
--- a/gui/src/main/daemon-rpc.ts
+++ b/gui/src/main/daemon-rpc.ts
@@ -15,14 +15,14 @@ import {
   ConnectionConfig,
   Constraint,
   DaemonEvent,
+  DeviceEvent,
+  DeviceState,
   ErrorStateCause,
   FirewallPolicyError,
   IAccountData,
   IAppVersionInfo,
   IBridgeConstraints,
   IDevice,
-  IDeviceConfig,
-  IDeviceEvent,
   IDeviceRemoval,
   IDnsOptions,
   IErrorState,
@@ -41,6 +41,8 @@ import {
   ITunnelStateRelayInfo,
   IWireguardConstraints,
   IWireguardTunnelData,
+  LoggedInDeviceState,
+  LoggedOutDeviceState,
   ObfuscationType,
   ProxySettings,
   ProxyType,
@@ -513,18 +515,9 @@ export class DaemonRpc {
     return response.getValue();
   }
 
-  public async getDevice(): Promise<IDeviceConfig | undefined> {
-    try {
-      const response = await this.callEmpty<grpcTypes.DeviceConfig>(this.client.getDevice);
-      return convertFromDeviceConfig(response);
-    } catch (e) {
-      const error = e as grpc.ServiceError;
-      if (error.code === grpc.status.NOT_FOUND) {
-        return undefined;
-      } else {
-        throw error;
-      }
-    }
+  public async getDevice(): Promise<DeviceState> {
+    const response = await this.callEmpty<grpcTypes.DeviceState>(this.client.getDevice);
+    return convertFromDeviceState(response);
   }
 
   public async updateDevice(): Promise<void> {
@@ -1409,21 +1402,40 @@ function convertToTransportProtocol(protocol: RelayProtocol): grpcTypes.Transpor
   }
 }
 
-function convertFromDeviceEvent(deviceEvent: grpcTypes.DeviceEvent): IDeviceEvent {
-  return {
-    deviceConfig: convertFromDeviceConfig(deviceEvent.getDevice()),
-    remote: deviceEvent.getRemote(),
-  };
+function convertFromDeviceEvent(deviceEvent: grpcTypes.DeviceEvent): DeviceEvent {
+  const deviceState = convertFromDeviceState(deviceEvent.getNewState()!);
+  switch (deviceEvent.getCause()) {
+    case grpcTypes.DeviceEvent.Cause.LOGGED_IN:
+      return { type: 'logged in', deviceState: deviceState as LoggedInDeviceState };
+    case grpcTypes.DeviceEvent.Cause.LOGGED_OUT:
+      return { type: 'logged out', deviceState: deviceState as LoggedOutDeviceState };
+    case grpcTypes.DeviceEvent.Cause.REVOKED:
+      return { type: 'revoked', deviceState: deviceState as LoggedOutDeviceState };
+    case grpcTypes.DeviceEvent.Cause.UPDATED:
+      return { type: 'updated', deviceState: deviceState as LoggedInDeviceState };
+    case grpcTypes.DeviceEvent.Cause.ROTATED_KEY:
+      return { type: 'rotated_key', deviceState: deviceState as LoggedInDeviceState };
+  }
 }
 
-function convertFromDeviceConfig(deviceConfig?: grpcTypes.DeviceConfig): IDeviceConfig | undefined {
-  const device = deviceConfig?.getDevice();
-  return (
-    deviceConfig && {
-      accountToken: deviceConfig.getAccountToken(),
-      device: device ? convertFromDevice(device) : undefined,
+function convertFromDeviceState(deviceState: grpcTypes.DeviceState): DeviceState {
+  switch (deviceState.getState()) {
+    case grpcTypes.DeviceState.State.LOGGED_IN: {
+      const accountAndDevice = deviceState.getDevice()!;
+      const device = accountAndDevice.getDevice();
+      return {
+        type: 'logged in',
+        accountAndDevice: {
+          accountToken: accountAndDevice.getAccountToken(),
+          device: device && convertFromDevice(device),
+        },
+      };
     }
-  );
+    case grpcTypes.DeviceState.State.LOGGED_OUT:
+      return { type: 'logged out' };
+    case grpcTypes.DeviceState.State.REVOKED:
+      return { type: 'revoked' };
+  }
 }
 
 function convertFromDeviceRemoval(deviceRemoval: grpcTypes.RemoveDeviceEvent): Array<IDevice> {
diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index c1fb5407f7..4d0f915bb9 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -27,10 +27,10 @@ import {
   BridgeSettings,
   BridgeState,
   DaemonEvent,
+  DeviceEvent,
+  DeviceState,
   IAccountData,
   IAppVersionInfo,
-  IDeviceConfig,
-  IDeviceEvent,
   IDeviceRemoval,
   IDnsOptions,
   IRelayList,
@@ -206,8 +206,7 @@ class ApplicationMain {
       },
     },
   };
-  private deviceConfig?: IDeviceConfig;
-  private hasReceivedDeviceConfig = false;
+  private deviceState?: DeviceState;
   private guiSettings = new GuiSettings();
   private tunnelStateExpectation?: Expectation;
 
@@ -681,10 +680,13 @@ class ApplicationMain {
 
     // fetch device
     try {
-      this.setDeviceConfig({ deviceConfig: await this.daemonRpc.getDevice() });
-      void this.daemonRpc
-        .updateDevice()
-        .catch((error: Error) => log.warn(`Failed to update device info: ${error.message}`));
+      const deviceState = await this.daemonRpc.getDevice();
+      this.handleDeviceEvent({ type: deviceState.type, deviceState } as DeviceEvent);
+      if (deviceState.type === 'logged in') {
+        void this.daemonRpc
+          .updateDevice()
+          .catch((error: Error) => log.warn(`Failed to update device info: ${error.message}`));
+      }
     } catch (e) {
       const error = e as Error;
       log.error(`Failed to fetch device: ${error.message}`);
@@ -749,7 +751,7 @@ class ApplicationMain {
     }
 
     // show window when account is not set
-    if (!this.deviceConfig) {
+    if (!this.isLoggedIn()) {
       this.windowController?.show();
     }
   };
@@ -822,7 +824,7 @@ class ApplicationMain {
         } else if ('appVersionInfo' in daemonEvent) {
           this.setLatestVersion(daemonEvent.appVersionInfo);
         } else if ('device' in daemonEvent) {
-          this.setDeviceConfig(daemonEvent.device);
+          this.handleDeviceEvent(daemonEvent.device);
         } else if ('deviceRemoval' in daemonEvent) {
           if (this.windowController) {
             IpcMainEventChannel.account.notifyDevices(
@@ -854,26 +856,14 @@ class ApplicationMain {
   }
 
   private connectTunnel = async (): Promise<void> => {
-    if (
-      connectEnabled(
-        this.connectedToDaemon,
-        this.deviceConfig?.accountToken,
-        this.tunnelState.state,
-      )
-    ) {
+    if (connectEnabled(this.connectedToDaemon, this.isLoggedIn(), this.tunnelState.state)) {
       this.setOptimisticTunnelState('connecting');
       await this.daemonRpc.connectTunnel();
     }
   };
 
   private reconnectTunnel = async (): Promise<void> => {
-    if (
-      reconnectEnabled(
-        this.connectedToDaemon,
-        this.deviceConfig?.accountToken,
-        this.tunnelState.state,
-      )
-    ) {
+    if (reconnectEnabled(this.connectedToDaemon, this.isLoggedIn(), this.tunnelState.state)) {
       this.setOptimisticTunnelState('connecting');
       await this.daemonRpc.reconnectTunnel();
     }
@@ -1166,28 +1156,41 @@ class ApplicationMain {
     }
   }
 
-  private setDeviceConfig(deviceEvent: IDeviceEvent) {
-    const oldDeviceConfig = this.deviceConfig;
-    this.deviceConfig = deviceEvent.deviceConfig;
-    this.hasReceivedDeviceConfig = true;
+  private handleDeviceEvent(deviceEvent: DeviceEvent) {
+    this.deviceState = deviceEvent.deviceState;
 
     if (this.isPerformingPostUpgrade) {
       void this.performPostUpgradeCheck();
     }
 
-    // make sure to invalidate the account data cache when account tokens change
-    this.updateAccountDataOnAccountChange(
-      oldDeviceConfig?.accountToken,
-      deviceEvent.deviceConfig?.accountToken,
-    );
+    switch (deviceEvent.deviceState.type) {
+      case 'logged in':
+        this.accountDataCache.fetch(deviceEvent.deviceState.accountAndDevice.accountToken);
+        break;
+      case 'logged out':
+      case 'revoked':
+        this.accountDataCache.invalidate();
+        break;
+    }
 
     void this.updateAccountHistory();
+    this.setTrayContextMenu();
 
     if (this.windowController) {
       IpcMainEventChannel.account.notifyDevice(this.windowController.webContents, deviceEvent);
     }
   }
 
+  private isLoggedIn(): boolean {
+    return this.deviceState?.type === 'logged in';
+  }
+
+  private getAccountToken(): AccountToken | undefined {
+    return this.deviceState?.type === 'logged in'
+      ? this.deviceState.accountAndDevice.accountToken
+      : undefined;
+  }
+
   private trayIconType(tunnelState: TunnelState, blockWhenDisconnected: boolean): TrayIconType {
     switch (tunnelState.state) {
       case 'connected':
@@ -1270,8 +1273,7 @@ class ApplicationMain {
       tunnelState: this.tunnelState,
       settings: this.settings,
       isPerformingPostUpgrade: this.isPerformingPostUpgrade,
-      deviceConfig: this.deviceConfig,
-      hasReceivedDeviceConfig: this.hasReceivedDeviceConfig,
+      deviceState: this.deviceState,
       relayListPair: {
         relays: this.processRelaysForPresentation(this.relays, this.settings.relaySettings),
         bridges: this.processBridgesForPresentation(this.relays, this.settings.bridgeState),
@@ -1363,7 +1365,7 @@ class ApplicationMain {
     IpcMainEventChannel.account.handleLogout(() => this.logout());
     IpcMainEventChannel.account.handleGetWwwAuthToken(() => this.daemonRpc.getWwwAuthToken());
     IpcMainEventChannel.account.handleSubmitVoucher(async (voucherCode: string) => {
-      const currentAccountToken = this.deviceConfig?.accountToken;
+      const currentAccountToken = this.getAccountToken();
       const response = await this.daemonRpc.submitVoucher(voucherCode);
 
       if (currentAccountToken) {
@@ -1374,9 +1376,9 @@ class ApplicationMain {
     });
     IpcMainEventChannel.account.handleUpdateData(() => this.updateAccountData());
 
-    IpcMainEventChannel.account.handleGetDevice(async () => {
-      const deviceConfig = await this.daemonRpc.getDevice();
-      return deviceConfig?.device;
+    IpcMainEventChannel.account.handleGetDeviceState(async () => {
+      await this.daemonRpc.updateDevice();
+      return this.daemonRpc.getDevice();
     });
     IpcMainEventChannel.account.handleListDevices((accountToken: AccountToken) => {
       return this.daemonRpc.listDevices(accountToken);
@@ -1540,7 +1542,7 @@ class ApplicationMain {
   private async autoConnect() {
     if (process.env.NODE_ENV === 'development') {
       log.info('Skip autoconnect in development');
-    } else if (this.deviceConfig && (!this.accountData || !hasExpired(this.accountData.expiry))) {
+    } else if (this.isLoggedIn() && (!this.accountData || !hasExpired(this.accountData.expiry))) {
       if (this.guiSettings.autoConnect) {
         try {
           log.info('Autoconnect the tunnel');
@@ -1571,21 +1573,9 @@ class ApplicationMain {
     }
   }
 
-  private updateAccountDataOnAccountChange(oldAccount?: string, newAccount?: string) {
-    if (oldAccount && !newAccount) {
-      this.accountDataCache.invalidate();
-    } else if (!oldAccount && newAccount) {
-      this.accountDataCache.fetch(newAccount);
-    } else if (oldAccount && newAccount && oldAccount !== newAccount) {
-      this.accountDataCache.fetch(newAccount);
-    }
-
-    this.setTrayContextMenu();
-  }
-
   private updateAccountData() {
-    if (this.connectedToDaemon && this.deviceConfig) {
-      this.accountDataCache.fetch(this.deviceConfig.accountToken);
+    if (this.connectedToDaemon && this.isLoggedIn()) {
+      this.accountDataCache.fetch(this.getAccountToken()!);
     }
   }
 
@@ -1971,7 +1961,7 @@ class ApplicationMain {
           this.tray?.on('right-click', () =>
             this.trayIconController?.popUpContextMenu(
               this.connectedToDaemon,
-              this.deviceConfig?.accountToken,
+              this.isLoggedIn(),
               this.tunnelState,
             ),
           );
@@ -2009,7 +1999,7 @@ class ApplicationMain {
   private setTrayContextMenu() {
     this.trayIconController?.setContextMenu(
       this.connectedToDaemon,
-      this.deviceConfig?.accountToken,
+      this.isLoggedIn(),
       this.tunnelState,
     );
   }
diff --git a/gui/src/main/tray-icon-controller.ts b/gui/src/main/tray-icon-controller.ts
index 81af2e174e..723228eac6 100644
--- a/gui/src/main/tray-icon-controller.ts
+++ b/gui/src/main/tray-icon-controller.ts
@@ -5,7 +5,7 @@ import { sprintf } from 'sprintf-js';
 import { promisify } from 'util';
 
 import { connectEnabled, disconnectEnabled, reconnectEnabled } from '../shared/connect-helper';
-import { AccountToken, ILocation, TunnelState } from '../shared/daemon-rpc-types';
+import { ILocation, TunnelState } from '../shared/daemon-rpc-types';
 import { messages, relayLocations } from '../shared/gettext';
 import log from '../shared/logging';
 import KeyframeAnimation from './keyframe-animation';
@@ -102,15 +102,9 @@ export default class TrayIconController {
     animation.play({ end: frame });
   }
 
-  public setContextMenu(
-    connectedToDaemon: boolean,
-    accountToken: AccountToken | undefined,
-    tunnelState: TunnelState,
-  ) {
+  public setContextMenu(connectedToDaemon: boolean, loggedIn: boolean, tunnelState: TunnelState) {
     if (process.platform === 'linux') {
-      this.tray.setContextMenu(
-        this.createContextMenu(connectedToDaemon, accountToken, tunnelState),
-      );
+      this.tray.setContextMenu(this.createContextMenu(connectedToDaemon, loggedIn, tunnelState));
     }
   }
 
@@ -119,14 +113,8 @@ export default class TrayIconController {
     this.tray?.setToolTip(tooltip);
   }
 
-  public popUpContextMenu(
-    connectedToDaemon: boolean,
-    accountToken: AccountToken | undefined,
-    tunnelState: TunnelState,
-  ) {
-    this.tray.popUpContextMenu(
-      this.createContextMenu(connectedToDaemon, accountToken, tunnelState),
-    );
+  public popUpContextMenu(connectedToDaemon: boolean, loggedIn: boolean, tunnelState: TunnelState) {
+    this.tray.popUpContextMenu(this.createContextMenu(connectedToDaemon, loggedIn, tunnelState));
   }
 
   private createTooltipText(connectedToDaemon: boolean, tunnelState: TunnelState): string {
@@ -267,7 +255,7 @@ export default class TrayIconController {
 
   private createContextMenu(
     connectedToDaemon: boolean,
-    accountToken: AccountToken | undefined,
+    loggedIn: boolean,
     tunnelState: TunnelState,
   ) {
     const template: Electron.MenuItemConstructorOptions[] = [
@@ -281,13 +269,13 @@ export default class TrayIconController {
       {
         id: 'connect',
         label: messages.gettext('Connect'),
-        enabled: connectEnabled(connectedToDaemon, accountToken, tunnelState.state),
+        enabled: connectEnabled(connectedToDaemon, loggedIn, tunnelState.state),
         click: this.connect,
       },
       {
         id: 'reconnect',
         label: messages.gettext('Reconnect'),
-        enabled: reconnectEnabled(connectedToDaemon, accountToken, tunnelState.state),
+        enabled: reconnectEnabled(connectedToDaemon, loggedIn, tunnelState.state),
         click: this.reconnect,
       },
       {