Merge branch 'fix-overwrite-relay-location-test-framework'

9 files changed, 313 insertions, 222 deletions

diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index be12c29443..e6c91f9793 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -9,7 +9,7 @@ mod relay_selector;
 // Re-exports
 pub use error::Error;
 pub use relay_selector::{
-    detailer, query, relays::WireguardConfig, AdditionalRelayConstraints,
-    AdditionalWireguardConstraints, GetRelay, RelaySelector, RuntimeParameters, SelectedBridge,
-    SelectedObfuscator, SelectorConfig, RETRY_ORDER,
+    detailer, matcher, matcher::filter_matching_relay_list, query, relays::WireguardConfig,
+    AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, RelaySelector,
+    RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig, RETRY_ORDER,
 };
diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index f9fa0d7a01..d46b275b0a 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/mullvad-relay-selector/src/relay_selector/matcher.rs
@@ -8,20 +8,17 @@ use mullvad_types::{
         GeographicLocationConstraint, InternalBridgeConstraints, LocationConstraint, Ownership,
         Providers, ShadowsocksSettings,
     },
-    relay_list::{Relay, RelayEndpointData, WireguardRelayEndpointData},
+    relay_list::{Relay, RelayEndpointData, RelayList, WireguardRelayEndpointData},
 };
 use talpid_types::net::{IpVersion, TunnelType};
 
-use super::{
-    parsed_relays::ParsedRelays,
-    query::{ObfuscationQuery, RelayQuery, WireguardRelayQuery},
-};
+use super::query::{ObfuscationQuery, RelayQuery, WireguardRelayQuery};
 
 /// Filter a list of relays and their endpoints based on constraints.
 /// Only relays with (and including) matching endpoints are returned.
 pub fn filter_matching_relay_list(
     query: &RelayQuery,
-    relay_list: &ParsedRelays,
+    relay_list: &RelayList,
     custom_lists: &CustomListsSettings,
 ) -> Vec<Relay> {
     let relays = relay_list.relays();
@@ -133,13 +130,13 @@ pub fn filter_on_daita(filter: &Constraint<bool>, relay: &Relay) -> bool {
 /// Returns whether `relay` satisfies the obfuscation settings.
 fn filter_on_obfuscation(
     query: &WireguardRelayQuery,
-    relay_list: &ParsedRelays,
+    relay_list: &RelayList,
     relay: &Relay,
 ) -> bool {
     match &query.obfuscation {
         // Shadowsocks has relay-specific constraints
         ObfuscationQuery::Shadowsocks(settings) => {
-            let wg_data = &relay_list.parsed_list().wireguard;
+            let wg_data = &relay_list.wireguard;
             filter_on_shadowsocks(
                 &wg_data.shadowsocks_port_ranges,
                 &query.ip_version,
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 86f0b300e5..1522491b50 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -2,7 +2,7 @@
 
 pub mod detailer;
 mod helpers;
-mod matcher;
+pub mod matcher;
 mod parsed_relays;
 pub mod query;
 pub mod relays;
@@ -312,6 +312,21 @@ impl Default for SelectorConfig {
     }
 }
 
+impl TryFrom<Settings> for RelayQuery {
+    type Error = crate::Error;
+
+    fn try_from(value: Settings) -> Result<Self, Self::Error> {
+        let selector_config = SelectorConfig::from_settings(&value);
+        let specilized_selector_config = SpecializedSelectorConfig::from(&selector_config);
+        let SpecializedSelectorConfig::Normal(normal_selector_config) = specilized_selector_config
+        else {
+            return Err(Error::InvalidConstraints);
+        };
+
+        RelayQuery::try_from(normal_selector_config)
+    }
+}
+
 impl<'a> From<&'a SelectorConfig> for SpecializedSelectorConfig<'a> {
     fn from(value: &'a SelectorConfig) -> SpecializedSelectorConfig<'a> {
         match &value.relay_settings {
@@ -485,7 +500,7 @@ impl RelaySelector {
     /// Returns a non-custom bridge based on the relay and bridge constraints, ignoring the bridge
     /// state.
     pub fn get_bridge_forced(&self) -> Option<Shadowsocks> {
-        let parsed_relays = &self.parsed_relays.lock().unwrap();
+        let parsed_relays = &self.parsed_relays.lock().unwrap().parsed_list().clone();
         let config = self.config.lock().unwrap();
         let specialized_config = SpecializedSelectorConfig::from(&*config);
 
@@ -530,8 +545,8 @@ impl RelaySelector {
                 Ok(GetRelay::Custom(custom_config.clone()))
             }
             SpecializedSelectorConfig::Normal(normal_config) => {
-                let parsed_relays = &self.parsed_relays.lock().unwrap();
-                Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
+                let relay_list = &self.parsed_relays.lock().unwrap().parsed_list().clone();
+                Self::get_relay_inner(&query, relay_list, normal_config.custom_lists)
             }
         }
     }
@@ -566,16 +581,16 @@ impl RelaySelector {
                 Ok(GetRelay::Custom(custom_config.clone()))
             }
             SpecializedSelectorConfig::Normal(normal_config) => {
-                let parsed_relays = &self.parsed_relays.lock().unwrap();
+                let relay_list = self.parsed_relays.lock().unwrap().parsed_list().clone();
                 // Merge user preferences with the relay selector's default preferences.
                 let query = Self::pick_and_merge_query(
                     retry_attempt,
                     retry_order,
                     runtime_params,
                     &normal_config,
-                    parsed_relays,
+                    &relay_list,
                 )?;
-                Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
+                Self::get_relay_inner(&query, &relay_list, normal_config.custom_lists)
             }
         }
     }
@@ -599,7 +614,7 @@ impl RelaySelector {
         retry_order: &[RelayQuery],
         runtime_params: RuntimeParameters,
         user_config: &NormalSelectorConfig<'_>,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<RelayQuery, Error> {
         let user_query = RelayQuery::try_from(user_config.clone())?;
         log::trace!("Merging user preferences {user_query:?} with default retry strategy");
@@ -633,7 +648,7 @@ impl RelaySelector {
     #[cfg(not(target_os = "android"))]
     fn get_relay_inner(
         query: &RelayQuery,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         custom_lists: &CustomListsSettings,
     ) -> Result<GetRelay, Error> {
         match query.tunnel_protocol() {
@@ -665,7 +680,7 @@ impl RelaySelector {
     #[cfg(target_os = "android")]
     fn get_relay_inner(
         query: &RelayQuery,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         custom_lists: &CustomListsSettings,
     ) -> Result<GetRelay, Error> {
         // FIXME: A bit of defensive programming - calling `get_wireguard_relay` with a query that
@@ -693,7 +708,7 @@ impl RelaySelector {
     fn get_wireguard_relay(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<GetRelay, Error> {
         assert_eq!(
             query.tunnel_protocol(),
@@ -720,7 +735,7 @@ impl RelaySelector {
     fn get_wireguard_relay_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<WireguardConfig, Error> {
         let inner = if query.singlehop() {
             match Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays) {
@@ -766,7 +781,7 @@ impl RelaySelector {
     fn get_wireguard_singlehop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Option<Singlehop> {
         let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
         helpers::pick_random_relay(&candidates)
@@ -782,7 +797,7 @@ impl RelaySelector {
     fn get_wireguard_auto_multihop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<Multihop, Error> {
         let mut exit_relay_query = query.clone();
 
@@ -833,7 +848,7 @@ impl RelaySelector {
     fn get_wireguard_multihop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<Multihop, Error> {
         // Here, we modify the original query just a bit.
         // The actual query for an entry relay is identical as for an exit relay, with the
@@ -882,12 +897,12 @@ impl RelaySelector {
     /// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
     fn get_wireguard_endpoint(
         query: &RelayQuery,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         relay: &WireguardConfig,
     ) -> Result<MullvadWireguardEndpoint, Error> {
         wireguard_endpoint(
             query.wireguard_constraints(),
-            &parsed_relays.parsed_list().wireguard,
+            &parsed_relays.wireguard,
             relay,
         )
         .map_err(|internal| Error::NoEndpoint {
@@ -900,7 +915,7 @@ impl RelaySelector {
         query: &RelayQuery,
         relay: WireguardConfig,
         endpoint: &MullvadWireguardEndpoint,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<Option<SelectedObfuscator>, Error> {
         let obfuscator_relay = match relay {
             WireguardConfig::Singlehop { exit } => exit,
@@ -911,17 +926,14 @@ impl RelaySelector {
         match &query.wireguard_constraints().obfuscation {
             ObfuscationQuery::Off | ObfuscationQuery::Auto => Ok(None),
             ObfuscationQuery::Udp2tcp(settings) => {
-                let udp2tcp_ports = &parsed_relays.parsed_list().wireguard.udp2tcp_ports;
+                let udp2tcp_ports = &parsed_relays.wireguard.udp2tcp_ports;
 
                 helpers::get_udp2tcp_obfuscator(settings, udp2tcp_ports, obfuscator_relay, endpoint)
                     .map(Some)
                     .map_err(box_obfsucation_error)
             }
             ObfuscationQuery::Shadowsocks(settings) => {
-                let port_ranges = &parsed_relays
-                    .parsed_list()
-                    .wireguard
-                    .shadowsocks_port_ranges;
+                let port_ranges = &parsed_relays.wireguard.shadowsocks_port_ranges;
                 let obfuscation = helpers::get_shadowsocks_obfuscator(
                     settings,
                     port_ranges,
@@ -952,7 +964,7 @@ impl RelaySelector {
     fn get_openvpn_relay(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<GetRelay, Error> {
         assert_eq!(
             query.tunnel_protocol(),
@@ -985,17 +997,14 @@ impl RelaySelector {
     fn get_openvpn_endpoint(
         query: &RelayQuery,
         relay: &Relay,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Result<Endpoint, Error> {
-        openvpn_endpoint(
-            query.openvpn_constraints(),
-            &parsed_relays.parsed_list().openvpn,
-            relay,
+        openvpn_endpoint(query.openvpn_constraints(), &parsed_relays.openvpn, relay).map_err(
+            |internal| Error::NoEndpoint {
+                internal,
+                relay: EndpointErrorDetails::from_openvpn(relay.clone()),
+            },
         )
-        .map_err(|internal| Error::NoEndpoint {
-            internal,
-            relay: EndpointErrorDetails::from_openvpn(relay.clone()),
-        })
     }
 
     /// Selects a suitable bridge based on the specified settings, relay information, and transport
@@ -1019,7 +1028,7 @@ impl RelaySelector {
         query: &RelayQuery,
         relay: &Relay,
         protocol: &TransportProtocol,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         custom_lists: &CustomListsSettings,
     ) -> Result<Option<SelectedBridge>, Error> {
         if !BridgeQuery::should_use_bridge(&query.openvpn_constraints().bridge_settings) {
@@ -1050,7 +1059,7 @@ impl RelaySelector {
         query: &BridgeQuery,
         location: &Location,
         transport_protocol: TransportProtocol,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         custom_lists: &CustomListsSettings,
     ) -> Result<Option<SelectedBridge>, Error> {
         match query {
@@ -1079,13 +1088,13 @@ impl RelaySelector {
     ///
     /// The connection details are returned alongside the relay hosting the bridge.
     fn get_proxy_settings<T: Into<Coordinates>>(
-        parsed_relays: &ParsedRelays,
+        relay_list: &RelayList,
         constraints: &InternalBridgeConstraints,
         location: Option<T>,
         custom_lists: &CustomListsSettings,
     ) -> Result<(Shadowsocks, Relay), Error> {
-        let bridges = filter_matching_bridges(constraints, parsed_relays.relays(), custom_lists);
-        let bridge_data = &parsed_relays.parsed_list().bridge;
+        let bridges = filter_matching_bridges(constraints, relay_list.relays(), custom_lists);
+        let bridge_data = &relay_list.bridge;
         let bridge = match location {
             Some(location) => Self::get_proximate_bridge(bridges, location),
             None => helpers::pick_random_relay(&bridges)
@@ -1133,7 +1142,7 @@ impl RelaySelector {
     /// relays match the constraints.
     fn get_relay_midpoint(
         query: &RelayQuery,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
         custom_lists: &CustomListsSettings,
     ) -> Option<Coordinates> {
         use std::ops::Not;
@@ -1161,7 +1170,7 @@ impl RelaySelector {
     fn choose_openvpn_relay(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
-        parsed_relays: &ParsedRelays,
+        parsed_relays: &RelayList,
     ) -> Option<Relay> {
         // Filter among all valid relays
         let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
diff --git a/test/test-manager/src/tests/daita.rs b/test/test-manager/src/tests/daita.rs
index 24f5f4b0df..91687002aa 100644
--- a/test/test-manager/src/tests/daita.rs
+++ b/test/test-manager/src/tests/daita.rs
@@ -3,8 +3,8 @@ use futures::StreamExt;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
-    relay_constraints::GeographicLocationConstraint, relay_list::RelayEndpointData,
-    states::TunnelState,
+    constraints::Constraint, relay_constraints::GeographicLocationConstraint,
+    relay_list::RelayEndpointData, states::TunnelState,
 };
 use talpid_types::{net::TunnelEndpoint, tunnel::ErrorStateCause};
 use test_macro::test_function;
@@ -29,13 +29,11 @@ pub async fn test_daita(
     _rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    let relay_list = mullvad_client.get_relay_locations().await?;
-    let wg_relays = relay_list
-        .relays()
-        .flat_map(|relay| match &relay.endpoint_data {
-            RelayEndpointData::Wireguard(wireguard) => Some((relay, wireguard)),
-            _ => None,
-        });
+    let relays = helpers::get_all_pickable_relays(&mut mullvad_client).await?;
+    let wg_relays = relays.iter().flat_map(|relay| match &relay.endpoint_data {
+        RelayEndpointData::Wireguard(wireguard) => Some((relay, wireguard)),
+        _ => None,
+    });
 
     // Select two relays to use for the test, one with DAITA and one without.
     let daita_relay = wg_relays
@@ -62,38 +60,24 @@ pub async fn test_daita(
     );
     log::info!("Selected non-daita relay: {}", non_daita_relay.hostname);
 
-    let non_daita_location_query = RelayQueryBuilder::new()
-        .wireguard()
-        .location(non_daita_relay_location.clone())
-        .build();
-
-    let daita_location_query = RelayQueryBuilder::new()
-        .wireguard()
-        .location(daita_relay_location.clone())
-        .build();
-
-    let daita_to_non_daita_multihop_query = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .entry(daita_relay_location.clone())
-        .location(non_daita_relay_location.clone())
-        .build();
-
-    let non_daita_multihop_query = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .entry(non_daita_relay_location.clone())
-        .build();
+    log::info!("Setting wireguard and DAITA");
+    let wireguard_query = RelayQueryBuilder::new().wireguard().build();
+    helpers::apply_settings_from_relay_query(&mut mullvad_client, wireguard_query.clone()).await?;
+    mullvad_client.set_enable_daita(true).await?;
 
     let mut events = mullvad_client
         .events_listen()
         .await?
         .inspect(|event| log::debug!("New daemon event: {event:?}"));
 
-    log::info!("Connecting to non-daita relay with DAITA by automatically using multihop");
+    log::info!("Connecting to non-daita relay with DAITA should automatically use multihop");
     {
-        helpers::set_relay_settings(&mut mullvad_client, non_daita_location_query.clone()).await?;
-        mullvad_client.set_enable_daita(true).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(false).await?;
+
         mullvad_client.connect_tunnel().await?;
         let state = wait_for_daemon_reconnect(&mut events)
             .await
@@ -106,8 +90,12 @@ pub async fn test_daita(
         log::info!("Successfully multihopped with 'direct only' disabled");
     }
 
-    log::info!("Connecting to non-daita relay with 'DAITA: direct only'");
+    log::info!("Connecting to non-daita relay with 'direct_only' shoud fail");
     {
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+        })
+        .await?;
         mullvad_client.set_daita_direct_only(true).await?;
 
         let result = wait_for_daemon_reconnect(&mut events).await;
@@ -121,9 +109,13 @@ pub async fn test_daita(
         log::info!("Failed to connect, this is expected!");
     }
 
-    log::info!("Connecting to daita relay with 'direct_only' disabled");
+    log::info!("Connecting to daita relay with 'direct_only' should not use multihop");
     {
-        helpers::set_relay_settings(&mut mullvad_client, daita_location_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(daita_relay_location.clone().into());
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
 
         let state = wait_for_daemon_reconnect(&mut events)
             .await
@@ -139,9 +131,17 @@ pub async fn test_daita(
         log::info!("Successfully singlehopped with 'direct_only' disabled");
     }
 
-    log::info!("Connecting to daita relay with multihop");
+    log::info!("Connecting to a daita relay as entry for multihop and `direct_only` should work");
     {
-        helpers::set_relay_settings(&mut mullvad_client, daita_to_non_daita_multihop_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+            constraint.wireguard_constraints.entry_location =
+                Constraint::Only(daita_relay_location.clone().into());
+            constraint.wireguard_constraints.use_multihop = true;
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
+
         let state = wait_for_daemon_reconnect(&mut events)
             .await
             .context("Failed to connect via daita location with multihop enabled")?;
@@ -153,9 +153,19 @@ pub async fn test_daita(
         log::info!("Successfully connected with multihop");
     }
 
-    log::info!("Connecting to non_daita relay with multihop");
+    log::info!(
+        "Connecting to a non daita relay as entry for multihop and `direct_only` should fail"
+    );
     {
-        helpers::set_relay_settings(&mut mullvad_client, non_daita_multihop_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(daita_relay_location.clone().into());
+            constraint.wireguard_constraints.entry_location =
+                Constraint::Only(non_daita_relay_location.into());
+            constraint.wireguard_constraints.use_multihop = true;
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
+
         let result = wait_for_daemon_reconnect(&mut events).await;
         let Err(Error::UnexpectedErrorState(state)) = result else {
             bail!("Connection failed unsuccessfully, reason: {:?}", result);
diff --git a/test/test-manager/src/tests/dns.rs b/test/test-manager/src/tests/dns.rs
index 28db345cf7..b54be0e02e 100644
--- a/test/test-manager/src/tests/dns.rs
+++ b/test/test-manager/src/tests/dns.rs
@@ -8,7 +8,6 @@ use std::{
 use itertools::Itertools;
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_types::{
-    relay_constraints::RelaySettings,
     settings,
     wireguard::{DaitaSettings, QuantumResistantState},
     ConnectionConfig, CustomTunnelEndpoint,
@@ -18,7 +17,7 @@ use test_macro::test_function;
 use test_rpc::ServiceClient;
 
 use super::{
-    helpers::{self, connect_and_wait, set_relay_settings},
+    helpers::{self, connect_and_wait, set_custom_endpoint},
     Error, TestContext,
 };
 use crate::{
@@ -658,7 +657,7 @@ async fn connect_local_wg_relay(mullvad_client: &mut MullvadProxyClient) -> Resu
         CUSTOM_TUN_REMOTE_REAL_PORT,
     );
 
-    let relay_settings = RelaySettings::CustomTunnelEndpoint(CustomTunnelEndpoint {
+    let custom_tunnel_endpoint = CustomTunnelEndpoint {
         host: peer_addr.ip().to_string(),
         config: ConnectionConfig::Wireguard(wireguard::ConnectionConfig {
             tunnel: wireguard::TunnelConfig {
@@ -678,9 +677,8 @@ async fn connect_local_wg_relay(mullvad_client: &mut MullvadProxyClient) -> Resu
             fwmark: None,
             ipv6_gateway: None,
         }),
-    });
-
-    set_relay_settings(mullvad_client, relay_settings)
+    };
+    set_custom_endpoint(mullvad_client, custom_tunnel_endpoint)
         .await
         .expect("failed to update relay settings");
 
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 195c3da26a..f4a5986755 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -13,7 +13,8 @@ use anyhow::{anyhow, bail, ensure, Context};
 use futures::StreamExt;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_relay_selector::{
-    query::RelayQuery, GetRelay, RelaySelector, SelectorConfig, WireguardConfig,
+    query::{OpenVpnRelayQuery, RelayQuery, WireguardRelayQuery},
+    GetRelay, RelaySelector, SelectorConfig, WireguardConfig,
 };
 use mullvad_types::{
     constraints::Constraint,
@@ -133,7 +134,7 @@ pub async fn reboot(rpc: &mut ServiceClient) -> Result<(), Error> {
     #[cfg(target_os = "macos")]
     crate::vm::network::macos::configure_tunnel()
         .await
-        .map_err(|error| Error::Other(format!("Failed to recreate custom wg tun: {error}")))?;
+        .context("Failed to recreate custom wg tun: {error}")?;
 
     Ok(())
 }
@@ -596,40 +597,70 @@ impl<T> Drop for AbortOnDrop<T> {
     }
 }
 
+/// Applies the given query to the daemon location selection. The query will be intersected with
+/// the current location settings, so that the default location custom list is still used.
 pub async fn apply_settings_from_relay_query(
     mullvad_client: &mut MullvadProxyClient,
     query: RelayQuery,
-) -> Result<(), Error> {
-    let (constraints, bridge_state, bridge_settings, obfuscation) = query.into_settings();
+) -> anyhow::Result<()> {
+    // To prevent overwriting default custom list location constraint, we make an intersection with
+    // a query containing only the current location constraint
+    let intersected_relay_query = intersect_with_current_location(mullvad_client, query.clone())
+        .await
+        .with_context(|| {
+            format!("Failed to join query with current daemon settings. Query: {query:#?}")
+        })?;
+    let (constraints, bridge_state, bridge_settings, obfuscation) =
+        intersected_relay_query.into_settings();
 
     mullvad_client
         .set_relay_settings(constraints.into())
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))?;
+        .context("Failed to set daemon settings")?;
     mullvad_client
         .set_bridge_state(bridge_state)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set bridge state: {}", error)))?;
+        .context("Failed to set bridge state")?;
     mullvad_client
         .set_bridge_settings(bridge_settings)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set bridge settings: {}", error)))?;
+        .context("Failed to set bridge settings")?;
     mullvad_client
         .set_obfuscation_settings(obfuscation)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set obfuscation settings: {}", error)))
+        .context("Failed to set obfuscation settings")?;
+    Ok(())
 }
 
-pub async fn set_relay_settings(
+pub async fn set_custom_endpoint(
     mullvad_client: &mut MullvadProxyClient,
-    relay_settings: impl Into<RelaySettings>,
+    custom_endpoint: mullvad_types::CustomTunnelEndpoint,
 ) -> Result<(), Error> {
     mullvad_client
-        .set_relay_settings(relay_settings.into())
+        .set_relay_settings(RelaySettings::CustomTunnelEndpoint(custom_endpoint))
         .await
         .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))
 }
 
+pub async fn update_relay_constraints(
+    mullvad_client: &mut MullvadProxyClient,
+    fn_mut: impl FnOnce(&mut RelayConstraints),
+) -> anyhow::Result<()> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get setting from daemon")?;
+    let RelaySettings::Normal(mut relay_constraints) = settings.relay_settings else {
+        bail!("Mutating custom endpoint not supported");
+    };
+    fn_mut(&mut relay_constraints);
+    mullvad_client
+        .set_relay_settings(RelaySettings::Normal(relay_constraints))
+        .await
+        .context("Failed to set relay settings")?;
+    Ok(())
+}
+
 /// Wait for the relay list to be updated, to make sure we have the overridden one.
 /// Time out after a while.
 pub async fn ensure_updated_relay_list(
@@ -695,18 +726,99 @@ pub async fn get_app_env() -> anyhow::Result<HashMap<String, String>> {
     ]))
 }
 
-/// Constrain the daemon to only select the relay selected with `query` when establishing all
-/// future tunnels (until relay settings are updated, see [`set_relay_settings`]). Returns the
-/// selected [`Relay`] for future reference.
+/// Constrain the daemon to only select the relay compatible with `query` and the current relay
+/// settings when establishing all future tunnels (until relay settings are updated, see [`set_relay_settings`]).
+/// Returns the selected [`Relay`] for future reference.
+pub async fn constrain_to_relay(
+    mullvad_client: &mut MullvadProxyClient,
+    query: RelayQuery,
+) -> anyhow::Result<Relay> {
+    let intersect_query = intersect_with_current_location(mullvad_client, query).await?;
+    let (exit, relay_constraints) =
+        get_single_relay_location_contraint(mullvad_client, intersect_query).await?;
+
+    update_relay_constraints(mullvad_client, |current_constraints| {
+        *current_constraints = relay_constraints
+    })
+    .await
+    .unwrap();
+
+    Ok(exit)
+}
+
+/// Intersects the given query with the current location constraints, to prevent accidentally
+/// overwriting the default location custom list
+async fn intersect_with_current_location(
+    mullvad_client: &mut MullvadProxyClient,
+    query: RelayQuery,
+) -> anyhow::Result<RelayQuery> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get settings")?;
+    let RelaySettings::Normal(constraint) = settings.relay_settings else {
+        unimplemented!("Setting location for a custom endpoint is not supported");
+    };
+
+    // Construct a relay query preserving only the information about the current location
+    let current_location_query = RelayQuery::new(
+        constraint.location,
+        Constraint::Any,
+        Constraint::Any,
+        Constraint::Any,
+        WireguardRelayQuery {
+            entry_location: constraint.wireguard_constraints.entry_location,
+            ..Default::default()
+        },
+        OpenVpnRelayQuery {
+            bridge_settings: mullvad_relay_selector::query::BridgeQuery::Normal(
+                settings.bridge_settings.normal,
+            ),
+            ..Default::default()
+        },
+    )?;
+    use mullvad_types::Intersection;
+    let intersect_query = query
+        .intersection(current_location_query)
+        .context("Relay query incompatible with default settings")?;
+    Ok(intersect_query)
+}
+
+/// Get a query representing the current daemon settings
+async fn get_query_from_current_settings(
+    mullvad_client: &mut MullvadProxyClient,
+) -> anyhow::Result<RelayQuery> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get settings")?;
+    RelayQuery::try_from(settings).context("Failed to convert settings to relay query")
+}
+
+pub async fn get_all_pickable_relays(
+    mullvad_client: &mut MullvadProxyClient,
+) -> anyhow::Result<Vec<Relay>> {
+    let settings = mullvad_client.get_settings().await?;
+    let relay_list = mullvad_client.get_relay_locations().await?;
+    let relays = mullvad_relay_selector::filter_matching_relay_list(
+        &helpers::get_query_from_current_settings(mullvad_client).await?,
+        &relay_list,
+        &settings.custom_lists,
+    );
+    Ok(relays)
+}
+
+/// Selects a relay compatible with the given query and relay list from the client, and returns a
+/// location constraint for only that relay, along with the relay itself.
 ///
 /// # Note
 /// This function does not handle bridges and multihop configurations (currently). There is no
 /// particular reason for this other than it not being needed at the time, so feel free to extend
 /// this function :).
-pub async fn constrain_to_relay(
+async fn get_single_relay_location_contraint(
     mullvad_client: &mut MullvadProxyClient,
     query: RelayQuery,
-) -> anyhow::Result<Relay> {
+) -> anyhow::Result<(Relay, RelayConstraints)> {
     /// Convert the result of invoking the relay selector to a relay constraint.
     fn convert_to_relay_constraints(
         query: RelayQuery,
@@ -726,18 +838,11 @@ pub async fn constrain_to_relay(
             unsupported => bail!("Can not constrain to a {unsupported:?}"),
         }
     }
-
     let settings = mullvad_client.get_settings().await?;
-    // Construct a relay selector with up-to-date information from the runnin daemon's relay list
     let relay_list = mullvad_client.get_relay_locations().await?;
     let relay_selector = get_daemon_relay_selector(&settings, relay_list);
-    // Select an(y) appropriate relay for the given query and constrain the daemon to only connect
-    // to that specific relay (when connecting).
     let relay = relay_selector.get_relay_by_query(query.clone())?;
-    let (exit, relay_constraints) = convert_to_relay_constraints(query, relay)?;
-    set_relay_settings(mullvad_client, RelaySettings::Normal(relay_constraints)).await?;
-
-    Ok(exit)
+    convert_to_relay_constraints(query, relay)
 }
 
 /// Get a mirror of the relay selector used by the daemon.
@@ -1200,36 +1305,32 @@ fn parse_am_i_mullvad(result: String) -> anyhow::Result<bool> {
     })
 }
 
-/// Set the location to the given [`LocationConstraint`]. This also includes
-/// entry location for multihop. It does not, however, affect bridge location for OpenVPN.
-/// This is for simplify, as bridges default to using the server closest to the exit anyway, and
-/// OpenVPN is slated for removal.
-///
-/// NOTE: Calling this from within a test will overwrite the default test lcoation specified in
-/// the settings.
+/// Set the location to the given [`LocationConstraint`]. The same location constraint will be set
+/// for the multihop entry and OpenVPN bridge location as well.
 pub async fn set_location(
     mullvad_client: &mut MullvadProxyClient,
     location: impl Into<LocationConstraint>,
 ) -> anyhow::Result<()> {
-    let constraints = get_location_relay_constraints(location.into());
+    let location_constraint: LocationConstraint = location.into();
+    let mut settings = mullvad_client
+        .get_settings()
+        .await
+        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))?;
 
+    settings.bridge_settings.normal.location = Constraint::Only(location_constraint.clone());
     mullvad_client
-        .set_relay_settings(constraints.into())
-        .await
-        .context("Failed to set relay settings")
-}
+        .set_bridge_settings(settings.bridge_settings)
+        .await?;
 
-fn get_location_relay_constraints(custom_list: LocationConstraint) -> RelayConstraints {
-    let wireguard_constraints = mullvad_types::relay_constraints::WireguardConstraints {
-        entry_location: Constraint::Only(custom_list.clone()),
-        ..Default::default()
+    let RelaySettings::Normal(mut constraint) = settings.relay_settings else {
+        unimplemented!("Setting location for a custom endpoint is not supported");
     };
-
-    RelayConstraints {
-        location: Constraint::Only(custom_list),
-        wireguard_constraints,
-        ..Default::default()
-    }
+    constraint.location = Constraint::Only(location_constraint.clone());
+    constraint.wireguard_constraints.entry_location = Constraint::Only(location_constraint);
+    mullvad_client
+        .set_relay_settings(RelaySettings::Normal(constraint))
+        .await?;
+    Ok(())
 }
 
 /// Dig out a custom list from the daemon settings based on the custom list's name.
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 4f5362167f..66635a10c1 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -72,9 +72,8 @@ pub enum Error {
     #[error("GUI test binary missing")]
     MissingGuiTest,
 
-    #[cfg(target_os = "macos")]
     #[error("An error occurred: {0}")]
-    Other(String),
+    Other(#[from] anyhow::Error),
 }
 
 #[derive(Clone)]
@@ -145,13 +144,13 @@ pub async fn prepare_daemon(
         .context("Failed to restart daemon")?;
 
     log::debug!("Resetting daemon settings before test");
+    helpers::disconnect_and_wait(&mut mullvad_client)
+        .await
+        .context("Failed to disconnect daemon after test")?;
     mullvad_client
         .reset_settings()
         .await
         .context("Failed to reset settings")?;
-    helpers::disconnect_and_wait(&mut mullvad_client)
-        .await
-        .context("Failed to disconnect daemon after test")?;
     helpers::ensure_logged_in(&mut mullvad_client).await?;
 
     Ok(mullvad_client)
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index b64a16d854..24357fdfad 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -1,14 +1,11 @@
 use super::{
     config::TEST_CONFIG,
-    helpers::{
-        self, apply_settings_from_relay_query, connect_and_wait, disconnect_and_wait,
-        set_relay_settings,
-    },
+    helpers::{self, apply_settings_from_relay_query, connect_and_wait, disconnect_and_wait},
     Error, TestContext,
 };
 use crate::{
     network_monitor::{start_packet_monitor, MonitorOptions},
-    tests::helpers::login_with_retries,
+    tests::helpers::{login_with_retries, update_relay_constraints},
 };
 
 use anyhow::Context;
@@ -17,8 +14,7 @@ use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
-        self, BridgeConstraints, BridgeSettings, BridgeType, OpenVpnConstraints, RelayConstraints,
-        RelaySettings, TransportPort,
+        self, BridgeConstraints, BridgeSettings, BridgeType, OpenVpnConstraints, TransportPort,
     },
     wireguard,
 };
@@ -63,15 +59,12 @@ pub async fn test_openvpn_tunnel(
     for (protocol, constraint) in CONSTRAINTS {
         log::info!("Connect to {protocol} OpenVPN endpoint");
 
-        let relay_settings = RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            openvpn_constraints: OpenVpnConstraints { port: constraint },
-            ..Default::default()
-        });
-
-        set_relay_settings(&mut mullvad_client, relay_settings)
-            .await
-            .expect("failed to update relay settings");
+        update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+            relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+            relay_constraints.openvpn_constraints = OpenVpnConstraints { port: constraint };
+        })
+        .await
+        .expect("failed to update relay constraints");
 
         connect_and_wait(&mut mullvad_client).await?;
 
@@ -353,15 +346,11 @@ pub async fn test_wireguard_autoconnect(
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
     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");
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
+    })
+    .await
+    .expect("failed to update relay constraints");
 
     mullvad_client
         .set_auto_connect(true)
@@ -396,14 +385,11 @@ pub async fn test_openvpn_autoconnect(
 ) -> Result<(), Error> {
     log::info!("Setting tunnel protocol to OpenVPN");
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
+    .await
+    .expect("failed to update relay constraints");
 
     mullvad_client
         .set_auto_connect(true)
@@ -596,15 +582,11 @@ pub async fn test_remote_socks_bridge(
         .await
         .expect("failed to update bridge settings");
 
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            ..Default::default()
-        }),
-    )
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
     .await
-    .expect("failed to update relay settings");
+    .expect("failed to update relay constraints");
 
     // Connect to VPN
     //
@@ -694,15 +676,11 @@ pub async fn test_local_socks_bridge(
         .await
         .expect("failed to update bridge settings");
 
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            ..Default::default()
-        }),
-    )
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
     .await
-    .expect("failed to update relay settings");
+    .expect("failed to update relay constraints");
 
     // Connect to VPN
     //
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index 5ccf0f863e..4005f76480 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -1,19 +1,20 @@
 use super::{
     helpers::{
-        self, connect_and_wait, send_guest_probes, set_relay_settings,
-        unreachable_wireguard_tunnel, wait_for_tunnel_state,
+        self, connect_and_wait, send_guest_probes, unreachable_wireguard_tunnel,
+        wait_for_tunnel_state,
     },
     ui, Error, TestContext,
 };
-use crate::{assert_tunnel_state, tests::helpers::ping_sized_with_timeout};
+use crate::{
+    assert_tunnel_state,
+    tests::helpers::{ping_sized_with_timeout, set_custom_endpoint, update_relay_constraints},
+};
 
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
-    relay_constraints::{
-        GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
-    },
+    relay_constraints::{GeographicLocationConstraint, LocationConstraint},
     states::TunnelState,
     CustomTunnelEndpoint,
 };
@@ -185,14 +186,15 @@ pub async fn test_connecting_state(
     log::info!("Verify tunnel state: disconnected");
     assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected { .. });
 
-    let relay_settings = RelaySettings::CustomTunnelEndpoint(CustomTunnelEndpoint {
-        host: "1.3.3.7".to_owned(),
-        config: mullvad_types::ConnectionConfig::Wireguard(unreachable_wireguard_tunnel()),
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    set_custom_endpoint(
+        &mut mullvad_client,
+        CustomTunnelEndpoint {
+            host: "1.3.3.7".to_owned(),
+            config: mullvad_types::ConnectionConfig::Wireguard(unreachable_wireguard_tunnel()),
+        },
+    )
+    .await
+    .expect("failed to update relay settings");
 
     mullvad_client
         .connect_tunnel()
@@ -273,21 +275,18 @@ pub async fn test_error_state(
 
     log::info!("Enter error state");
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: Constraint::Only(LocationConstraint::from(
-            GeographicLocationConstraint::country("xx"),
-        )),
-        ..Default::default()
-    });
-
     mullvad_client
         .set_allow_lan(false)
         .await
         .expect("failed to disable LAN sharing");
 
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    update_relay_constraints(&mut mullvad_client, |constraints| {
+        constraints.location = Constraint::Only(LocationConstraint::from(
+            GeographicLocationConstraint::country("xx"),
+        ))
+    })
+    .await
+    .expect("Failed to set invalid location");
 
     let _ = connect_and_wait(&mut mullvad_client).await;
     assert_tunnel_state!(&mut mullvad_client, TunnelState::Error { .. });