Fix actually constrain test to a single relay

test/test-manager/src/tests/helpers.rs

@@ -5,13 +5,17 @@ use crate::network_monitor::{
 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,
+};
 use mullvad_types::{
     constraints::Constraint,
     location::Location,
     relay_constraints::{
-        BridgeSettings, GeographicLocationConstraint, LocationConstraint, RelaySettings,
+        BridgeSettings, GeographicLocationConstraint, LocationConstraint, RelayConstraints,
+        RelaySettings,
     },
-    relay_list::{Relay, RelayList},
+    relay_list::Relay,
     states::TunnelState,
 };
 use pcap::Direction;
@@ -499,35 +503,57 @@ pub fn get_app_env() -> HashMap<String, String> {
     ])
 }
 
-/// Return a filtered version of the daemon's relay list.
+/// 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.
 ///
-/// * `mullvad_client` - An interface to the Mullvad daemon.
-/// * `critera` - A function used to determine which relays to return.
-pub async fn filter_relays<Filter>(
+/// # 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(
     mullvad_client: &mut MullvadProxyClient,
-    criteria: Filter,
-) -> Result<Vec<Relay>, Error>
-where
-    Filter: Fn(&Relay) -> bool,
-{
-    let relay_list: RelayList = mullvad_client
-        .get_relay_locations()
-        .await
-        .map_err(|error| Error::Daemon(format!("Failed to obtain relay list: {}", error)))?;
+    query: RelayQuery,
+) -> anyhow::Result<Relay> {
+    /// Convert the result of invoking the relay selector to a relay constraint.
+    fn convert_to_relay_constraints(
+        selected_relay: GetRelay,
+    ) -> anyhow::Result<(Relay, RelayConstraints)> {
+        match selected_relay {
+            GetRelay::Wireguard {
+                inner: WireguardConfig::Singlehop { exit },
+                ..
+            }
+            | GetRelay::OpenVpn { exit, .. } => {
+                let location = into_constraint(&exit)?;
+                let relay_constraints = RelayConstraints {
+                    location,
+                    ..Default::default()
+                };
+                Ok((exit, relay_constraints))
+            }
+            unsupported => bail!("Can not constrain to a {unsupported:?}"),
+        }
+    }
+
+    // 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 = RelaySelector::from_list(SelectorConfig::default(), 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)?;
+    let (exit, relay_constraints) = convert_to_relay_constraints(relay)?;
+    set_relay_settings(mullvad_client, RelaySettings::Normal(relay_constraints)).await?;
 
-    Ok(relay_list
-        .relays()
-        .filter(|relay| criteria(relay))
-        .cloned()
-        .collect())
+    Ok(exit)
 }
 
 /// Convenience function for constructing a constraint from a given [`Relay`].
 ///
 /// # Panics
 ///
 /// The relay must have a location set.
-pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
+pub fn into_constraint(relay: &Relay) -> anyhow::Result<Constraint<LocationConstraint>> {
     relay
         .location
         .as_ref()
@@ -537,16 +563,12 @@ pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
                  city_code,
                  ..
              }| {
-                GeographicLocationConstraint::Hostname(
-                    country_code.to_string(),
-                    city_code.to_string(),
-                    relay.hostname.to_string(),
-                )
+                GeographicLocationConstraint::hostname(country_code, city_code, &relay.hostname)
             },
         )
         .map(LocationConstraint::Location)
         .map(Constraint::Only)
-        .expect("relay is missing location")
+        .ok_or(anyhow!("relay is missing location"))
 }
 
 /// Ping monitoring made easy!

test/test-manager/src/tests/tunnel.rs

@@ -811,15 +811,13 @@ pub async fn test_mul_02_002(
     // Step 0 - Disconnect from any active tunnel connection
     helpers::disconnect_and_wait(&mut mullvad_client).await?;
     // Step 1 - Choose a relay
-    let relay_constraints = RelayQueryBuilder::new()
-        .openvpn()
-        .transport_protocol(TransportProtocol::Tcp)
-        .port(443)
-        .into_constraint();
-
-    set_relay_settings(
+    helpers::constrain_to_relay(
         &mut mullvad_client,
-        RelaySettings::Normal(relay_constraints),
+        RelayQueryBuilder::new()
+            .openvpn()
+            .transport_protocol(TransportProtocol::Tcp)
+            .port(443)
+            .build(),
     )
     .await?;
 
@@ -829,7 +827,7 @@ pub async fn test_mul_02_002(
         bail!("Expected tunnel state to be `Connected` - instead it was {tunnel_state:?}");
     };
     helpers::disconnect_and_wait(&mut mullvad_client).await?;
-    let gateway = endpoint.endpoint.address;
+    let target_endpoint = endpoint.endpoint.address;
 
     // Step 2 - Start a network monitor snooping the outbound network interface for some
     // identifiable payload
@@ -846,7 +844,7 @@ pub async fn test_mul_02_002(
         start_packet_monitor(identify_rogue_packet, MonitorOptions::default()).await;
 
     // Step 3 - Start the rogue program which will try to leak traffic to the chosen relay endpoint
-    let mut checker = ConnChecker::new(rpc.clone(), mullvad_client.clone(), gateway);
+    let mut checker = ConnChecker::new(rpc.clone(), mullvad_client.clone(), target_endpoint);
     let mut conn_artist = checker.spawn().await?;
     // Before proceeding, assert that the method of detecting identifiable packets work.
     conn_artist.check_connection().await?;

test/test-manager/src/tests/tunnel_state.rs

@@ -11,12 +11,12 @@ use crate::{
 };
 
 use mullvad_management_interface::MullvadProxyClient;
+use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
         GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
     },
-    relay_list::{Relay, RelayEndpointData},
     states::TunnelState,
     CustomTunnelEndpoint,
 };
@@ -340,40 +340,21 @@ pub async fn test_connected_state(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     let inet_destination = "1.1.1.1:1337".parse().unwrap();
 
     // Set relay to use
-    //
-
     log::info!("Select relay");
-
-    let relay_filter = |relay: &Relay| {
-        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
-    };
-
-    let relay = helpers::filter_relays(&mut mullvad_client, relay_filter)
-        .await?
-        .pop()
-        .unwrap();
-
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: helpers::into_constraint(&relay),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    let relay = helpers::constrain_to_relay(
+        &mut mullvad_client,
+        RelayQueryBuilder::new().wireguard().build(),
+    )
+    .await?;
 
     // Connect
-    //
-
     connect_and_wait(&mut mullvad_client).await?;
 
     // Verify that endpoint was selected
-    //
-
     match mullvad_client.get_tunnel_state().await? {
         TunnelState::Connected {
             endpoint:

test/test-manager/src/tests/ui.rs

@@ -1,5 +1,6 @@
 use super::{config::TEST_CONFIG, helpers, Error, TestContext};
 use mullvad_management_interface::MullvadProxyClient;
+use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     relay_constraints::{RelayConstraints, RelaySettings},
     relay_list::{Relay, RelayEndpointData},
@@ -85,25 +86,14 @@ pub async fn test_ui_tunnel_settings(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     // tunnel-state.spec precondition: a single WireGuard relay should be selected
     log::info!("Select WireGuard relay");
-    let entry = helpers::filter_relays(&mut mullvad_client, |relay: &Relay| {
-        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
-    })
-    .await?
-    .pop()
-    .unwrap();
-
-    // The test expects us to be disconnected and logged in but to have a specific relay selected
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: helpers::into_constraint(&entry),
-        ..Default::default()
-    });
-
-    helpers::set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    let entry = helpers::constrain_to_relay(
+        &mut mullvad_client,
+        RelayQueryBuilder::new().wireguard().build(),
+    )
+    .await?;
 
     let ui_result = run_test_env(
         &rpc,