Make relay location broadcast optional

6 files changed, 10 insertions, 14 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/TunnelState.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/TunnelState.kt
index 2551cfe978..15a698c6f4 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/TunnelState.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/TunnelState.kt
@@ -3,7 +3,7 @@ package net.mullvad.mullvadvpn.model
 sealed class TunnelState() {
     class Disconnected() : TunnelState()
     class Connecting(val location: GeoIpLocation?) : TunnelState()
-    class Connected(val location: GeoIpLocation) : TunnelState()
+    class Connected(val location: GeoIpLocation?) : TunnelState()
     class Disconnecting() : TunnelState()
     class Blocked() : TunnelState()
 }
diff --git a/gui/src/main/daemon-rpc.ts b/gui/src/main/daemon-rpc.ts
index 2c467a0d05..91f9e8771c 100644
--- a/gui/src/main/daemon-rpc.ts
+++ b/gui/src/main/daemon-rpc.ts
@@ -250,7 +250,7 @@ const tunnelStateSchema = oneOf(
           }),
         ),
       }),
-      location: locationSchema,
+      location: maybe(locationSchema),
     }),
   }),
   object({
diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index 9013c7c5cd..bdc1c4f5ec 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -729,7 +729,7 @@ class ApplicationMain {
     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) {
+      if (tunnelState.details && tunnelState.details.location) {
         this.setLocation(tunnelState.details.location);
       }
     } else if (tunnelState.state === 'disconnected') {
diff --git a/gui/src/shared/daemon-rpc-types.ts b/gui/src/shared/daemon-rpc-types.ts
index ec52d0ad5f..5d2c77d223 100644
--- a/gui/src/shared/daemon-rpc-types.ts
+++ b/gui/src/shared/daemon-rpc-types.ts
@@ -77,7 +77,7 @@ export type DaemonEvent =
 
 export interface ITunnelStateRelayInfo {
   endpoint: ITunnelEndpoint;
-  location: ILocation;
+  location?: ILocation;
 }
 
 export type TunnelState =
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 0e66149915..0394a0fa54 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -458,15 +458,11 @@ where
             TunnelStateTransition::Disconnected => TunnelState::Disconnected,
             TunnelStateTransition::Connecting(endpoint) => TunnelState::Connecting {
                 endpoint,
-                location: self
-                    .build_location_from_relay()
-                    .expect("No relay to get location from"),
+                location: self.build_location_from_relay(),
             },
             TunnelStateTransition::Connected(endpoint) => TunnelState::Connected {
                 endpoint,
-                location: self
-                    .build_location_from_relay()
-                    .expect("No relay to get location from"),
+                location: self.build_location_from_relay(),
             },
             TunnelStateTransition::Disconnecting(after_disconnect) => {
                 TunnelState::Disconnecting(after_disconnect)
@@ -793,7 +789,7 @@ where
         let get_location: Box<dyn Future<Item = Option<GeoIpLocation>, Error = ()> + Send> =
             match &self.tunnel_state {
                 Disconnected => Box::new(self.get_geo_location().map(Some)),
-                Connecting { location, .. } => Box::new(future::result(Ok(Some(location.clone())))),
+                Connecting { location, .. } => Box::new(future::result(Ok(location.clone()))),
                 Disconnecting(..) => match self.build_location_from_relay() {
                     Some(relay_location) => Box::new(future::result(Ok(Some(relay_location)))),
                     // Custom relay is set, no location is known
@@ -806,7 +802,7 @@ where
                             .map(|fetched_location| GeoIpLocation {
                                 ipv4: fetched_location.ipv4,
                                 ipv6: fetched_location.ipv6,
-                                ..relay_location
+                                ..relay_location.unwrap_or(fetched_location)
                             })
                             .map(Some),
                     )
diff --git a/mullvad-types/src/states.rs b/mullvad-types/src/states.rs
index 91618025ac..9c3ffa42a9 100644
--- a/mullvad-types/src/states.rs
+++ b/mullvad-types/src/states.rs
@@ -23,11 +23,11 @@ pub enum TunnelState {
     Disconnected,
     Connecting {
         endpoint: TunnelEndpoint,
-        location: GeoIpLocation,
+        location: Option<GeoIpLocation>,
     },
     Connected {
         endpoint: TunnelEndpoint,
-        location: GeoIpLocation,
+        location: Option<GeoIpLocation>,
     },
     Disconnecting(ActionAfterDisconnect),
     Blocked(BlockReason),