Merge branch 'fix-disconnecting-info'

7 files changed, 177 insertions, 121 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 12e0df97ea..f30be7bfd7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -62,6 +62,9 @@ Line wrap the file at 100 chars.                                              Th
 - Fix in-app notification button not working for some notifications.
 - Fix incorrectly positioned navigation bar title when navigating back to a scrolled down view.
 - Fix connectivity check for WireGuard multihop when the exit hop is down.
+- Fix incorrect location and connection status while disconnecting and incorrect location in the
+  beginning while connecting in the desktop app.
+- Improve responsiveness of the controls and status text in the main view in the desktop app.
 
 #### Linux
 - Make offline monitor aware of routing table changes.
diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index 8845422e6b..8fa09576c7 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -26,7 +26,6 @@ import {
   IAccountData,
   IAppVersionInfo,
   IDnsOptions,
-  ILocation,
   IRelayList,
   ISettings,
   IWireguardPublicKey,
@@ -125,6 +124,9 @@ class ApplicationMain {
   private accountData?: IAccountData = undefined;
   private accountHistory?: AccountToken = undefined;
   private tunnelState: TunnelState = { state: 'disconnected' };
+  private lastIgnoredTunnelState?: TunnelState;
+  private ignoreTunnelStatesUntil?: TunnelState['state'];
+  private ignoreTunnelStateFallbackScheduler = new Scheduler();
   private settings: ISettings = {
     accountToken: undefined,
     allowLan: false,
@@ -179,10 +181,7 @@ class ApplicationMain {
     },
   };
   private guiSettings = new GuiSettings();
-  private location?: ILocation;
-  private lastDisconnectedLocation?: ILocation;
   private tunnelStateExpectation?: Expectation;
-  private getLocationPromise?: Promise<ILocation>;
 
   private relays: IRelayList = { countries: [] };
 
@@ -624,6 +623,10 @@ class ApplicationMain {
     // Reset the daemon event listener since it's going to be invalidated on disconnect
     this.daemonEventListener = undefined;
 
+    this.ignoreTunnelStatesUntil = undefined;
+    this.lastIgnoredTunnelState = undefined;
+    this.autoConnectFallbackScheduler.cancel();
+
     if (wasConnected) {
       this.connectedToDaemon = false;
 
@@ -729,10 +732,30 @@ class ApplicationMain {
     this.wireguardKeygenEventAutoConnect();
   }
 
+  private setIgnoreTunnelStatesUntil(state: TunnelState['state']) {
+    this.ignoreTunnelStatesUntil = state;
+    this.ignoreTunnelStateFallbackScheduler.schedule(() => {
+      if (this.lastIgnoredTunnelState) {
+        this.ignoreTunnelStatesUntil = undefined;
+        this.setTunnelState(this.lastIgnoredTunnelState);
+        this.lastIgnoredTunnelState = undefined;
+      }
+    }, 3000);
+  }
+
   private setTunnelState(newState: TunnelState) {
+    if (this.ignoreTunnelStatesUntil) {
+      if (this.ignoreTunnelStatesUntil === newState.state) {
+        this.ignoreTunnelStatesUntil = undefined;
+        this.lastIgnoredTunnelState = undefined;
+      } else {
+        this.lastIgnoredTunnelState = newState;
+        return;
+      }
+    }
+
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
-    void this.updateLocation();
 
     if (process.platform === 'linux') {
       this.tray?.setContextMenu(this.createTrayContextMenu());
@@ -799,14 +822,6 @@ class ApplicationMain {
     }
   }
 
-  private setLocation(newLocation: ILocation) {
-    this.location = newLocation;
-
-    if (this.windowController) {
-      IpcMainEventChannel.location.notify(this.windowController.webContents, newLocation);
-    }
-  }
-
   private setRelays(
     newRelayList: IRelayList,
     relaySettings: RelaySettings,
@@ -981,50 +996,6 @@ class ApplicationMain {
     }
   }
 
-  private async updateLocation() {
-    const tunnelState = this.tunnelState;
-
-    if (tunnelState.state === 'connected' || tunnelState.state === 'connecting') {
-      // Location was broadcasted with the tunnel state, but it doesn't contain the relay out IP
-      // address, so it will have to be fetched afterwards
-      if (tunnelState.details && tunnelState.details.location) {
-        this.setLocation(tunnelState.details.location);
-      }
-    } else if (tunnelState.state === 'disconnected') {
-      // It may take some time to fetch the new user location.
-      // So take the user to the last known location when disconnected.
-      if (this.lastDisconnectedLocation) {
-        this.setLocation(this.lastDisconnectedLocation);
-      }
-    }
-
-    if (tunnelState.state === 'connected' || tunnelState.state === 'disconnected') {
-      try {
-        // Fetch the new user location
-        const getLocationPromise = (this.getLocationPromise = this.daemonRpc.getLocation());
-        const location = await getLocationPromise;
-        // If the location is currently unavailable, do nothing! This only ever
-        // happens when a custom relay is set or we are in a blocked state.
-        if (!location) {
-          return;
-        }
-
-        // Cache the user location
-        // Note: hostname is only set for relay servers.
-        if (location.hostname === null) {
-          this.lastDisconnectedLocation = location;
-        }
-
-        // Broadcast the new location if it is the result of the most recent call to getLocation.
-        if (getLocationPromise === this.getLocationPromise) {
-          this.setLocation(location);
-        }
-      } catch (error) {
-        log.error(`Failed to update the location: ${error.message}`);
-      }
-    }
-  }
-
   private trayIconType(tunnelState: TunnelState, blockWhenDisconnected: boolean): TrayIconType {
     switch (tunnelState.state) {
       case 'connected':
@@ -1102,7 +1073,6 @@ class ApplicationMain {
       accountHistory: this.accountHistory,
       tunnelState: this.tunnelState,
       settings: this.settings,
-      location: this.location,
       relayListPair: {
         relays: this.processRelaysForPresentation(
           this.relays,
@@ -1158,9 +1128,20 @@ class ApplicationMain {
       return this.setAutoStart(autoStart);
     });
 
-    IpcMainEventChannel.tunnel.handleConnect(() => this.daemonRpc.connectTunnel());
-    IpcMainEventChannel.tunnel.handleDisconnect(() => this.daemonRpc.disconnectTunnel());
-    IpcMainEventChannel.tunnel.handleReconnect(() => this.daemonRpc.reconnectTunnel());
+    IpcMainEventChannel.location.handleGet(() => this.daemonRpc.getLocation());
+
+    IpcMainEventChannel.tunnel.handleConnect(() => {
+      this.setIgnoreTunnelStatesUntil('connecting');
+      return this.daemonRpc.connectTunnel();
+    });
+    IpcMainEventChannel.tunnel.handleDisconnect(() => {
+      this.setIgnoreTunnelStatesUntil('disconnecting');
+      return this.daemonRpc.disconnectTunnel();
+    });
+    IpcMainEventChannel.tunnel.handleReconnect(() => {
+      this.setIgnoreTunnelStatesUntil('connecting');
+      return this.daemonRpc.reconnectTunnel();
+    });
 
     IpcMainEventChannel.guiSettings.handleSetEnableSystemNotifications((flag: boolean) => {
       this.guiSettings.enableSystemNotifications = flag;
diff --git a/gui/src/renderer/app.tsx b/gui/src/renderer/app.tsx
index 28d2c8b480..dd8a6543ca 100644
--- a/gui/src/renderer/app.tsx
+++ b/gui/src/renderer/app.tsx
@@ -1,5 +1,5 @@
 import * as React from 'react';
-import { Provider } from 'react-redux';
+import { batch, Provider } from 'react-redux';
 import { Router } from 'react-router';
 import { bindActionCreators } from 'redux';
 
@@ -116,14 +116,17 @@ export default class AppRenderer {
   };
 
   private locale = 'en';
-  private location?: ILocation;
+  private location?: Partial<ILocation>;
+  private lastDisconnectedLocation?: Partial<ILocation>;
   private relayListPair!: IRelayListPair;
   private tunnelState!: TunnelState;
+  private optimisticTunnelState?: TunnelState['state'];
   private settings!: ISettings;
   private guiSettings!: IGuiSettingsState;
   private doingLogin = false;
   private loginScheduler = new Scheduler();
   private connectedToDaemon = false;
+  private getLocationPromise?: Promise<ILocation>;
 
   constructor() {
     log.addOutput(new ConsoleOutput(LogLevel.debug));
@@ -164,10 +167,6 @@ export default class AppRenderer {
       this.updateBlockedState(this.tunnelState, newSettings.blockWhenDisconnected);
     });
 
-    IpcRendererEventChannel.location.listen((newLocation: ILocation) => {
-      this.setLocation(newLocation);
-    });
-
     IpcRendererEventChannel.relays.listen((relayListPair: IRelayListPair) => {
       this.setRelayListPair(relayListPair);
     });
@@ -231,10 +230,6 @@ export default class AppRenderer {
     this.setTunnelState(initialState.tunnelState);
     this.updateBlockedState(initialState.tunnelState, initialState.settings.blockWhenDisconnected);
 
-    if (initialState.location) {
-      this.setLocation(initialState.location);
-    }
-
     this.setRelayListPair(initialState.relayListPair);
     this.setCurrentVersion(initialState.currentVersion);
     this.setUpgradeVersion(initialState.upgradeVersion);
@@ -254,6 +249,8 @@ export default class AppRenderer {
       );
     }
 
+    void this.updateLocation();
+
     const navigationBase = this.getNavigationBase(
       initialState.isConnected,
       initialState.settings.accountToken,
@@ -327,38 +324,48 @@ export default class AppRenderer {
   }
 
   public async connectTunnel(): Promise<void> {
-    const state = this.tunnelState.state;
+    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.resetLocationToConstraints();
-      this.reduxActions.connection.connecting();
+      this.optimisticTunnelState = 'connecting';
+      batch(() => {
+        void this.updateLocation({ state: 'connecting' });
+        this.reduxActions.connection.connecting();
+      });
 
       return IpcRendererEventChannel.tunnel.connect();
     }
   }
 
   public async disconnectTunnel(): Promise<void> {
-    const state = this.tunnelState.state;
+    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.reduxActions.connection.disconnecting('nothing');
+      this.optimisticTunnelState = 'disconnecting';
+      batch(() => {
+        void this.updateLocation({ state: 'disconnecting', details: 'nothing' });
+        this.reduxActions.connection.disconnecting('nothing');
+      });
 
       return IpcRendererEventChannel.tunnel.disconnect();
     }
   }
 
   public async reconnectTunnel(): Promise<void> {
-    const state = this.tunnelState.state;
+    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.resetLocationToConstraints();
-      this.reduxActions.connection.connecting();
+      this.optimisticTunnelState = 'connecting';
+      batch(() => {
+        void this.updateLocation({ state: 'connecting' });
+        this.reduxActions.connection.connecting();
+      });
 
       return IpcRendererEventChannel.tunnel.reconnect();
     }
@@ -605,37 +612,6 @@ export default class AppRenderer {
     this.reduxActions.userInterface.updateLocale(locale);
   }
 
-  private resetLocationToConstraints() {
-    const relaySettings = this.settings.relaySettings;
-    if ('normal' in relaySettings) {
-      const location = relaySettings.normal.location;
-      if (location !== 'any' && 'only' in location) {
-        const constraint = location.only;
-
-        const relayLocations = this.reduxStore.getState().settings.relayLocations;
-        if ('country' in constraint) {
-          const country = relayLocations.find(({ code }) => constraint.country === code);
-
-          this.reduxActions.connection.newLocation({ country: country?.name });
-        } else if ('city' in constraint) {
-          const country = relayLocations.find(({ code }) => constraint.city[0] === code);
-          const city = country?.cities.find(({ code }) => constraint.city[1] === code);
-
-          this.reduxActions.connection.newLocation({ country: country?.name, city: city?.name });
-        } else if ('hostname' in constraint) {
-          const country = relayLocations.find(({ code }) => constraint.hostname[0] === code);
-          const city = country?.cities.find((location) => location.code === constraint.hostname[1]);
-
-          this.reduxActions.connection.newLocation({
-            country: country?.name,
-            city: city?.name,
-            hostname: constraint.hostname[2],
-          });
-        }
-      }
-    }
-  }
-
   private setRelaySettings(relaySettings: RelaySettings) {
     const actions = this.reduxActions;
 
@@ -777,6 +753,9 @@ export default class AppRenderer {
         actions.connection.blocked(tunnelState.details);
         break;
     }
+
+    // Update the location when entering a new tunnel state since it's likely changed.
+    void this.updateLocation();
   }
 
   private setSettings(newSettings: ISettings) {
@@ -841,7 +820,7 @@ export default class AppRenderer {
     this.doingLogin = false;
   }
 
-  private setLocation(location: ILocation) {
+  private setLocation(location: Partial<ILocation>) {
     this.location = location;
     this.propagateLocationToRedux();
   }
@@ -852,7 +831,7 @@ export default class AppRenderer {
     }
   }
 
-  private translateLocation(inputLocation: ILocation): ILocation {
+  private translateLocation(inputLocation: Partial<ILocation>): Partial<ILocation> {
     const location = { ...inputLocation };
 
     if (location.city) {
@@ -933,4 +912,94 @@ export default class AppRenderer {
   private setWireguardPublicKey(publicKey?: IWireguardPublicKey) {
     this.reduxActions.settings.setWireguardKey(publicKey);
   }
+
+  private async updateLocation(tunnelState = this.tunnelState) {
+    if (
+      (tunnelState.state === 'disconnected' || tunnelState.state === 'disconnecting') &&
+      this.lastDisconnectedLocation
+    ) {
+      // If we have a previous location for the disconnected state we use that when disconnecting and
+      // during the location fetch while disconnected.
+      this.setLocation(this.lastDisconnectedLocation);
+    } else if (tunnelState.state === 'disconnecting') {
+      // If there's no previous location while disconnecting we remove the location. We keep the
+      // coordinates to prevent the map from jumping around.
+      const { longitude, latitude } = this.reduxStore.getState().connection;
+      this.setLocation({ longitude, latitude });
+    }
+
+    if (
+      (tunnelState.state === 'connected' || tunnelState.state === 'connecting') &&
+      tunnelState.details?.location
+    ) {
+      this.setLocation(tunnelState.details.location);
+    } else if (tunnelState.state === 'connecting') {
+      this.setLocation(this.getLocationFromConstraints());
+    } else if (tunnelState.state === 'connected' || tunnelState.state === 'disconnected') {
+      const location = await this.fetchLocation();
+      if (location) {
+        this.setLocation(location);
+        // Cache the user location
+        if (tunnelState.state === 'disconnected') {
+          this.lastDisconnectedLocation = location;
+        }
+      }
+    }
+  }
+
+  private async fetchLocation(): Promise<ILocation | void> {
+    try {
+      // Fetch the new user location
+      const getLocationPromise = IpcRendererEventChannel.location.get();
+      this.getLocationPromise = getLocationPromise;
+      const location = await getLocationPromise;
+      // If the location is currently unavailable, do nothing! This only ever happens when a
+      // custom relay is set or we are in a blocked state.
+      if (location && getLocationPromise === this.getLocationPromise) {
+        return location;
+      }
+    } catch (error) {
+      log.error(`Failed to update the location: ${error.message}`);
+    }
+  }
+
+  private getLocationFromConstraints(): Partial<ILocation> {
+    const state = this.reduxStore.getState();
+    const coordinates = {
+      longitude: state.connection.longitude,
+      latitude: state.connection.latitude,
+    };
+
+    const relaySettings = this.settings.relaySettings;
+    if ('normal' in relaySettings) {
+      const location = relaySettings.normal.location;
+      if (location !== 'any' && 'only' in location) {
+        const constraint = location.only;
+
+        const relayLocations = state.settings.relayLocations;
+        if ('country' in constraint) {
+          const country = relayLocations.find(({ code }) => constraint.country === code);
+
+          return { country: country?.name, ...coordinates };
+        } else if ('city' in constraint) {
+          const country = relayLocations.find(({ code }) => constraint.city[0] === code);
+          const city = country?.cities.find(({ code }) => constraint.city[1] === code);
+
+          return { country: country?.name, city: city?.name, ...coordinates };
+        } else if ('hostname' in constraint) {
+          const country = relayLocations.find(({ code }) => constraint.hostname[0] === code);
+          const city = country?.cities.find((location) => location.code === constraint.hostname[1]);
+
+          return {
+            country: country?.name,
+            city: city?.name,
+            hostname: constraint.hostname[2],
+            ...coordinates,
+          };
+        }
+      }
+    }
+
+    return coordinates;
+  }
 }
diff --git a/gui/src/renderer/components/Connect.tsx b/gui/src/renderer/components/Connect.tsx
index e48abf8d96..a150da90f1 100644
--- a/gui/src/renderer/components/Connect.tsx
+++ b/gui/src/renderer/components/Connect.tsx
@@ -213,9 +213,6 @@ export default class Connect extends React.Component<IProps, IState> {
   private showMarkerOrSpinner(): MarkerOrSpinner {
     const status = this.props.connection.status;
 
-    return status.state === 'connecting' ||
-      (status.state === 'disconnecting' && status.details === 'reconnect')
-      ? 'spinner'
-      : 'marker';
+    return status.state === 'connecting' || status.state === 'disconnecting' ? 'spinner' : 'marker';
   }
 }
diff --git a/gui/src/renderer/components/SecuredLabel.tsx b/gui/src/renderer/components/SecuredLabel.tsx
index 2d80991ed3..c00edc73de 100644
--- a/gui/src/renderer/components/SecuredLabel.tsx
+++ b/gui/src/renderer/components/SecuredLabel.tsx
@@ -8,11 +8,13 @@ export enum SecuredDisplayStyle {
   blocked,
   securing,
   unsecured,
+  unsecuring,
   failedToSecure,
 }
 
 const securedDisplayStyleColorMap = {
   [SecuredDisplayStyle.securing]: colors.white,
+  [SecuredDisplayStyle.unsecuring]: colors.white,
   [SecuredDisplayStyle.secured]: colors.green,
   [SecuredDisplayStyle.blocked]: colors.green,
   [SecuredDisplayStyle.unsecured]: colors.red,
@@ -20,6 +22,8 @@ const securedDisplayStyleColorMap = {
 };
 
 const StyledSecuredLabel = styled.span((props: { displayStyle: SecuredDisplayStyle }) => ({
+  display: 'inline-block',
+  minHeight: '22px',
   color: securedDisplayStyleColorMap[props.displayStyle],
 }));
 
@@ -50,6 +54,9 @@ function getLabelText(displayStyle: SecuredDisplayStyle) {
     case SecuredDisplayStyle.unsecured:
       return messages.gettext('UNSECURE CONNECTION');
 
+    case SecuredDisplayStyle.unsecuring:
+      return '';
+
     case SecuredDisplayStyle.failedToSecure:
       return messages.gettext('FAILED TO SECURE CONNECTION');
   }
diff --git a/gui/src/renderer/components/TunnelControl.tsx b/gui/src/renderer/components/TunnelControl.tsx
index 016c360e4a..de3906c012 100644
--- a/gui/src/renderer/components/TunnelControl.tsx
+++ b/gui/src/renderer/components/TunnelControl.tsx
@@ -163,7 +163,7 @@ export default class TunnelControl extends React.Component<ITunnelControlProps>
         return (
           <Wrapper>
             <Body>
-              <Secured displayStyle={SecuredDisplayStyle.secured} />
+              <Secured displayStyle={SecuredDisplayStyle.unsecuring} />
               <Location>{this.renderCountry()}</Location>
             </Body>
             <Footer>
diff --git a/gui/src/shared/ipc-schema.ts b/gui/src/shared/ipc-schema.ts
index 93f6b66b92..a20213f424 100644
--- a/gui/src/shared/ipc-schema.ts
+++ b/gui/src/shared/ipc-schema.ts
@@ -46,7 +46,6 @@ export interface IAppStateSnapshot {
   accountHistory?: AccountToken;
   tunnelState: TunnelState;
   settings: ISettings;
-  location?: ILocation;
   relayListPair: IRelayListPair;
   currentVersion: ICurrentAppVersionInfo;
   upgradeVersion: IAppVersionInfo;
@@ -112,9 +111,6 @@ export const ipcSchema = {
     connected: notifyRenderer<void>(),
     disconnected: notifyRenderer<void>(),
   },
-  location: {
-    '': notifyRenderer<ILocation>(),
-  },
   relays: {
     '': notifyRenderer<IRelayListPair>(),
   },
@@ -129,6 +125,9 @@ export const ipcSchema = {
     openUrl: invoke<string, void>(),
     showOpenDialog: invoke<Electron.OpenDialogOptions, Electron.OpenDialogReturnValue>(),
   },
+  location: {
+    get: invoke<void, ILocation>(),
+  },
   tunnel: {
     '': notifyRenderer<TunnelState>(),
     connect: invoke<void, void>(),