Further improve get_current_location RPC

6 files changed, 55 insertions, 39 deletions

diff --git a/gui/packages/desktop/src/main/daemon-rpc.js b/gui/packages/desktop/src/main/daemon-rpc.js
index bb55e5d335..1716190351 100644
--- a/gui/packages/desktop/src/main/daemon-rpc.js
+++ b/gui/packages/desktop/src/main/daemon-rpc.js
@@ -34,15 +34,17 @@ export type Location = {
   mullvadExitIp: boolean,
   hostname: ?string,
 };
-const LocationSchema = partialObject({
-  ip: maybe(string),
-  country: string,
-  city: maybe(string),
-  latitude: number,
-  longitude: number,
-  mullvad_exit_ip: boolean,
-  hostname: maybe(string),
-});
+const LocationSchema = maybe(
+  partialObject({
+    ip: maybe(string),
+    country: string,
+    city: maybe(string),
+    latitude: number,
+    longitude: number,
+    mullvad_exit_ip: boolean,
+    hostname: maybe(string),
+  }),
+);
 
 export type BlockReason =
   | {
@@ -413,7 +415,7 @@ export interface DaemonRpcProtocol {
   setAutoConnect(boolean): Promise<void>;
   connectTunnel(): Promise<void>;
   disconnectTunnel(): Promise<void>;
-  getLocation(): Promise<Location>;
+  getLocation(): Promise<?Location>;
   getState(): Promise<TunnelStateTransition>;
   getSettings(): Promise<Settings>;
   subscribeStateListener(listener: SubscriptionListener<TunnelStateTransition>): Promise<void>;
@@ -532,10 +534,13 @@ export class DaemonRpc implements DaemonRpcProtocol {
     await this._transport.send('disconnect');
   }
 
-  async getLocation(): Promise<Location> {
+  async getLocation(): Promise<?Location> {
     const response = await this._transport.send('get_current_location', [], NETWORK_CALL_TIMEOUT);
     try {
-      return camelCaseObjectKeys(validate(LocationSchema, response));
+      const validatedResponse = validate(LocationSchema, response);
+      if (validatedResponse) {
+        return camelCaseObjectKeys(validatedResponse);
+      }
     } catch (error) {
       throw new ResponseParseError('Invalid response from get_current_location', error);
     }
diff --git a/gui/packages/desktop/src/main/index.js b/gui/packages/desktop/src/main/index.js
index 3165729458..5da25f39de 100644
--- a/gui/packages/desktop/src/main/index.js
+++ b/gui/packages/desktop/src/main/index.js
@@ -660,6 +660,11 @@ const ApplicationMain = {
 
         // Fetch the new user location
         const location = await this._daemonRpc.getLocation();
+        // 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.
diff --git a/mullvad-cli/src/cmds/status.rs b/mullvad-cli/src/cmds/status.rs
index 26b2cf7f06..9c022a38d8 100644
--- a/mullvad-cli/src/cmds/status.rs
+++ b/mullvad-cli/src/cmds/status.rs
@@ -65,7 +65,13 @@ fn print_blocked_reason(reason: &BlockReason) {
 }
 
 fn print_location(rpc: &mut DaemonRpcClient) -> Result<()> {
-    let location = rpc.get_current_location()?;
+    let location = match rpc.get_current_location()? {
+        Some(loc) => loc,
+        None => {
+            println!("Location data unavailable");
+            return Ok(());
+        }
+    };
     let city_and_country = if let Some(city) = location.city {
         format!("{}, {}", city, location.country)
     } else {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index cf1a5f7443..d3ae39676a 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -456,28 +456,31 @@ impl Daemon {
         Self::oneshot_send(tx, self.tunnel_state.clone(), "current state");
     }
 
-    fn on_get_current_location(&self, tx: oneshot::Sender<GeoIpLocation>) {
+    fn on_get_current_location(&self, tx: oneshot::Sender<Option<GeoIpLocation>>) {
         use self::TunnelStateTransition::*;
-        let get_location: Box<dyn Future<Item = GeoIpLocation, Error = ()> + Send> =
+        let get_location: Box<dyn Future<Item = Option<GeoIpLocation>, Error = ()> + Send> =
             match self.tunnel_state {
-                Disconnected => Box::new(self.get_geo_location()),
-                Connecting(_) | Disconnecting(..) => {
-                    Box::new(future::result(Ok(self.build_location_from_relay())))
-                }
-                Connected(_) => {
-                    let location_from_relay = self.build_location_from_relay();
-                    Box::new(
+                Disconnected => Box::new(self.get_geo_location().map(Some)),
+                Connecting(_) | 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
+                    None => Box::new(future::result(Ok(None))),
+                },
+                Connected(_) => match self.build_location_from_relay() {
+                    Some(location_from_relay) => Box::new(
                         self.get_geo_location()
                             .map(|fetched_location| GeoIpLocation {
                                 ip: fetched_location.ip,
                                 ..location_from_relay
-                            }),
-                    )
-                }
+                            })
+                            .map(Some),
+                    ),
+                    // Custom relay is set, no location is known intrinsicly
+                    None => Box::new(self.get_geo_location().map(Some)),
+                },
                 Blocked(..) => {
-                    // We are not online at all at this stage. Return error.
-                    mem::drop(tx);
-                    return;
+                    // We are not online at all at this stage so no location data is available.
+                    Box::new(future::result(Ok(None)))
                 }
             };
 
@@ -494,15 +497,12 @@ impl Daemon {
         })
     }
 
-    fn build_location_from_relay(&self) -> GeoIpLocation {
-        let relay = self
-            .last_generated_relay
-            .as_ref()
-            .expect("Can't build location from relay in disconnected state");
+    fn build_location_from_relay(&self) -> Option<GeoIpLocation> {
+        let relay = self.last_generated_relay.as_ref()?;
         let location = relay.location.as_ref().cloned().unwrap();
         let hostname = relay.hostname.clone();
 
-        GeoIpLocation {
+        Some(GeoIpLocation {
             ip: None,
             country: location.country,
             city: Some(location.city),
@@ -510,7 +510,7 @@ impl Daemon {
             longitude: location.longitude,
             mullvad_exit_ip: true,
             hostname: Some(hostname),
-        }
+        })
     }
 
     fn on_get_account_data(
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 59e29381a2..c68c8ff967 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -95,7 +95,7 @@ build_rpc_trait! {
         /// Performs a geoIP lookup and returns the current location as perceived by the public
         /// internet.
         #[rpc(meta, name = "get_current_location")]
-        fn get_current_location(&self, Self::Metadata) -> BoxFuture<GeoIpLocation, Error>;
+        fn get_current_location(&self, Self::Metadata) -> BoxFuture<Option<GeoIpLocation>, Error>;
 
         /// Makes the daemon exit its main loop and quit.
         #[rpc(meta, name = "shutdown")]
@@ -168,7 +168,7 @@ pub enum ManagementCommand {
     /// Request the current state.
     GetState(OneshotSender<TunnelStateTransition>),
     /// Get the current geographical location.
-    GetCurrentLocation(OneshotSender<GeoIpLocation>),
+    GetCurrentLocation(OneshotSender<Option<GeoIpLocation>>),
     /// Request the metadata for an account.
     GetAccountData(
         OneshotSender<BoxFuture<AccountData, mullvad_rpc::Error>>,
@@ -529,7 +529,7 @@ impl<T: From<ManagementCommand> + 'static + Send> ManagementInterfaceApi
         Box::new(future)
     }
 
-    fn get_current_location(&self, _: Self::Metadata) -> BoxFuture<GeoIpLocation, Error> {
+    fn get_current_location(&self, _: Self::Metadata) -> BoxFuture<Option<GeoIpLocation>, Error> {
         log::debug!("get_current_location");
         let (tx, rx) = sync::oneshot::channel();
         let future = self
diff --git a/mullvad-ipc-client/src/lib.rs b/mullvad-ipc-client/src/lib.rs
index 066bd8250d..99f9ff6815 100644
--- a/mullvad-ipc-client/src/lib.rs
+++ b/mullvad-ipc-client/src/lib.rs
@@ -148,7 +148,7 @@ impl DaemonRpcClient {
         self.call("get_auto_connect", &NO_ARGS)
     }
 
-    pub fn get_current_location(&mut self) -> Result<GeoIpLocation> {
+    pub fn get_current_location(&mut self) -> Result<Option<GeoIpLocation>> {
         self.call("get_current_location", &NO_ARGS)
     }