Merge branch 'implement-use-anywhere-daita-setting-des-1068'

6 files changed, 216 insertions, 58 deletions

diff --git a/test/test-manager/src/tests/daita.rs b/test/test-manager/src/tests/daita.rs
new file mode 100644
index 0000000000..912c811ecc
--- /dev/null
+++ b/test/test-manager/src/tests/daita.rs
@@ -0,0 +1,200 @@
+use anyhow::{anyhow, bail, ensure, Context};
+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,
+};
+use talpid_types::{net::TunnelEndpoint, tunnel::ErrorStateCause};
+use test_macro::test_function;
+use test_rpc::ServiceClient;
+
+use super::{helpers, Error, TestContext};
+
+/// Test that daita and daita_smart_routing works by connecting
+/// - to a non-DAITA relay with singlehop (should block)
+/// - to a DAITA relay with singlehop
+/// - to a DAITA relay with auto-multihop using smart_routing
+/// - to a DAITA relay with explicit multihop
+/// - to a non-DAITA relay with multihop (should block)
+///
+/// # Limitations
+///
+/// The test does not analyze any traffic, nor verify that DAITA is in use in any way except
+/// by looking at [TunnelEndpoint::daita].
+#[test_function]
+pub async fn test_daita(
+    _ctx: TestContext,
+    _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,
+        });
+
+    // Select two relays to use for the test, one with DAITA and one without.
+    let daita_relay = wg_relays
+        .clone()
+        .find(|(_relay, wireguard_data)| wireguard_data.daita)
+        .map(|(relay, _)| relay)
+        .context("Failed to find a daita wireguard relay")?;
+    log::info!("Selected daita relay: {}", daita_relay.hostname);
+    let daita_relay_location = GeographicLocationConstraint::hostname(
+        &daita_relay.location.country_code,
+        &daita_relay.location.city_code,
+        &daita_relay.hostname,
+    );
+
+    let non_daita_relay = wg_relays
+        .clone()
+        .find(|(_relay, wireguard_data)| !wireguard_data.daita)
+        .map(|(relay, _)| relay)
+        .context("Failed to find a non-daita wireguard relay")?;
+    let non_daita_relay_location = GeographicLocationConstraint::hostname(
+        &non_daita_relay.location.country_code,
+        &non_daita_relay.location.city_code,
+        &non_daita_relay.hostname,
+    );
+    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();
+
+    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 smart routing");
+    {
+        helpers::set_relay_settings(&mut mullvad_client, non_daita_location_query.clone()).await?;
+        mullvad_client.set_enable_daita(true).await?;
+        mullvad_client.connect_tunnel().await?;
+        let state = wait_for_daemon_reconnect(&mut events)
+            .await
+            .context("Failed to connect with smart_routing enabled")?;
+
+        let endpoint: &TunnelEndpoint = state.endpoint().ok_or(anyhow!("No endpoint"))?;
+        ensure!(endpoint.daita, "DAITA must be used");
+        ensure!(endpoint.entry_endpoint.is_some(), "multihop must be used");
+
+        log::info!("Successfully multihopped with use smart_routing");
+    }
+
+    log::info!("Connecting to non-daita relay with DAITA but no smart routing");
+    {
+        mullvad_client.set_daita_smart_routing(false).await?;
+
+        let result = wait_for_daemon_reconnect(&mut events).await;
+        let Err(Error::UnexpectedErrorState(state)) = result else {
+            bail!("Connection failed unsuccessfully, reason: {:?}", result);
+        };
+        let ErrorStateCause::TunnelParameterError(_) = state.cause() else {
+            bail!("Connection failed unsuccessfully, cause: {}", state.cause());
+        };
+
+        log::info!("Failed to connect, this is expected!");
+    }
+
+    log::info!("Connecting to daita relay with smart_routing");
+    {
+        helpers::set_relay_settings(&mut mullvad_client, daita_location_query).await?;
+
+        let state = wait_for_daemon_reconnect(&mut events)
+            .await
+            .context("Failed to connect to daita location with smart_routing enabled")?;
+
+        let endpoint = state.endpoint().context("No endpoint")?;
+        ensure!(endpoint.daita, "DAITA must be used");
+        ensure!(
+            endpoint.entry_endpoint.is_none(),
+            "multihop must not be used"
+        );
+
+        log::info!("Successfully singlehopped with smart_routing");
+    }
+
+    log::info!("Connecting to daita relay with multihop");
+    {
+        helpers::set_relay_settings(&mut mullvad_client, daita_to_non_daita_multihop_query).await?;
+        let state = wait_for_daemon_reconnect(&mut events)
+            .await
+            .context("Failed to connect via daita location with multihop enabled")?;
+
+        let endpoint = state.endpoint().context("No endpoint")?;
+        ensure!(endpoint.daita, "DAITA must be used");
+        ensure!(endpoint.entry_endpoint.is_some(), "multihop must be used");
+
+        log::info!("Successfully connected with multihop");
+    }
+
+    log::info!("Connecting to non_daita relay with multihop");
+    {
+        helpers::set_relay_settings(&mut mullvad_client, non_daita_multihop_query).await?;
+        let result = wait_for_daemon_reconnect(&mut events).await;
+        let Err(Error::UnexpectedErrorState(state)) = result else {
+            bail!("Connection failed unsuccessfully, reason: {:?}", result);
+        };
+        let ErrorStateCause::TunnelParameterError(_) = state.cause() else {
+            bail!("Connection failed unsuccessfully, cause: {}", state.cause());
+        };
+
+        log::info!("Failed to connect, this is expected!");
+    }
+
+    Ok(())
+}
+
+async fn wait_for_daemon_reconnect(
+    mut event_stream: impl futures::Stream<Item = Result<DaemonEvent, mullvad_management_interface::Error>>
+        + Unpin,
+) -> Result<TunnelState, Error> {
+    // wait until the daemon informs us that it's trying to connect
+    helpers::find_daemon_event(&mut event_stream, |event| match event {
+        DaemonEvent::TunnelState(state) => Some(match state {
+            TunnelState::Connecting { .. } => Ok(state),
+            TunnelState::Connected { .. } => return None,
+            TunnelState::Disconnecting { .. } => return None,
+            TunnelState::Disconnected { .. } => Err(Error::UnexpectedTunnelState(Box::new(state))),
+            TunnelState::Error(state) => Err(Error::UnexpectedErrorState(state)),
+        }),
+        _ => None,
+    })
+    .await??;
+
+    // then wait until the daemon informs us that it connected (or failed)
+    helpers::find_daemon_event(&mut event_stream, |event| match event {
+        DaemonEvent::TunnelState(state) => match state {
+            TunnelState::Connecting { .. } => None,
+            TunnelState::Connected { .. } => Some(Ok(state)),
+            _ => Some(Err(Error::UnexpectedTunnelState(Box::new(state)))),
+        },
+        _ => None,
+    })
+    .await?
+}
diff --git a/test/test-manager/src/tests/dns.rs b/test/test-manager/src/tests/dns.rs
index d7ea3d021d..69f98450ca 100644
--- a/test/test-manager/src/tests/dns.rs
+++ b/test/test-manager/src/tests/dns.rs
@@ -642,7 +642,7 @@ async fn connect_local_wg_relay(mullvad_client: &mut MullvadProxyClient) -> Resu
         .set_quantum_resistant_tunnel(QuantumResistantState::Off)
         .await?;
     mullvad_client
-        .set_daita_settings(DaitaSettings { enabled: false })
+        .set_daita_settings(DaitaSettings::default())
         .await?;
 
     let peer_addr: SocketAddr = SocketAddr::new(
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 6c51b1f185..3eb0be0cc2 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -17,7 +17,6 @@ use mullvad_relay_selector::{
 };
 use mullvad_types::{
     constraints::Constraint,
-    location::Location,
     relay_constraints::{
         GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
     },
@@ -710,7 +709,7 @@ pub async fn constrain_to_relay(
                 ..
             }
             | GetRelay::OpenVpn { exit, .. } => {
-                let location = into_constraint(&exit)?;
+                let location = into_constraint(&exit);
                 let (mut relay_constraints, ..) = query.into_settings();
                 relay_constraints.location = location;
                 Ok((exit, relay_constraints))
@@ -736,22 +735,14 @@ pub async fn constrain_to_relay(
 /// # Panics
 ///
 /// The relay must have a location set.
-pub fn into_constraint(relay: &Relay) -> anyhow::Result<Constraint<LocationConstraint>> {
-    relay
-        .location
-        .as_ref()
-        .map(
-            |Location {
-                 country_code,
-                 city_code,
-                 ..
-             }| {
-                GeographicLocationConstraint::hostname(country_code, city_code, &relay.hostname)
-            },
-        )
-        .map(LocationConstraint::Location)
-        .map(Constraint::Only)
-        .ok_or(anyhow!("relay is missing location"))
+pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
+    let constraint = GeographicLocationConstraint::hostname(
+        relay.location.country_code.clone(),
+        relay.location.city_code.clone(),
+        &relay.hostname,
+    );
+
+    Constraint::Only(LocationConstraint::Location(constraint))
 }
 
 /// Ping monitoring made easy!
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 7e4cbc9eb6..bc17a7f3f6 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -2,6 +2,7 @@ mod access_methods;
 mod account;
 pub mod config;
 mod cve_2019_14899;
+mod daita;
 mod dns;
 mod helpers;
 mod install;
@@ -57,7 +58,10 @@ pub enum Error {
     #[error("The daemon returned an error: {0}")]
     Daemon(String),
 
-    #[error("The daemon ended up in the error state")]
+    #[error("The daemon ended up in the the wrong tunnel-state: {0:?}")]
+    UnexpectedTunnelState(Box<mullvad_types::states::TunnelState>),
+
+    #[error("The daemon ended up in the error state: {0:?}")]
     UnexpectedErrorState(talpid_types::tunnel::ErrorState),
 
     #[error("The gRPC client ran into an error: {0}")]
diff --git a/test/test-manager/src/tests/relay_ip_overrides.rs b/test/test-manager/src/tests/relay_ip_overrides.rs
index 48df8ff0e4..3805412ee9 100644
--- a/test/test-manager/src/tests/relay_ip_overrides.rs
+++ b/test/test-manager/src/tests/relay_ip_overrides.rs
@@ -298,12 +298,7 @@ async fn pick_a_relay(
 
     let relay_ip = relay.ipv4_addr_in;
     let hostname = relay.hostname.clone();
-    let city = relay
-        .location
-        .as_ref()
-        .ok_or(anyhow!("Got Relay with an unknown location"))?
-        .city_code
-        .clone();
+    let city = relay.location.city_code.clone();
 
     log::info!("selected {hostname} ({relay_ip})");
     let location = GeographicLocationConstraint::Hostname(country, city, hostname.clone()).into();
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index 575339c5a8..6fb57adc14 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -390,38 +390,6 @@ pub async fn test_wireguard_autoconnect(
     Ok(())
 }
 
-/// Test connecting to a WireGuard relay using DAITA.
-///
-/// # Limitations
-///
-/// The test does not analyze any traffic, nor verify that DAITA is in use.
-#[test_function]
-pub async fn test_daita(
-    _: TestContext,
-    rpc: ServiceClient,
-    mut mullvad_client: MullvadProxyClient,
-) -> anyhow::Result<()> {
-    log::info!("Connecting to relay with DAITA");
-
-    apply_settings_from_relay_query(
-        &mut mullvad_client,
-        RelayQueryBuilder::new().wireguard().build(),
-    )
-    .await?;
-
-    mullvad_client
-        .set_daita_settings(wireguard::DaitaSettings { enabled: true })
-        .await
-        .context("Failed to enable daita")?;
-
-    connect_and_wait(&mut mullvad_client).await?;
-
-    log::info!("Check that the connection works");
-    let _ = helpers::geoip_lookup_with_retries(&rpc).await?;
-
-    Ok(())
-}
-
 /// Test whether the daemon automatically connects on reboot when using
 /// OpenVPN.
 ///