Refactor optimistic tunnel state

3 files changed, 130 insertions, 105 deletions

diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index 05ee5cb3b3..b5743b7c76 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -76,6 +76,7 @@ import ReconnectionBackoff from './reconnection-backoff';
 import TrayIconController, { TrayIconType } from './tray-icon-controller';
 import WindowController from './window-controller';
 import { ITranslations, MacOsScrollbarVisibility } from '../shared/ipc-schema';
+import { connectEnabled, disconnectEnabled, reconnectEnabled } from '../shared/connect-helper';
 
 const execAsync = util.promisify(exec);
 
@@ -132,10 +133,16 @@ class ApplicationMain {
 
   private accountData?: IAccountData = undefined;
   private accountHistory?: AccountToken = undefined;
+
+  // The current tunnel state
   private tunnelState: TunnelState = { state: 'disconnected' };
-  private lastIgnoredTunnelState?: TunnelState;
-  private optimisticTunnelState?: TunnelState['state'];
-  private optimisticTunnelStateFallbackScheduler = new Scheduler();
+  // When pressing connect/disconnect/reconnect the app assumes what the next state will be before
+  // it get's the new state from the daemon. The latest state from the daemon is saved as fallback
+  // if the assumed state isn't reached.
+  private tunnelStateFallback?: TunnelState;
+  // Scheduler for discarding the assumed next state.
+  private tunnelStateFallbackScheduler = new Scheduler();
+
   private settings: ISettings = {
     accountToken: undefined,
     allowLan: false,
@@ -701,8 +708,7 @@ class ApplicationMain {
     // Reset the daemon event listener since it's going to be invalidated on disconnect
     this.daemonEventListener = undefined;
 
-    this.optimisticTunnelState = undefined;
-    this.lastIgnoredTunnelState = undefined;
+    this.tunnelStateFallback = undefined;
     this.autoConnectFallbackScheduler.cancel();
 
     if (wasConnected) {
@@ -772,6 +778,31 @@ class ApplicationMain {
     return daemonEventListener;
   }
 
+  private connectTunnel = async (): Promise<void> => {
+    if (
+      connectEnabled(this.connectedToDaemon, this.settings.accountToken, this.tunnelState.state)
+    ) {
+      this.setOptimisticTunnelState('connecting');
+      await this.daemonRpc.connectTunnel();
+    }
+  };
+
+  private reconnectTunnel = async (): Promise<void> => {
+    if (
+      reconnectEnabled(this.connectedToDaemon, this.settings.accountToken, this.tunnelState.state)
+    ) {
+      this.setOptimisticTunnelState('connecting');
+      await this.daemonRpc.reconnectTunnel();
+    }
+  };
+
+  private disconnectTunnel = async (): Promise<void> => {
+    if (disconnectEnabled(this.connectedToDaemon, this.tunnelState.state)) {
+      this.setOptimisticTunnelState('disconnecting');
+      await this.daemonRpc.disconnectTunnel();
+    }
+  };
+
   private setAccountHistory(accountHistory?: AccountToken) {
     this.accountHistory = accountHistory;
 
@@ -811,33 +842,47 @@ class ApplicationMain {
     this.wireguardKeygenEventAutoConnect();
   }
 
+  // This function sets a new tunnel state as an assumed next state and saves the current state as
+  // fallback. The fallback is used if the assumed next state isn't reached.
   private setOptimisticTunnelState(state: 'connecting' | 'disconnecting') {
-    const tunnelState =
-      state === 'disconnecting' ? ({ state, details: 'nothing' } as const) : { state };
-    this.updateTrayIcon(tunnelState, this.settings.blockWhenDisconnected);
+    this.tunnelStateFallback = this.tunnelState;
+
+    this.setTunnelStateImpl(
+      state === 'disconnecting' ? { state, details: 'nothing' as const } : { state },
+    );
 
-    this.optimisticTunnelState = state;
-    this.optimisticTunnelStateFallbackScheduler.schedule(() => {
-      if (this.lastIgnoredTunnelState) {
-        this.optimisticTunnelState = undefined;
-        this.setTunnelState(this.lastIgnoredTunnelState);
-        this.lastIgnoredTunnelState = undefined;
+    this.tunnelStateFallbackScheduler.schedule(() => {
+      if (this.tunnelStateFallback) {
+        this.setTunnelStateImpl(this.tunnelStateFallback);
+        this.tunnelStateFallback = undefined;
       }
     }, 3000);
   }
 
   private setTunnelState(newState: TunnelState) {
-    if (this.optimisticTunnelState) {
-      if (this.optimisticTunnelState === newState.state) {
-        this.optimisticTunnelStateFallbackScheduler.cancel();
-        this.optimisticTunnelState = undefined;
-        this.lastIgnoredTunnelState = undefined;
+    // If there's a fallback state set then the app is in an assumed next state and need to check
+    // if it's now reached or if the current state should be ignored and set as the fallback state.
+    if (this.tunnelStateFallback) {
+      if (this.tunnelState.state === newState.state) {
+        this.tunnelStateFallbackScheduler.cancel();
+        this.tunnelStateFallback = undefined;
       } else {
-        this.lastIgnoredTunnelState = newState;
+        this.tunnelStateFallback = newState;
         return;
       }
     }
 
+    if (newState.state === 'disconnecting' && newState.details === 'reconnect') {
+      // When reconnecting there's no need of showing the disconnecting state. This switches to the
+      // connecting state immediately.
+      this.setOptimisticTunnelState('connecting');
+      this.tunnelStateFallback = newState;
+    } else {
+      this.setTunnelStateImpl(newState);
+    }
+  }
+
+  private setTunnelStateImpl(newState: TunnelState) {
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
 
@@ -1209,18 +1254,9 @@ class ApplicationMain {
 
     IpcMainEventChannel.location.handleGet(() => this.daemonRpc.getLocation());
 
-    IpcMainEventChannel.tunnel.handleConnect(() => {
-      this.setOptimisticTunnelState('connecting');
-      return this.daemonRpc.connectTunnel();
-    });
-    IpcMainEventChannel.tunnel.handleDisconnect(() => {
-      this.setOptimisticTunnelState('disconnecting');
-      return this.daemonRpc.disconnectTunnel();
-    });
-    IpcMainEventChannel.tunnel.handleReconnect(() => {
-      this.setOptimisticTunnelState('connecting');
-      return this.daemonRpc.reconnectTunnel();
-    });
+    IpcMainEventChannel.tunnel.handleConnect(this.connectTunnel);
+    IpcMainEventChannel.tunnel.handleReconnect(this.reconnectTunnel);
+    IpcMainEventChannel.tunnel.handleDisconnect(this.disconnectTunnel);
 
     IpcMainEventChannel.guiSettings.handleSetEnableSystemNotifications((flag: boolean) => {
       this.guiSettings.enableSystemNotifications = flag;
diff --git a/gui/src/renderer/app.tsx b/gui/src/renderer/app.tsx
index 0419bb0919..7fc2411855 100644
--- a/gui/src/renderer/app.tsx
+++ b/gui/src/renderer/app.tsx
@@ -91,7 +91,6 @@ export default class AppRenderer {
   private lastDisconnectedLocation?: Partial<ILocation>;
   private relayListPair!: IRelayListPair;
   private tunnelState!: TunnelState;
-  private optimisticTunnelState?: TunnelState['state'];
   private settings!: ISettings;
   private guiSettings!: IGuiSettingsState;
   private doingLogin = false;
@@ -312,51 +311,15 @@ export default class AppRenderer {
   }
 
   public async connectTunnel(): Promise<void> {
-    const state = this.optimisticTunnelState ?? this.tunnelState.state;
-
-    // connect only if tunnel is disconnected or blocked.
-    if (state === 'disconnecting' || state === 'disconnected' || state === 'error') {
-      // switch to the connecting state ahead of time to make the app look more responsive
-      this.optimisticTunnelState = 'connecting';
-      batch(() => {
-        void this.updateLocation({ state: 'connecting' });
-        this.reduxActions.connection.connecting();
-      });
-
-      return IpcRendererEventChannel.tunnel.connect();
-    }
+    return IpcRendererEventChannel.tunnel.connect();
   }
 
   public async disconnectTunnel(): Promise<void> {
-    const state = this.optimisticTunnelState ?? this.tunnelState.state;
-
-    // disconnect only if tunnel is connected, connecting or blocked.
-    if (state === 'connecting' || state === 'connected' || state === 'error') {
-      // switch to the disconnecting state ahead of time to make the app look more responsive
-      this.optimisticTunnelState = 'disconnecting';
-      batch(() => {
-        void this.updateLocation({ state: 'disconnecting', details: 'nothing' });
-        this.reduxActions.connection.disconnecting('nothing');
-      });
-
-      return IpcRendererEventChannel.tunnel.disconnect();
-    }
+    return IpcRendererEventChannel.tunnel.disconnect();
   }
 
   public async reconnectTunnel(): Promise<void> {
-    const state = this.optimisticTunnelState ?? this.tunnelState.state;
-
-    // reconnect only if tunnel is connected or connecting.
-    if (state === 'connecting' || state === 'connected') {
-      // switch to the connecting state ahead of time to make the app look more responsive
-      this.optimisticTunnelState = 'connecting';
-      batch(() => {
-        void this.updateLocation({ state: 'connecting' });
-        this.reduxActions.connection.connecting();
-      });
-
-      return IpcRendererEventChannel.tunnel.reconnect();
-    }
+    return IpcRendererEventChannel.tunnel.reconnect();
   }
 
   public updateRelaySettings(relaySettings: RelaySettingsUpdate) {
@@ -724,41 +687,33 @@ export default class AppRenderer {
     log.verbose(`Tunnel state: ${tunnelState.state}`);
 
     this.tunnelState = tunnelState;
-    // The main process doesn't notify the tunnel state while waiting for a new one (unless it times
-    // out). Therefore the first tunnel state update will be the one we're waiting for.
-    this.optimisticTunnelState = undefined;
 
-    switch (tunnelState.state) {
-      case 'connecting':
-        actions.connection.connecting(tunnelState.details);
-        break;
+    batch(() => {
+      switch (tunnelState.state) {
+        case 'connecting':
+          actions.connection.connecting(tunnelState.details);
+          break;
 
-      case 'connected':
-        actions.connection.connected(tunnelState.details);
-        break;
+        case 'connected':
+          actions.connection.connected(tunnelState.details);
+          break;
 
-      case 'disconnecting':
-        if (tunnelState.details === 'reconnect') {
-          this.optimisticTunnelState = 'connecting';
-          this.reduxActions.connection.connecting();
-        } else {
+        case 'disconnecting':
           actions.connection.disconnecting(tunnelState.details);
-        }
-        break;
+          break;
 
-      case 'disconnected':
-        actions.connection.disconnected();
-        break;
+        case 'disconnected':
+          actions.connection.disconnected();
+          break;
 
-      case 'error':
-        actions.connection.blocked(tunnelState.details);
-        break;
-    }
+        case 'error':
+          actions.connection.blocked(tunnelState.details);
+          break;
+      }
 
-    // Update the location when entering a new tunnel state since it's likely changed.
-    void this.updateLocation(
-      this.optimisticTunnelState === undefined ? undefined : { state: this.optimisticTunnelState },
-    );
+      // Update the location when entering a new tunnel state since it's likely changed.
+      void this.updateLocation();
+    });
   }
 
   private setSettings(newSettings: ISettings) {
@@ -876,8 +831,8 @@ export default class AppRenderer {
     this.reduxActions.settings.setWireguardKey(publicKey);
   }
 
-  private async updateLocation(tunnelState = this.tunnelState) {
-    switch (tunnelState.state) {
+  private async updateLocation() {
+    switch (this.tunnelState.state) {
       case 'disconnected': {
         if (this.lastDisconnectedLocation) {
           this.setLocation(this.lastDisconnectedLocation);
@@ -900,11 +855,11 @@ export default class AppRenderer {
         }
         break;
       case 'connecting':
-        this.setLocation(tunnelState.details?.location ?? this.getLocationFromConstraints());
+        this.setLocation(this.tunnelState.details?.location ?? this.getLocationFromConstraints());
         break;
       case 'connected': {
-        if (tunnelState.details?.location) {
-          this.setLocation(tunnelState.details.location);
+        if (this.tunnelState.details?.location) {
+          this.setLocation(this.tunnelState.details.location);
         }
         const location = await this.fetchLocation();
         if (location) {
diff --git a/gui/src/shared/connect-helper.ts b/gui/src/shared/connect-helper.ts
new file mode 100644
index 0000000000..214d22a59a
--- /dev/null
+++ b/gui/src/shared/connect-helper.ts
@@ -0,0 +1,34 @@
+import { AccountToken, TunnelState } from './daemon-rpc-types';
+
+export function connectEnabled(
+  connectedToDaemon: boolean,
+  accountToken: AccountToken | undefined,
+  tunnelState: TunnelState['state'],
+) {
+  return (
+    connectedToDaemon &&
+    accountToken !== undefined &&
+    (tunnelState === 'disconnected' || tunnelState === 'disconnecting' || tunnelState === 'error')
+  );
+}
+
+export function reconnectEnabled(
+  connectedToDaemon: boolean,
+  accountToken: AccountToken | undefined,
+  tunnelState: TunnelState['state'],
+) {
+  return (
+    connectedToDaemon &&
+    accountToken !== undefined &&
+    (tunnelState === 'connected' || tunnelState === 'connecting')
+  );
+}
+
+// Disconnecting while logged out is allowed since it's possible to "connect" and end up in the
+// blocked state with the CLI.
+export function disconnectEnabled(connectedToDaemon: boolean, tunnelState: TunnelState['state']) {
+  return (
+    connectedToDaemon &&
+    (tunnelState === 'connected' || tunnelState === 'connecting' || tunnelState === 'error')
+  );
+}