Move all location handling to renderer

3 files changed, 109 insertions, 116 deletions

diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index 03121170e7..bdaf33a7c8 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,
@@ -179,10 +178,7 @@ class ApplicationMain {
     },
   };
   private guiSettings = new GuiSettings();
-  private location?: ILocation;
-  private lastDisconnectedLocation?: ILocation;
   private tunnelStateExpectation?: Expectation;
-  private getLocationPromise?: Promise<ILocation>;
 
   private relays: IRelayList = { countries: [] };
 
@@ -732,7 +728,6 @@ class ApplicationMain {
   private setTunnelState(newState: TunnelState) {
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
-    void this.updateLocation();
 
     if (process.platform === 'linux') {
       this.tray?.setContextMenu(this.createTrayContextMenu());
@@ -799,14 +794,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,39 +968,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);
-      }
-    }
-
-    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. Save and broadcast the new location if
-        // it is the result of the most recent call to getLocation.
-        if (location && getLocationPromise === this.getLocationPromise) {
-          this.setLocation(location);
-
-          // Cache the user location
-          if (tunnelState.state === 'disconnected') {
-            this.lastDisconnectedLocation = location;
-          }
-        }
-      } catch (error) {
-        log.error(`Failed to update the location: ${error.message}`);
-      }
-    }
-  }
-
   private trayIconType(tunnelState: TunnelState, blockWhenDisconnected: boolean): TrayIconType {
     switch (tunnelState.state) {
       case 'connected':
@@ -1091,7 +1045,6 @@ class ApplicationMain {
       accountHistory: this.accountHistory,
       tunnelState: this.tunnelState,
       settings: this.settings,
-      location: this.location,
       relayListPair: {
         relays: this.processRelaysForPresentation(
           this.relays,
@@ -1147,18 +1100,10 @@ class ApplicationMain {
       return this.setAutoStart(autoStart);
     });
 
+    IpcMainEventChannel.location.handleGet(() => this.daemonRpc.getLocation());
+
     IpcMainEventChannel.tunnel.handleConnect(() => this.daemonRpc.connectTunnel());
-    IpcMainEventChannel.tunnel.handleDisconnect(async () => {
-      // It may take some time to fetch the new user location.
-      // So take the user to the last known location when disconnected.
-      if (this.windowController && this.lastDisconnectedLocation) {
-        IpcMainEventChannel.location.notify(
-          this.windowController.webContents,
-          this.lastDisconnectedLocation,
-        );
-      }
-      return this.daemonRpc.disconnectTunnel();
-    });
+    IpcMainEventChannel.tunnel.handleDisconnect(() => this.daemonRpc.disconnectTunnel());
     IpcMainEventChannel.tunnel.handleReconnect(() => this.daemonRpc.reconnectTunnel());
 
     IpcMainEventChannel.guiSettings.handleSetEnableSystemNotifications((flag: boolean) => {
diff --git a/gui/src/renderer/app.tsx b/gui/src/renderer/app.tsx
index 4ee5dc51e7..41e39b0214 100644
--- a/gui/src/renderer/app.tsx
+++ b/gui/src/renderer/app.tsx
@@ -116,7 +116,8 @@ 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 settings!: ISettings;
@@ -124,6 +125,7 @@ export default class AppRenderer {
   private doingLogin = false;
   private loginScheduler = new Scheduler();
   private connectedToDaemon = false;
+  private getLocationPromise?: Promise<ILocation>;
 
   constructor() {
     log.addOutput(new ConsoleOutput(LogLevel.debug));
@@ -164,10 +166,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 +229,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 +248,8 @@ export default class AppRenderer {
       );
     }
 
+    void this.updateLocation();
+
     const navigationBase = this.getNavigationBase(
       initialState.isConnected,
       initialState.settings.accountToken,
@@ -332,7 +328,7 @@ export default class AppRenderer {
     // 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();
+      void this.updateLocation({ state: 'connecting' });
       this.reduxActions.connection.connecting();
 
       return IpcRendererEventChannel.tunnel.connect();
@@ -345,6 +341,7 @@ export default class AppRenderer {
     // 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
+      void this.updateLocation({ state: 'disconnecting', details: 'nothing' });
       this.reduxActions.connection.disconnecting('nothing');
 
       return IpcRendererEventChannel.tunnel.disconnect();
@@ -357,7 +354,7 @@ export default class AppRenderer {
     // 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();
+      void this.updateLocation({ state: 'connecting' });
       this.reduxActions.connection.connecting();
 
       return IpcRendererEventChannel.tunnel.reconnect();
@@ -605,47 +602,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 state = this.reduxStore.getState();
-        const coordinates = {
-          longitude: state.connection.longitude,
-          latitude: state.connection.latitude,
-        };
-        const relayLocations = state.settings.relayLocations;
-        if ('country' in constraint) {
-          const country = relayLocations.find(({ code }) => constraint.country === code);
-
-          this.reduxActions.connection.newLocation({ 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);
-
-          this.reduxActions.connection.newLocation({
-            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]);
-
-          this.reduxActions.connection.newLocation({
-            country: country?.name,
-            city: city?.name,
-            hostname: constraint.hostname[2],
-            ...coordinates,
-          });
-        }
-      }
-    }
-  }
-
   private setRelaySettings(relaySettings: RelaySettings) {
     const actions = this.reduxActions;
 
@@ -787,6 +743,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) {
@@ -851,7 +810,7 @@ export default class AppRenderer {
     this.doingLogin = false;
   }
 
-  private setLocation(location: ILocation) {
+  private setLocation(location: Partial<ILocation>) {
     this.location = location;
     this.propagateLocationToRedux();
   }
@@ -862,7 +821,7 @@ export default class AppRenderer {
     }
   }
 
-  private translateLocation(inputLocation: ILocation): ILocation {
+  private translateLocation(inputLocation: Partial<ILocation>): Partial<ILocation> {
     const location = { ...inputLocation };
 
     if (location.city) {
@@ -943,4 +902,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/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>(),