Set address change listener after TSM has spawned

5 files changed, 84 insertions, 116 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index cffef475d4..21ca9298b8 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -24,7 +24,7 @@ mod version_check;
 use futures::{
     channel::{mpsc, oneshot},
     future::{abortable, AbortHandle, Future},
-    SinkExt, StreamExt,
+    StreamExt,
 };
 use log::{debug, error, info, warn};
 use mullvad_rpc::availability::ApiAvailabilityHandle;
@@ -559,50 +559,6 @@ where
         let runtime = tokio::runtime::Handle::current();
 
         let (internal_event_tx, internal_event_rx) = command_channel.destructure();
-        let (address_change_tx, mut address_change_rx) = mpsc::channel(0);
-        let address_change_tx = std::sync::Mutex::new(address_change_tx);
-        let address_change_runtime = runtime.clone();
-
-        let mut rpc_runtime = mullvad_rpc::MullvadRpcRuntime::with_cache(
-            runtime.clone(),
-            Some(&resource_dir),
-            &cache_dir,
-            true,
-            move |address| {
-                let (result_tx, result_rx) = oneshot::channel();
-
-                let mut tx = address_change_tx.lock().unwrap().clone();
-                address_change_runtime.block_on(async move {
-                    let tunnel_command = TunnelCommand::AllowEndpoint(
-                        Endpoint::from_socket_address(address, TransportProtocol::Tcp),
-                        result_tx,
-                    );
-                    let _ = tx.send(tunnel_command).await;
-                    result_rx.await.map_err(|_| ())
-                })
-            },
-            #[cfg(target_os = "android")]
-            Self::create_bypass_tx(&internal_event_tx),
-        )
-        .await
-        .map_err(Error::InitRpcFactory)?;
-        let rpc_handle = rpc_runtime.mullvad_rest_handle();
-        let api_availability = rpc_runtime.availability_handle();
-        api_availability.suspend();
-
-        let relay_list_listener = event_listener.clone();
-        let on_relay_list_update = move |relay_list: &RelayList| {
-            relay_list_listener.notify_relay_list(relay_list.clone());
-        };
-
-        let relay_selector = relays::RelaySelector::new(
-            rpc_handle.clone(),
-            on_relay_list_update,
-            &resource_dir,
-            &cache_dir,
-            api_availability.clone(),
-        );
-
 
         if let Err(error) = migrations::migrate_all(&cache_dir, &settings_dir).await {
             log::error!(
@@ -616,21 +572,6 @@ where
             let _ = settings.set_show_beta_releases(true).await;
         }
 
-        let app_version_info = version_check::load_cache(&cache_dir).await;
-        let (version_updater, version_updater_handle) = version_check::VersionUpdater::new(
-            rpc_handle.clone(),
-            api_availability.clone(),
-            cache_dir.clone(),
-            internal_event_tx.to_specialized_sender(),
-            app_version_info.clone(),
-            settings.show_beta_releases,
-        );
-        tokio::spawn(version_updater.run());
-        let account_history =
-            account_history::AccountHistory::new(&settings_dir, settings.get_account_token())
-                .await
-                .map_err(Error::LoadAccountHistory)?;
-
         // Restore the tunnel to a previous state
         let target_cache = cache_dir.join(TARGET_START_STATE_FILE);
         let cached_target_state: Option<TargetState> =
@@ -677,10 +618,6 @@ where
         };
         Self::cache_target_state(&cache_dir, initial_target_state).await;
 
-        let initial_api_endpoint = Endpoint::from_socket_address(
-            rpc_runtime.address_cache.peek_address(),
-            TransportProtocol::Tcp,
-        );
         #[cfg(windows)]
         let exclude_paths = if settings.split_tunnel.enable_exclusions {
             settings
@@ -693,8 +630,26 @@ where
             vec![]
         };
 
-        let (offline_state_tx, offline_state_rx) = mpsc::unbounded();
+        let mut rpc_runtime = mullvad_rpc::MullvadRpcRuntime::with_cache(
+            runtime.clone(),
+            Some(&resource_dir),
+            &cache_dir,
+            true,
+            #[cfg(target_os = "android")]
+            Self::create_bypass_tx(&internal_event_tx),
+        )
+        .await
+        .map_err(Error::InitRpcFactory)?;
+
+        let api_availability = rpc_runtime.availability_handle();
+        api_availability.suspend();
 
+        let initial_api_endpoint = Endpoint::from_socket_address(
+            rpc_runtime.address_cache.peek_address(),
+            TransportProtocol::Tcp,
+        );
+
+        let (offline_state_tx, offline_state_rx) = mpsc::unbounded();
         let tunnel_command_tx = tunnel_state_machine::spawn(
             runtime.clone(),
             tunnel_state_machine::InitialTunnelState {
@@ -708,7 +663,7 @@ where
             },
             tunnel_parameters_generator,
             log_dir,
-            resource_dir,
+            resource_dir.clone(),
             cache_dir.clone(),
             internal_event_tx.to_specialized_sender(),
             offline_state_tx,
@@ -719,21 +674,56 @@ where
         .await
         .map_err(Error::TunnelError)?;
 
-        api_availability.unsuspend();
-
-        Self::forward_offline_state(&runtime, api_availability.clone(), offline_state_rx).await;
-
-        let tsm_api_address_change_tx = Arc::downgrade(&tunnel_command_tx);
-        tokio::spawn(async move {
-            while let Some(address_change) = address_change_rx.next().await {
-                if let Some(tx) = tsm_api_address_change_tx.upgrade() {
-                    let _ = tx.unbounded_send(address_change);
+        let address_change_runtime = runtime.clone();
+        let tunnel_cmd_weak_tx = Arc::downgrade(&tunnel_command_tx);
+        rpc_runtime.set_address_change_listener(move |address| {
+            let (result_tx, result_rx) = oneshot::channel();
+            let tx = tunnel_cmd_weak_tx.clone();
+            address_change_runtime.block_on(async move {
+                if let Some(tx) = tx.upgrade() {
+                    let _ = tx.unbounded_send(TunnelCommand::AllowEndpoint(
+                        Endpoint::from_socket_address(address, TransportProtocol::Tcp),
+                        result_tx,
+                    ));
+                    result_rx.await.map_err(|_| ())
                 } else {
-                    return;
+                    Err(())
                 }
-            }
+            })
         });
 
+        let rpc_handle = rpc_runtime.mullvad_rest_handle();
+
+        Self::forward_offline_state(&runtime, api_availability.clone(), offline_state_rx).await;
+
+        let relay_list_listener = event_listener.clone();
+        let on_relay_list_update = move |relay_list: &RelayList| {
+            relay_list_listener.notify_relay_list(relay_list.clone());
+        };
+
+        let relay_selector = relays::RelaySelector::new(
+            rpc_handle.clone(),
+            on_relay_list_update,
+            &resource_dir,
+            &cache_dir,
+            api_availability.clone(),
+        );
+
+        let app_version_info = version_check::load_cache(&cache_dir).await;
+        let (version_updater, version_updater_handle) = version_check::VersionUpdater::new(
+            rpc_handle.clone(),
+            api_availability.clone(),
+            cache_dir.clone(),
+            internal_event_tx.to_specialized_sender(),
+            app_version_info.clone(),
+            settings.show_beta_releases,
+        );
+        tokio::spawn(version_updater.run());
+        let account_history =
+            account_history::AccountHistory::new(&settings_dir, settings.get_account_token())
+                .await
+                .map_err(Error::LoadAccountHistory)?;
+
         let wireguard_key_manager = wireguard::KeyManager::new(
             internal_event_tx.clone(),
             api_availability.clone(),
@@ -784,6 +774,8 @@ where
 
         daemon.ensure_wireguard_keys_for_current_account().await;
 
+        api_availability.unsuspend();
+
         Ok(daemon)
     }
 
diff --git a/mullvad-problem-report/src/lib.rs b/mullvad-problem-report/src/lib.rs
index e6312062ff..9a5bca1a61 100644
--- a/mullvad-problem-report/src/lib.rs
+++ b/mullvad-problem-report/src/lib.rs
@@ -283,7 +283,6 @@ pub fn send_problem_report(
             None,
             cache_dir,
             false,
-            |_| Ok(()),
             #[cfg(target_os = "android")]
             None,
         ))
diff --git a/mullvad-rpc/src/address_cache.rs b/mullvad-rpc/src/address_cache.rs
index b681265dc5..44a7127f7d 100644
--- a/mullvad-rpc/src/address_cache.rs
+++ b/mullvad-rpc/src/address_cache.rs
@@ -44,11 +44,7 @@ pub struct AddressCache {
 
 impl AddressCache {
     /// Initialize cache using the given list, and write changes to `write_path`.
-    pub fn new(
-        addresses: Vec<SocketAddr>,
-        write_path: Option<Box<Path>>,
-        change_listener: Arc<Box<CurrentAddressChangeListener>>,
-    ) -> Result<Self, Error> {
+    pub fn new(addresses: Vec<SocketAddr>, write_path: Option<Box<Path>>) -> Result<Self, Error> {
         let mut cache = AddressCacheInner::from_addresses(addresses)?;
         cache.shuffle_tail();
         log::trace!("API address cache: {:?}", cache.addresses);
@@ -57,23 +53,15 @@ impl AddressCache {
         let address_cache = Self {
             inner: Arc::new(Mutex::new(cache)),
             write_path: write_path.map(|cache| Arc::from(cache)),
-            change_listener,
+            change_listener: Arc::new(Box::new(|_| Ok(()))),
         };
         Ok(address_cache)
     }
 
     /// Initialize cache using `read_path`, and write changes to `write_path`.
-    pub async fn from_file(
-        read_path: &Path,
-        write_path: Option<Box<Path>>,
-        change_listener: Arc<Box<CurrentAddressChangeListener>>,
-    ) -> Result<Self, Error> {
+    pub async fn from_file(read_path: &Path, write_path: Option<Box<Path>>) -> Result<Self, Error> {
         log::debug!("Loading API addresses from {:?}", read_path);
-        Self::new(
-            read_address_file(read_path).await?,
-            write_path,
-            change_listener,
-        )
+        Self::new(read_address_file(read_path).await?, write_path)
     }
 
     pub fn set_change_listener(&mut self, change_listener: Arc<Box<CurrentAddressChangeListener>>) {
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index 38f1499d02..61654cab1c 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -117,11 +117,7 @@ impl MullvadRpcRuntime {
     ) -> Result<Self, Error> {
         Ok(MullvadRpcRuntime {
             handle,
-            address_cache: AddressCache::new(
-                vec![API_ADDRESS.clone()],
-                None,
-                Arc::new(Box::new(|_| Ok(()))),
-            )?,
+            address_cache: AddressCache::new(vec![API_ADDRESS.clone()], None)?,
             api_availability: ApiAvailability::new(availability::State::default()),
             #[cfg(target_os = "android")]
             socket_bypass_tx,
@@ -136,7 +132,6 @@ impl MullvadRpcRuntime {
         resource_dir: Option<&Path>,
         cache_dir: &Path,
         write_changes: bool,
-        address_change_listener: impl Fn(SocketAddr) -> Result<(), ()> + Send + Sync + 'static,
         #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
     ) -> Result<Self, Error> {
         #[cfg(feature = "api-override")]
@@ -155,16 +150,7 @@ impl MullvadRpcRuntime {
             None
         };
 
-        let address_change_listener =
-            Arc::<Box<CurrentAddressChangeListener>>::new(Box::new(address_change_listener));
-
-        let address_cache = match AddressCache::from_file(
-            &cache_file,
-            write_file.clone(),
-            address_change_listener.clone(),
-        )
-        .await
-        {
+        let address_cache = match AddressCache::from_file(&cache_file, write_file.clone()).await {
             Ok(cache) => cache,
             Err(error) => {
                 let cache_exists = cache_file.exists();
@@ -181,12 +167,8 @@ impl MullvadRpcRuntime {
                 match resource_dir {
                     Some(resource_dir) => {
                         let read_file = resource_dir.join(API_IP_CACHE_FILENAME);
-                        let empty_listener =
-                            Arc::<Box<CurrentAddressChangeListener>>::new(Box::new(|_| Ok(())));
-                        let mut cache =
-                            AddressCache::from_file(&read_file, write_file, empty_listener).await?;
+                        let cache = AddressCache::from_file(&read_file, write_file).await?;
                         cache.randomize().await?;
-                        cache.set_change_listener(address_change_listener);
                         cache
                     }
                     None => return Err(Error::AddressCacheError(error)),
@@ -203,6 +185,14 @@ impl MullvadRpcRuntime {
         })
     }
 
+    pub fn set_address_change_listener(
+        &mut self,
+        address_change_listener: impl Fn(SocketAddr) -> Result<(), ()> + Send + Sync + 'static,
+    ) {
+        self.address_cache
+            .set_change_listener(Arc::new(Box::new(address_change_listener)));
+    }
+
     /// Creates a new request service and returns a handle to it.
     fn new_request_service(&mut self, sni_hostname: Option<String>) -> rest::RequestServiceHandle {
         let https_connector = HttpsConnectorWithSni::new(
diff --git a/mullvad-setup/src/main.rs b/mullvad-setup/src/main.rs
index 396d456afa..eb09513a1f 100644
--- a/mullvad-setup/src/main.rs
+++ b/mullvad-setup/src/main.rs
@@ -178,7 +178,6 @@ async fn remove_wireguard_key() -> Result<(), Error> {
                 None,
                 &cache_path,
                 false,
-                |_| Ok(()),
             )
             .await
             .map_err(Error::RpcInitializationError)?;