Include known exit IP address in current location

4 files changed, 61 insertions, 33 deletions

diff --git a/gui/packages/desktop/src/renderer/lib/daemon-rpc.js b/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
index c4ac506f36..594a33fdf6 100644
--- a/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
+++ b/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
@@ -25,6 +25,7 @@ export type AccountData = { expiry: string };
 export type AccountToken = string;
 export type Ip = string;
 export type Location = {
+  ip: ?string,
   country: string,
   city: ?string,
   latitude: number,
@@ -33,6 +34,7 @@ export type Location = {
   hostname: ?string,
 };
 const LocationSchema = object({
+  ip: maybe(string),
   country: string,
   city: maybe(string),
   latitude: number,
diff --git a/mullvad-daemon/src/geoip.rs b/mullvad-daemon/src/geoip.rs
index 01f42ec6c9..2faed344a6 100644
--- a/mullvad-daemon/src/geoip.rs
+++ b/mullvad-daemon/src/geoip.rs
@@ -21,13 +21,13 @@ error_chain! {
 
 pub fn send_location_request(
     request_sender: mullvad_rpc::rest::RequestSender,
-) -> Box<Future<Item = GeoIpLocation, Error = Error>> {
+) -> impl Future<Item = GeoIpLocation, Error = Error> {
     let (response_tx, response_rx) = futures::sync::oneshot::channel();
     let request = mullvad_rpc::rest::create_get_request(URI.parse().unwrap());
-    let future = futures::Sink::send(request_sender.clone(), (request, response_tx))
+
+    futures::Sink::send(request_sender, (request, response_tx))
         .map_err(|e| Error::with_chain(e, ErrorKind::NoResponse))
         .and_then(|_| response_rx.map_err(|e| Error::with_chain(e, ErrorKind::NoResponse)))
         .and_then(|response_result| response_result.map_err(Error::from))
-        .and_then(|response| serde_json::from_slice(&response).map_err(Error::from));
-    Box::new(future)
+        .and_then(|response| serde_json::from_slice(&response).map_err(Error::from))
 }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index a65a34dd8e..e88c4cd0e1 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -45,6 +45,7 @@ mod rpc_uniqueness_check;
 
 use error_chain::ChainedError;
 use futures::{
+    future,
     sync::{mpsc::UnboundedSender, oneshot},
     Future, Sink,
 };
@@ -444,36 +445,58 @@ impl Daemon {
 
     fn on_get_current_location(&self, tx: oneshot::Sender<GeoIpLocation>) {
         use self::TunnelStateTransition::*;
-        match self.tunnel_state {
-            Disconnected => {
-                let https_handle = self.https_handle.clone();
-                self.tokio_remote.spawn(move |_| {
-                    geoip::send_location_request(https_handle)
-                        .map(move |location| Self::oneshot_send(tx, location, "current location"))
-                        .map_err(|e| {
-                            warn!("Unable to fetch GeoIP location: {}", e.display_chain());
-                        })
-                });
-            }
-            Connecting(_) | Connected(_) | Disconnecting(..) => {
-                if let Some(ref relay) = self.last_generated_relay {
-                    let location = relay.location.as_ref().cloned().unwrap();
-                    let hostname = relay.hostname.clone();
-                    let geo_ip_location = GeoIpLocation {
-                        country: location.country,
-                        city: Some(location.city),
-                        latitude: location.latitude,
-                        longitude: location.longitude,
-                        mullvad_exit_ip: true,
-                        hostname: Some(hostname),
-                    };
-                    Self::oneshot_send(tx, geo_ip_location, "current location");
+        let get_location: Box<dyn Future<Item = 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())))
                 }
-            }
-            Blocked(..) => {
-                // We are not online at all at this stage. Return error.
-                mem::drop(tx);
-            }
+                Connected(_) => {
+                    let location_from_relay = self.build_location_from_relay();
+                    Box::new(
+                        self.get_geo_location()
+                            .map(|fetched_location| GeoIpLocation {
+                                ip: fetched_location.ip,
+                                ..location_from_relay
+                            }),
+                    )
+                }
+                Blocked(..) => {
+                    // We are not online at all at this stage. Return error.
+                    mem::drop(tx);
+                    return;
+                }
+            };
+
+        self.tokio_remote.spawn(move |_| {
+            get_location.map(|location| Self::oneshot_send(tx, location, "current location"))
+        });
+    }
+
+    fn get_geo_location(&self) -> impl Future<Item = GeoIpLocation, Error = ()> {
+        let https_handle = self.https_handle.clone();
+
+        geoip::send_location_request(https_handle).map_err(|e| {
+            warn!("Unable to fetch GeoIP location: {}", e.display_chain());
+        })
+    }
+
+    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");
+        let location = relay.location.as_ref().cloned().unwrap();
+        let hostname = relay.hostname.clone();
+
+        GeoIpLocation {
+            ip: None,
+            country: location.country,
+            city: Some(location.city),
+            latitude: location.latitude,
+            longitude: location.longitude,
+            mullvad_exit_ip: true,
+            hostname: Some(hostname),
         }
     }
 
diff --git a/mullvad-types/src/location.rs b/mullvad-types/src/location.rs
index 57d1e31269..2c420fafc5 100644
--- a/mullvad-types/src/location.rs
+++ b/mullvad-types/src/location.rs
@@ -1,3 +1,5 @@
+use std::net::IpAddr;
+
 pub type CountryCode = String;
 pub type CityCode = String;
 pub type Hostname = String;
@@ -14,6 +16,7 @@ pub struct Location {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct GeoIpLocation {
+    pub ip: Option<IpAddr>,
     pub country: String,
     pub city: Option<String>,
     pub latitude: f64,