Add address change listener to AddressCache

5 files changed, 78 insertions, 12 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index b182686e19..9a55798f1e 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -48,7 +48,7 @@ use std::{
     io,
     marker::PhantomData,
     mem,
-    net::IpAddr,
+    net::{IpAddr, SocketAddr},
     path::PathBuf,
     sync::{mpsc as sync_mpsc, Arc, Weak},
     time::Duration,
@@ -260,6 +260,8 @@ pub(crate) enum InternalDaemonEvent {
     ),
     /// The background job fetching new `AppVersionInfo`s got a new info object.
     NewAppVersionInfo(AppVersionInfo),
+    /// A new API endpoint is being used
+    NewApiAddress(SocketAddr),
 }
 
 impl From<TunnelStateTransition> for InternalDaemonEvent {
@@ -491,11 +493,23 @@ where
         let (tunnel_state_machine_shutdown_tx, tunnel_state_machine_shutdown_signal) =
             oneshot::channel();
 
+        let (internal_event_tx, internal_event_rx) = command_channel.destructure();
+        let address_change_tx = std::sync::Mutex::new(internal_event_tx.clone());
+
         let mut rpc_runtime = mullvad_rpc::MullvadRpcRuntime::with_cache(
             tokio::runtime::Handle::current(),
             Some(&resource_dir),
             &user_cache_dir,
             true,
+            move |address| {
+                let tx = address_change_tx.lock().unwrap();
+                if tx
+                    .send(InternalDaemonEvent::NewApiAddress(address))
+                    .is_err()
+                {
+                    log::error!("Failed to send API address daemon event");
+                }
+            },
         )
         .await
         .map_err(Error::InitRpcFactory)?;
@@ -512,8 +526,6 @@ where
             &cache_dir,
         );
 
-        let (internal_event_tx, internal_event_rx) = command_channel.destructure();
-
 
         let mut settings = SettingsPersister::load(&settings_dir);
 
@@ -751,6 +763,10 @@ where
             NewAppVersionInfo(app_version_info) => {
                 self.handle_new_app_version_info(app_version_info)
             }
+            NewApiAddress(address) => {
+                // TODO
+                log::info!("ADDRESS! {:?}", address);
+            }
         }
     }
 
diff --git a/mullvad-problem-report/src/lib.rs b/mullvad-problem-report/src/lib.rs
index 4601a810ea..869661eb4e 100644
--- a/mullvad-problem-report/src/lib.rs
+++ b/mullvad-problem-report/src/lib.rs
@@ -281,6 +281,7 @@ pub fn send_problem_report(
             None,
             user_cache_dir,
             false,
+            |_| {},
         ))
         .map_err(Error::CreateRpcClientError)?;
     let rpc_client = mullvad_rpc::ProblemReportProxy::new(rpc_manager.mullvad_rest_handle());
diff --git a/mullvad-rpc/src/address_cache.rs b/mullvad-rpc/src/address_cache.rs
index 73b85e4998..5da1f09359 100644
--- a/mullvad-rpc/src/address_cache.rs
+++ b/mullvad-rpc/src/address_cache.rs
@@ -28,15 +28,22 @@ pub enum Error {
     EmptyAddressCache,
 }
 
+pub type CurrentAddressChangeListener = dyn Fn(SocketAddr) + Send + Sync + 'static;
+
 #[derive(Clone)]
 pub struct AddressCache {
     inner: Arc<Mutex<AddressCacheInner>>,
     write_path: Option<Arc<Path>>,
+    change_listener: Arc<Box<CurrentAddressChangeListener>>,
 }
 
 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>>) -> Result<Self, Error> {
+    pub fn new(
+        addresses: Vec<SocketAddr>,
+        write_path: Option<Box<Path>>,
+        change_listener: Arc<Box<CurrentAddressChangeListener>>,
+    ) -> Result<Self, Error> {
         let mut cache = AddressCacheInner::from_addresses(addresses)?;
         cache.shuffle_tail();
         log::trace!("API address cache: {:?}", cache.addresses);
@@ -45,14 +52,23 @@ impl AddressCache {
         let address_cache = Self {
             inner: Arc::new(Mutex::new(cache)),
             write_path: write_path.map(|cache| Arc::from(cache)),
+            change_listener,
         };
         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>>) -> Result<Self, Error> {
+    pub async fn from_file(
+        read_path: &Path,
+        write_path: Option<Box<Path>>,
+        change_listener: Arc<Box<CurrentAddressChangeListener>>,
+    ) -> Result<Self, Error> {
         log::debug!("Loading API addresses from {:?}", read_path);
-        Self::new(read_address_file(read_path).await?, write_path)
+        Self::new(
+            read_address_file(read_path).await?,
+            write_path,
+            change_listener,
+        )
     }
 
     /// Returns the currently selected address.
@@ -85,17 +101,19 @@ impl AddressCache {
     }
 
     pub async fn select_new_address(&self) {
-        let (new_choice, old_choice) = {
+        let (new_address, new_choice, old_choice) = {
             let mut inner = self.inner.lock().unwrap();
             let old_choice = inner.choice;
             inner.choice = inner.choice.wrapping_add(1);
-            (inner.choice, old_choice)
+            (Self::get_address_inner(&inner), inner.choice, old_choice)
         };
 
         if new_choice == old_choice {
             return;
         }
 
+        (*self.change_listener)(new_address);
+
         if let Err(error) = self.save_to_disk().await {
             log::error!("{}", error.display_chain());
         }
@@ -109,6 +127,9 @@ impl AddressCache {
             inner.shuffle();
             inner.choice = 0;
             inner.tried_current = false;
+
+            let new_address = Self::get_address_inner(&inner);
+            (*self.change_listener)(new_address);
         }
         self.save_to_disk().await.map_err(Error::WriteAddressCache)
     }
@@ -126,12 +147,18 @@ impl AddressCache {
                 inner.shuffle();
 
                 // Prefer a likely-working address
-                let choice = inner.addresses.iter().position(|&addr| addr == current_address);
+                let choice = inner
+                    .addresses
+                    .iter()
+                    .position(|&addr| addr == current_address);
                 if let Some(choice) = choice {
                     inner.choice = choice;
                 } else {
                     inner.choice = 0;
                     inner.tried_current = false;
+
+                    let new_address = Self::get_address_inner(&inner);
+                    (*self.change_listener)(new_address);
                 }
 
                 true
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index 10c5ca74a8..ca24288696 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -11,6 +11,7 @@ use std::{
     future::Future,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::Path,
+    sync::Arc,
 };
 use talpid_types::{net::wireguard, ErrorExt};
 
@@ -23,6 +24,7 @@ use crate::https_client_with_sni::HttpsConnectorWithSni;
 mod address_cache;
 mod relay_list;
 use address_cache::AddressCache;
+pub use address_cache::CurrentAddressChangeListener;
 pub use hyper::StatusCode;
 pub use relay_list::RelayListProxy;
 
@@ -60,7 +62,11 @@ impl MullvadRpcRuntime {
         Ok(MullvadRpcRuntime {
             https_connector: HttpsConnectorWithSni::new(),
             handle,
-            address_cache: AddressCache::new(vec![API_ADDRESS.into()], None)?,
+            address_cache: AddressCache::new(
+                vec![API_ADDRESS.into()],
+                None,
+                Arc::new(Box::new(|_| {})),
+            )?,
         })
     }
 
@@ -72,6 +78,7 @@ impl MullvadRpcRuntime {
         resource_dir: Option<&Path>,
         cache_dir: &Path,
         write_changes: bool,
+        address_change_listener: impl Fn(SocketAddr) + Send + Sync + 'static,
     ) -> Result<Self, Error> {
         let cache_file = cache_dir.join(API_IP_CACHE_FILENAME);
         let write_file = if write_changes {
@@ -80,7 +87,16 @@ impl MullvadRpcRuntime {
             None
         };
 
-        let address_cache = match AddressCache::from_file(&cache_file, write_file.clone()).await {
+        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
+        {
             Ok(cache) => cache,
             Err(error) => {
                 let cache_exists = cache_file.exists();
@@ -97,7 +113,12 @@ impl MullvadRpcRuntime {
                 match resource_dir {
                     Some(resource_dir) => {
                         let read_file = resource_dir.join(API_IP_CACHE_FILENAME);
-                        let cache = AddressCache::from_file(&read_file, write_file).await?;
+                        let cache = AddressCache::from_file(
+                            &read_file,
+                            write_file,
+                            address_change_listener,
+                        )
+                        .await?;
                         cache.randomize().await?;
                         cache
                     }
diff --git a/mullvad-setup/src/main.rs b/mullvad-setup/src/main.rs
index 7afb1fad16..7e1021f4ea 100644
--- a/mullvad-setup/src/main.rs
+++ b/mullvad-setup/src/main.rs
@@ -158,6 +158,7 @@ async fn clear_history() -> Result<(), Error> {
         None,
         &user_cache_path,
         false,
+        |_| {},
     )
     .await
     .map_err(Error::RpcInitializationError)?;