From 5546da8785d01f28dfa599a526fa661a70599b1c Mon Sep 17 00:00:00 2001
From: Markus Pettersson <markus.pettersson@mullvad.net>
Date: Thu, 17 Oct 2024 09:52:18 +0200
Subject: [PATCH] Add test case for smart routing overriding multihop

---
 .../src/relay_selector/query.rs               |  4 +-
 .../tests/relay_selector.rs                   | 74 +++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index a066d9859888..a14a5a2a1b16 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -932,7 +932,9 @@ impl RelayQueryExtensions for RelayQuery {
     fn use_multihop_if_necessary(&self) -> bool {
         self.wireguard_constraints()
             .daita_use_multihop_if_necessary
-            .is_only_and(|enabled| enabled)
+            // The default value is `Any`, which means that we need to check the intersection.
+            .intersection(Constraint::Only(true))
+            .is_some()
     }
     fn singlehop(&self) -> bool {
         !self.wireguard_constraints().multihop()
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 1ce69b0591fb..f7d5529623ce 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -95,6 +95,27 @@ static RELAYS: LazyLock<RelayList> = LazyLock::new(|| RelayList {
                     }),
                     location: DUMMY_LOCATION.clone(),
                 },
+                Relay {
+                    hostname: "se11-wireguard".to_string(),
+                    ipv4_addr_in: "185.213.154.69".parse().unwrap(),
+                    ipv6_addr_in: Some("2a03:1b20:5:f011::a11f".parse().unwrap()),
+                    overridden_ipv4: false,
+                    overridden_ipv6: false,
+                    include_in_country: true,
+                    active: true,
+                    owned: false,
+                    provider: "provider2".to_string(),
+                    weight: 1,
+                    endpoint_data: RelayEndpointData::Wireguard(WireguardRelayEndpointData {
+                        public_key: PublicKey::from_base64(
+                            "BLNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=",
+                        )
+                        .unwrap(),
+                        daita: true,
+                        shadowsocks_extra_addr_in: vec![],
+                    }),
+                    location: DUMMY_LOCATION.clone(),
+                },
                 Relay {
                     hostname: "se-got-001".to_string(),
                     ipv4_addr_in: "185.213.154.131".parse().unwrap(),
@@ -1401,6 +1422,59 @@ fn test_daita_any_tunnel_protocol() {
     );
 }
 
+/// Always use smart routing to select a DAITA-enabled entry relay if both smart routing and multihop is enabled.
+/// This applies even if the entry is set explicitly.
+/// DAITA is a core privacy feature
+#[test]
+fn test_daita_smart_routing_overrides_multihop() {
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), RELAYS.clone());
+    let query = RelayQueryBuilder::new()
+        .wireguard()
+        .daita()
+        .daita_use_multihop_if_necessary(true)
+        .multihop()
+        // Set the entry to a relay that explicitly does *not* support DAITA.
+        // Later, we check that the smart routing diregards this choice and selects a DAITA-enabled
+        // relay instead.
+        .entry(NON_DAITA_RELAY_LOCATION.clone())
+        .build();
+
+    for _ in 0..100 {
+        // Make sure a DAITA-enabled relay is always selected due to smart routing.
+        let relay = relay_selector
+            .get_relay_by_query(query.clone())
+            .expect("Expected to find a relay with daita_use_multihop_if_necessary");
+        match relay {
+                GetRelay::Wireguard {
+                    inner: WireguardConfig::Multihop { entry, exit: _ },
+                    ..
+                } => {
+                    assert!(supports_daita(&entry), "entry relay must support DAITA");
+                }
+                wrong_relay => panic!(
+                "Relay selector should have picked two Wireguard relays, instead chose {wrong_relay:?}"
+            ),
+            }
+    }
+
+    // Assert that disabling smart routing for this query will fail to generate a valid multihop
+    // config, thus blocking the user.
+    let query = RelayQueryBuilder::new()
+        .wireguard()
+        .daita()
+        .daita_use_multihop_if_necessary(false)
+        .multihop()
+        .entry(NON_DAITA_RELAY_LOCATION.clone())
+        .build();
+
+    let relay = relay_selector.get_relay_by_query(query);
+
+    assert!(
+        relay.is_err(),
+        "expected there to be no valid multihop configuration! Instead got {relay:#?}"
+    );
+}
+
 /// Always select a WireGuard relay when multihop is enabled
 /// Multihop is a core privacy feature
 #[test]