From ac2262b640316a6b5797bbc59bb05fc318aca702 Mon Sep 17 00:00:00 2001
From: Sebastian Holmin <sebastian.holmin@mullvad.net>
Date: Thu, 7 Dec 2023 11:13:02 +0100
Subject: [PATCH] Move Rest and availability handle into GeoIpHandler

---
 mullvad-daemon/src/geoip.rs | 12 ++++++++----
 mullvad-daemon/src/lib.rs   | 17 ++++++++---------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/mullvad-daemon/src/geoip.rs b/mullvad-daemon/src/geoip.rs
index 9f5b3854eb98..ec1d989dc642 100644
--- a/mullvad-daemon/src/geoip.rs
+++ b/mullvad-daemon/src/geoip.rs
@@ -49,27 +49,29 @@ static MULLVAD_CONNCHECK_HOST: Lazy<String> = Lazy::new(|| {
 const LOCATION_RETRY_STRATEGY: Jittered<ExponentialBackoff> =
     Jittered::jitter(ExponentialBackoff::new(Duration::from_secs(1), 4));
 
-/// Handler for request to am.i.mullvad.net
+/// Handler for request to am.i.mullvad.net, manages in-flight request and validity of responses.
 pub struct GeoIpHandler {
     /// Unique ID for each request, used to verify
     pub request_id: usize,
     /// Handle to abort any in-flight request.
     request_handle: Option<JoinHandle<()>>,
+    availability: ApiAvailabilityHandle,
+    rest_service: RequestServiceHandle,
 }
 
 impl GeoIpHandler {
-    pub fn new() -> Self {
+    pub fn new(availability: ApiAvailabilityHandle, rest_service: RequestServiceHandle) -> Self {
         Self {
             request_id: 0,
             request_handle: None,
+            availability,
+            rest_service,
         }
     }
 
     pub fn send_geo_location_request(
         &mut self,
         use_ipv6: bool,
-        availability: ApiAvailabilityHandle,
-        rest_service: RequestServiceHandle,
         daemon_event_sender: DaemonEventSender,
     ) {
         // Increment request ID
@@ -77,6 +79,8 @@ impl GeoIpHandler {
         let request_id_copy = self.request_id;
 
         self.abort_current_request();
+        let rest_service = self.rest_service.clone();
+        let availability = self.availability.clone();
         self.request_handle = Some(tokio::spawn(async move {
             if let Ok(merged_location) =
                 get_geo_location_with_retry(use_ipv6, availability, rest_service).await
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 610b196904cf..464994fdafda 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -832,6 +832,11 @@ where
         // Attempt to download a fresh relay list
         relay_list_updater.update().await;
 
+        let location_handler = GeoIpHandler::new(
+            api_handle.availability.clone(),
+            api_runtime.rest_handle().await,
+        );
+
         let daemon = Daemon {
             tunnel_state: TunnelState::Disconnected(None),
             target_state,
@@ -859,7 +864,7 @@ where
             tunnel_state_machine_handle,
             #[cfg(target_os = "windows")]
             volume_update_tx,
-            location_handler: GeoIpHandler::new(),
+            location_handler,
         };
 
         api_availability.unsuspend();
@@ -1050,15 +1055,9 @@ where
         }
 
         let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
-        let rest_service = self.api_runtime.rest_handle().await;
-        let availability = self.api_handle.availability.clone();
         let location_sender = self.tx.clone().to_specialized_sender();
-        self.location_handler.send_geo_location_request(
-            use_ipv6,
-            availability,
-            rest_service,
-            location_sender,
-        );
+        self.location_handler
+            .send_geo_location_request(use_ipv6, location_sender);
     }
 
     async fn handle_location_event(&mut self, request_id: usize, fetched_location: GeoIpLocation) {