Merge branch 'fix-connection-timeout-errors-caused-by-far-away-relays-des-1419'

3 files changed, 167 insertions, 10 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 2ebc5eb7ad..10842985ef 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -1181,3 +1181,127 @@ fn parse_am_i_mullvad(result: String) -> anyhow::Result<bool> {
         bail!("Unexpected output from connection-checker: {result:?}")
     })
 }
+
+pub mod custom_lists {
+    use super::*;
+    use mullvad_types::custom_list::{CustomList, Id};
+    use std::sync::{LazyLock, Mutex};
+
+    // Expose all custom list variants as a shorthand.
+    pub use List::*;
+
+    /// Mapping between [List] to daemon custom lists. Since custom list ids are assigned by the daemon at the creation
+    /// of the custom list settings object, we can't map a custom list name to a specific list before runtime.
+    static IDS: LazyLock<Mutex<HashMap<List, Id>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
+
+    /// Pre-defined (well-typed) custom lists which may be useuful in different test scenarios.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    pub enum List {
+        /// A selection of Nordic servers
+        Nordic,
+        /// A selection of European servers
+        Europe,
+        /// This custom list contains relays which are close geographically to the computer running
+        /// the test scenarios, which hopefully means there will be little latency between the test
+        /// machine and these relays
+        LowLatency,
+        /// Antithesis of [List::LowLatency], these relays are located far away from the test
+        /// server. Use this custom list if you want to simulate scenarios where the probability
+        /// of experiencing high latencies is desirable.
+        HighLatency,
+    }
+
+    impl List {
+        pub fn name(self) -> String {
+            use List::*;
+            match self {
+                Nordic => "Nordic".to_string(),
+                Europe => "Europe".to_string(),
+                LowLatency => "Low Latency".to_string(),
+                HighLatency => "High Latency".to_string(),
+            }
+        }
+
+        /// Iterator over all custom lists.
+        pub fn all() -> impl Iterator<Item = List> {
+            use List::*;
+            [Nordic, Europe, LowLatency, HighLatency].into_iter()
+        }
+
+        pub fn locations(self) -> impl Iterator<Item = GeographicLocationConstraint> {
+            use List::*;
+            let country = GeographicLocationConstraint::country;
+            let city = GeographicLocationConstraint::city;
+            match self {
+                Nordic => {
+                    vec![country("no"), country("se"), country("fi"), country("dk")].into_iter()
+                }
+                Europe => vec![
+                    // North
+                    country("se"),
+                    // West
+                    country("fr"),
+                    // East
+                    country("ro"),
+                    // South
+                    country("it"),
+                ]
+                .into_iter(),
+                LowLatency => {
+                    // Assumption: Test server is located in Gothenburg, Sweden.
+                    vec![city("se", "got")].into_iter()
+                }
+                HighLatency => {
+                    // Assumption: Test server is located in Gothenburg, Sweden.
+                    vec![country("au"), country("ca"), country("za")].into_iter()
+                }
+            }
+        }
+
+        pub fn to_constraint(self) -> Option<LocationConstraint> {
+            let ids = IDS.lock().unwrap();
+            let id = ids.get(&self)?;
+            Some(LocationConstraint::CustomList { list_id: *id })
+        }
+    }
+
+    impl From<List> for LocationConstraint {
+        fn from(custom_list: List) -> Self {
+            // TODO: Is this _too_ unsound ??
+            custom_list.to_constraint().unwrap()
+        }
+    }
+
+    /// Add a set of custom lists which can be used in different test scenarios.
+    ///
+    /// See [`List`] for available custom lists.
+    pub async fn add_default_lists(mullvad_client: &mut MullvadProxyClient) -> anyhow::Result<()> {
+        for custom_list in List::all() {
+            let id = mullvad_client
+                .create_custom_list(custom_list.name())
+                .await?;
+            let mut daemon_dito = find_custom_list(mullvad_client, &custom_list.name()).await?;
+            for locations in custom_list.locations() {
+                daemon_dito.locations.insert(locations);
+            }
+            mullvad_client.update_custom_list(daemon_dito).await?;
+            // Associate this custom list variant with a specific, runtime custom list id.
+            IDS.lock().unwrap().insert(custom_list, id);
+        }
+        Ok(())
+    }
+
+    /// Dig out a custom list from the daemon settings based on the custom list's name.
+    /// There should be an rpc for this.
+    async fn find_custom_list(
+        rpc: &mut MullvadProxyClient,
+        name: &str,
+    ) -> anyhow::Result<CustomList> {
+        rpc.get_settings()
+            .await?
+            .custom_lists
+            .into_iter()
+            .find(|list| list.name == name)
+            .ok_or(anyhow!("List '{name}' not found"))
+    }
+}
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index bc17a7f3f6..14c2a1e3a5 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -121,6 +121,7 @@ pub async fn prepare_daemon(
         .await
         .context("Failed to disconnect daemon after test")?;
     helpers::ensure_logged_in(&mut mullvad_client).await?;
+    helpers::custom_lists::add_default_lists(&mut mullvad_client).await?;
 
     Ok(())
 }
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index 6fb57adc14..fde8c73049 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -195,7 +195,15 @@ pub async fn test_wireguard_over_shadowsocks(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    let query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
+    // This is an attempt to try to reduce this type of flakiness.
+    use helpers::custom_lists::LowLatency;
+
+    let query = RelayQueryBuilder::new()
+        .wireguard()
+        .shadowsocks()
+        .location(LowLatency)
+        .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
@@ -294,7 +302,16 @@ pub async fn test_multihop(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
-    let query = RelayQueryBuilder::new().wireguard().multihop().build();
+    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
+    // This is an attempt to try to reduce this type of flakiness.
+    use helpers::custom_lists::LowLatency;
+
+    let query = RelayQueryBuilder::new()
+        .wireguard()
+        .multihop()
+        .location(LowLatency)
+        .entry(LowLatency)
+        .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
@@ -445,7 +462,11 @@ pub async fn test_quantum_resistant_tunnel(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
+    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
+    // This is an attempt to try to reduce this type of flakiness.
+    use helpers::custom_lists::LowLatency;
+
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::Off)
         .await
@@ -459,13 +480,12 @@ pub async fn test_quantum_resistant_tunnel(
 
     log::info!("Setting tunnel protocol to WireGuard");
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-        ..Default::default()
-    });
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("Failed to update relay settings");
+    let query = RelayQueryBuilder::new()
+        .wireguard()
+        .location(LowLatency)
+        .build();
+
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::On)
@@ -524,6 +544,10 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
+    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
+    // This is an attempt to try to reduce this type of flakiness.
+    use helpers::custom_lists::LowLatency;
+
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::On)
         .await
@@ -533,6 +557,8 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
         .wireguard()
         .multihop()
         .udp2tcp()
+        .entry(LowLatency)
+        .location(LowLatency)
         .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
@@ -559,6 +585,10 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
+    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
+    // This is an attempt to try to reduce this type of flakiness.
+    use helpers::custom_lists::LowLatency;
+
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::On)
         .await
@@ -568,6 +598,8 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
         .wireguard()
         .multihop()
         .shadowsocks()
+        .entry(LowLatency)
+        .location(LowLatency)
         .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;