Merge remote-tracking branch 'origin/out-ip-sometimes-missing-118' into main

3 files changed, 63 insertions, 38 deletions

diff --git a/mullvad-api/src/lib.rs b/mullvad-api/src/lib.rs
index 5a436fe1bf..c0024b22ee 100644
--- a/mullvad-api/src/lib.rs
+++ b/mullvad-api/src/lib.rs
@@ -352,7 +352,7 @@ impl Runtime {
     }
 
     /// Returns a new request service handle
-    pub async fn rest_handle(&mut self) -> rest::RequestServiceHandle {
+    pub async fn rest_handle(&self) -> rest::RequestServiceHandle {
         self.new_request_service(
             None,
             ApiConnectionMode::Direct.into_repeat(),
diff --git a/mullvad-daemon/src/geoip.rs b/mullvad-daemon/src/geoip.rs
index 43ffaf1054..02ee09d1c0 100644
--- a/mullvad-daemon/src/geoip.rs
+++ b/mullvad-daemon/src/geoip.rs
@@ -1,12 +1,19 @@
+use std::time::Duration;
+
 use futures::join;
 use mullvad_api::{
     self,
+    availability::ApiAvailabilityHandle,
     rest::{Error, RequestServiceHandle},
 };
 use mullvad_types::location::{AmIMullvad, GeoIpLocation};
 use once_cell::sync::Lazy;
+use talpid_core::future_retry::{retry_future, ConstantInterval};
 use talpid_types::ErrorExt;
 
+/// Retry interval for fetching location
+const RETRY_LOCATION_STRATEGY: ConstantInterval = ConstantInterval::new(Duration::ZERO, Some(3));
+
 // Define the Mullvad connection checking api endpoint.
 //
 // In a development build the host name for the connection checking endpoint can
@@ -36,7 +43,32 @@ static MULLVAD_CONNCHECK_HOST: Lazy<String> = Lazy::new(|| {
     host.to_string()
 });
 
-pub async fn send_location_request(
+/// Fetch the current `GeoIpLocation` with retrys
+pub async fn get_geo_location(
+    rest_service: RequestServiceHandle,
+    use_ipv6: bool,
+    api_handle: ApiAvailabilityHandle,
+) -> Option<GeoIpLocation> {
+    log::debug!("Fetching GeoIpLocation");
+    match retry_future(
+        move || send_location_request(rest_service.clone(), use_ipv6),
+        move |result| match result {
+            Err(error) if error.is_network_error() => !api_handle.get_state().is_offline(),
+            _ => false,
+        },
+        RETRY_LOCATION_STRATEGY,
+    )
+    .await
+    {
+        Ok(loc) => Some(loc),
+        Err(e) => {
+            log::warn!("Unable to fetch GeoIP location: {}", e.display_chain());
+            None
+        }
+    }
+}
+
+async fn send_location_request(
     request_sender: RequestServiceHandle,
     use_ipv6: bool,
 ) -> Result<GeoIpLocation, Error> {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index a540791838..65f2c56654 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -27,7 +27,7 @@ mod tunnel;
 pub mod version;
 mod version_check;
 
-use crate::target_state::PersistentTargetState;
+use crate::{geoip::get_geo_location, target_state::PersistentTargetState};
 use device::{AccountEvent, PrivateAccountAndDevice, PrivateDeviceEvent};
 use futures::{
     channel::{mpsc, oneshot},
@@ -1346,12 +1346,7 @@ where
         use self::TunnelState::*;
 
         match &self.tunnel_state {
-            Disconnected => {
-                let location = self.get_geo_location().await;
-                tokio::spawn(async {
-                    Self::oneshot_send(tx, location.await.ok(), "current location");
-                });
-            }
+            Disconnected => self.update_and_send_geo_location(tx, None).await,
             Connecting { location, .. } => {
                 Self::oneshot_send(tx, location.clone(), "current location")
             }
@@ -1361,38 +1356,36 @@ where
                 "current location",
             ),
             Connected { location, .. } => {
-                let relay_location = location.clone();
-                let location_future = self.get_geo_location().await;
-                tokio::spawn(async {
-                    let location = location_future.await;
-                    Self::oneshot_send(
-                        tx,
-                        location.ok().map(|fetched_location| GeoIpLocation {
-                            ipv4: fetched_location.ipv4,
-                            ipv6: fetched_location.ipv6,
-                            ..relay_location.unwrap_or(fetched_location)
-                        }),
-                        "current location",
-                    );
-                });
-            }
-            Error(_) => {
-                // We are not online at all at this stage so no location data is available.
-                Self::oneshot_send(tx, None, "current location");
+                self.update_and_send_geo_location(tx, location.clone())
+                    .await
             }
-        }
+            // We are not online at all at this stage so no location data is available.
+            Error(_) => Self::oneshot_send(tx, None, "current location"),
+        };
     }
 
-    async fn get_geo_location(&mut self) -> impl Future<Output = Result<GeoIpLocation, ()>> {
+    /// Fetch the current `GeoIpLocation` and send it on the return channel,
+    /// in a non-blocking fashion. Optionally give a chached previous location.
+    async fn update_and_send_geo_location(
+        &mut self,
+        tx: oneshot::Sender<Option<GeoIpLocation>>,
+        current_relay_location: Option<GeoIpLocation>,
+    ) {
         let rest_service = self.api_runtime.rest_handle().await;
         let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
-        async move {
-            geoip::send_location_request(rest_service, use_ipv6)
-                .await
-                .map_err(|e| {
-                    log::warn!("Unable to fetch GeoIP location: {}", e.display_chain());
-                })
-        }
+        let api_handle = self.api_handle.availability.clone();
+        tokio::spawn(async move {
+            let new_location = get_geo_location(rest_service, use_ipv6, api_handle).await;
+            Self::oneshot_send(
+                tx,
+                new_location.map(|fetched_location| GeoIpLocation {
+                    ipv4: fetched_location.ipv4,
+                    ipv6: fetched_location.ipv6,
+                    ..current_relay_location.unwrap_or(fetched_location)
+                }),
+                "current location",
+            );
+        });
     }
 
     fn on_create_new_account(&mut self, tx: ResponseTx<String, Error>) {
@@ -2588,8 +2581,8 @@ fn new_selector_config(settings: &Settings) -> SelectorConfig {
     }
 }
 
-/// Consume a oneshot sender of `T1` and return a sender that takes a different type `T2`. `forwarder` should map `T1` back to `T2` and
-/// send the result back to the original receiver.
+/// Consume a oneshot sender of `T1` and return a sender that takes a different type `T2`.
+/// `forwarder` should map `T1` back to `T2` and send the result back to the original receiver.
 fn oneshot_map<T1: Send + 'static, T2: Send + 'static>(
     tx: oneshot::Sender<T1>,
     forwarder: impl Fn(oneshot::Sender<T1>, T2) + Send + 'static,