3 files changed, 45 insertions, 39 deletions

diff --git a/gui/src/main/daemon-rpc.ts b/gui/src/main/daemon-rpc.ts
index 215807b2a7..2c467a0d05 100644
--- a/gui/src/main/daemon-rpc.ts
+++ b/gui/src/main/daemon-rpc.ts
@@ -8,7 +8,7 @@ import {
   IRelayList,
   ISettings,
   RelaySettingsUpdate,
-  TunnelStateTransition,
+  TunnelState,
 } from '../shared/daemon-rpc-types';
 import { CommunicationError, InvalidAccountError, NoDaemonError } from './errors';
 import JsonRpcClient, {
@@ -230,24 +230,27 @@ const accountDataSchema = partialObject({
   expiry: string,
 });
 
-const tunnelStateTransitionSchema = oneOf(
+const tunnelStateSchema = oneOf(
   object({
     state: enumeration('disconnecting'),
     details: enumeration('nothing', 'block', 'reconnect'),
   }),
   object({
     state: enumeration('connecting', 'connected'),
-    details: partialObject({
-      address: string,
-      protocol: enumeration('tcp', 'udp'),
-      tunnel_type: enumeration('wireguard', 'openvpn'),
-      proxy: maybe(
-        partialObject({
-          address: string,
-          protocol: enumeration('tcp', 'udp'),
-          proxy_type: enumeration('shadowsocks', 'custom'),
-        }),
-      ),
+    details: object({
+      endpoint: partialObject({
+        address: string,
+        protocol: enumeration('tcp', 'udp'),
+        tunnel_type: enumeration('wireguard', 'openvpn'),
+        proxy: maybe(
+          partialObject({
+            address: string,
+            protocol: enumeration('tcp', 'udp'),
+            proxy_type: enumeration('shadowsocks', 'custom'),
+          }),
+        ),
+      }),
+      location: locationSchema,
     }),
   }),
   object({
@@ -329,7 +332,7 @@ const settingsSchema = partialObject({
 
 const daemonEventSchema = oneOf(
   object({
-    state_transition: tunnelStateTransitionSchema,
+    tunnel_state: tunnelStateSchema,
   }),
   object({
     settings: settingsSchema,
@@ -468,12 +471,10 @@ export class DaemonRpc {
     }
   }
 
-  public async getState(): Promise<TunnelStateTransition> {
+  public async getState(): Promise<TunnelState> {
     const response = await this.transport.send('get_state');
     try {
-      return camelCaseObjectKeys(
-        validate(tunnelStateTransitionSchema, response),
-      ) as TunnelStateTransition;
+      return camelCaseObjectKeys(validate(tunnelStateSchema, response)) as TunnelState;
     } catch (error) {
       throw new ResponseParseError('Invalid response from get_state', error);
     }
diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index b34cc87fbc..9013c7c5cd 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -15,7 +15,7 @@ import {
   ISettings,
   RelaySettings,
   RelaySettingsUpdate,
-  TunnelStateTransition,
+  TunnelState,
 } from '../shared/daemon-rpc-types';
 import { loadTranslations, messages } from '../shared/gettext';
 import { IpcMainEventChannel } from '../shared/ipc-event-channel';
@@ -73,7 +73,7 @@ class ApplicationMain {
   private quitStage = AppQuitStage.unready;
 
   private accountHistory: AccountToken[] = [];
-  private tunnelState: TunnelStateTransition = { state: 'disconnected' };
+  private tunnelState: TunnelState = { state: 'disconnected' };
   private settings: ISettings = {
     accountToken: undefined,
     allowLan: false,
@@ -480,8 +480,8 @@ class ApplicationMain {
   private async subscribeEvents(): Promise<void> {
     const daemonEventListener = new SubscriptionListener(
       (daemonEvent: DaemonEvent) => {
-        if ('stateTransition' in daemonEvent) {
-          this.setTunnelState(daemonEvent.stateTransition);
+        if ('tunnelState' in daemonEvent) {
+          this.setTunnelState(daemonEvent.tunnelState);
         } else if ('settings' in daemonEvent) {
           this.setSettings(daemonEvent.settings);
         } else if ('relayList' in daemonEvent) {
@@ -512,7 +512,7 @@ class ApplicationMain {
     }
   }
 
-  private setTunnelState(newState: TunnelStateTransition) {
+  private setTunnelState(newState: TunnelState) {
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
     this.updateLocation();
@@ -724,16 +724,24 @@ class ApplicationMain {
   }
 
   private async updateLocation() {
-    const state = this.tunnelState.state;
+    const tunnelState = this.tunnelState;
 
-    if (state === 'connected' || state === 'disconnected' || state === 'connecting') {
-      try {
-        // It may take some time to fetch the new user location.
-        // So take the user to the last known location when disconnected.
-        if (state === 'disconnected' && this.lastDisconnectedLocation) {
-          this.setLocation(this.lastDisconnectedLocation);
-        }
+    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) {
+        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 location = await this.daemonRpc.getLocation();
         // If the location is currently unavailable, do nothing! This only ever
@@ -751,7 +759,7 @@ class ApplicationMain {
         // Broadcast the new location.
         // There is a chance that the location is not stale if the tunnel state before the location
         // request is the same as after receiving the response.
-        if (this.tunnelState.state === state) {
+        if (this.tunnelState.state === tunnelState.state) {
           this.setLocation(location);
         }
       } catch (error) {
@@ -760,10 +768,7 @@ class ApplicationMain {
     }
   }
 
-  private trayIconType(
-    tunnelState: TunnelStateTransition,
-    blockWhenDisconnected: boolean,
-  ): TrayIconType {
+  private trayIconType(tunnelState: TunnelState, blockWhenDisconnected: boolean): TrayIconType {
     switch (tunnelState.state) {
       case 'connected':
         return 'secured';
@@ -791,7 +796,7 @@ class ApplicationMain {
     }
   }
 
-  private updateTrayIcon(tunnelState: TunnelStateTransition, blockWhenDisconnected: boolean) {
+  private updateTrayIcon(tunnelState: TunnelState, blockWhenDisconnected: boolean) {
     const type = this.trayIconType(tunnelState, blockWhenDisconnected);
 
     if (this.trayIconController) {
diff --git a/gui/src/main/notification-controller.ts b/gui/src/main/notification-controller.ts
index a2e1e00a0b..771e788722 100644
--- a/gui/src/main/notification-controller.ts
+++ b/gui/src/main/notification-controller.ts
@@ -2,7 +2,7 @@ import { app, nativeImage, NativeImage, Notification, shell } from 'electron';
 import path from 'path';
 import { sprintf } from 'sprintf-js';
 import config from '../config.json';
-import { TunnelStateTransition } from '../shared/daemon-rpc-types';
+import { TunnelState } from '../shared/daemon-rpc-types';
 import { messages } from '../shared/gettext';
 
 export default class NotificationController {
@@ -22,7 +22,7 @@ export default class NotificationController {
     }
   }
 
-  public notifyTunnelState(tunnelState: TunnelStateTransition) {
+  public notifyTunnelState(tunnelState: TunnelState) {
     switch (tunnelState.state) {
       case 'connecting':
         if (!this.reconnecting) {