Merge branch 'make-sure-to-use-a-daita-compatible-entry-even-if-the-exit-des-1303'

9 files changed, 273 insertions, 113 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4a08634efd..9c9c19736d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,8 @@ Line wrap the file at 100 chars.                                              Th
 ### Changed
 - Replace the draft key encapsulation mechanism Kyber (round 3) with the standardized
   ML-KEM (FIPS 203) dito in the handshake for Quantum-resistant tunnels.
+- Make Smart Routing override multihop if both are enabled. To manually set the entry relay,
+  explicitly enable the "Direct only" option in the DAITA settings.
 
 #### Windows
 - Enable quantum-resistant tunnels by default (when set to `auto`).
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index dd6c828802..08c523ffd7 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -2403,10 +2403,9 @@ impl Daemon {
                     .intersection(Constraint::Only(TunnelType::Wireguard))
                     .is_some();
 
-                let multihop_enabled = constraints.wireguard_constraints.use_multihop;
                 let daita_enabled = self.settings.tunnel_options.wireguard.daita.enabled;
 
-                if settings_changed && wireguard_enabled && daita_enabled && !multihop_enabled {
+                if settings_changed && wireguard_enabled && daita_enabled {
                     log::info!("Reconnecting because DAITA settings changed");
                     self.reconnect_tunnel();
                 }
diff --git a/mullvad-relay-selector/src/error.rs b/mullvad-relay-selector/src/error.rs
index dea6f81c3e..9f7a50bda3 100644
--- a/mullvad-relay-selector/src/error.rs
+++ b/mullvad-relay-selector/src/error.rs
@@ -3,7 +3,7 @@
 
 use mullvad_types::{relay_constraints::MissingCustomBridgeSettings, relay_list::Relay};
 
-use crate::{detailer, WireguardConfig};
+use crate::{detailer, relay_selector::relays::WireguardConfig};
 
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 09708a059e..be12c29443 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, AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay,
-    RelaySelector, RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig,
-    WireguardConfig, RETRY_ORDER,
+    detailer, query, relays::WireguardConfig, AdditionalRelayConstraints,
+    AdditionalWireguardConstraints, GetRelay, RelaySelector, RuntimeParameters, SelectedBridge,
+    SelectedObfuscator, SelectorConfig, RETRY_ORDER,
 };
diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index c8b529e0cd..a8a2ba7fcb 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -207,8 +207,8 @@ fn port_if_in_range<R: RangeBounds<u16>>(port_ranges: &[R], port: u16) -> Result
 /// - `port_ranges`: A slice of port numbers.
 ///
 /// # Returns
-/// - On success, a randomly selected port number within the given ranges. Otherwise,
-///   an error is returned.
+/// - On success, a randomly selected port number within the given ranges. Otherwise, an error is
+///   returned.
 pub fn select_random_port<R: RangeBounds<u16> + Iterator<Item = u16> + Clone>(
     port_ranges: &[R],
 ) -> Result<u16, Error> {
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index b94a5812a0..e7a6922727 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -5,16 +5,30 @@ mod helpers;
 mod matcher;
 mod parsed_relays;
 pub mod query;
+pub mod relays;
+
+use matcher::{filter_matching_bridges, filter_matching_relay_list};
+use parsed_relays::ParsedRelays;
+use relays::{Multihop, Singlehop, WireguardConfig};
+
+use crate::{
+    detailer::{openvpn_endpoint, wireguard_endpoint},
+    error::{EndpointErrorDetails, Error},
+    query::{
+        BridgeQuery, ObfuscationQuery, OpenVpnRelayQuery, RelayQuery, RelayQueryExt,
+        WireguardRelayQuery,
+    },
+};
 
-use chrono::{DateTime, Local};
-use itertools::Itertools;
-use query::ObfuscationQuery;
 use std::{
     path::Path,
     sync::{Arc, LazyLock, Mutex},
     time::SystemTime,
 };
 
+use chrono::{DateTime, Local};
+use itertools::Itertools;
+
 use mullvad_types::{
     constraints::Constraint,
     custom_list::CustomListsSettings,
@@ -37,15 +51,6 @@ use talpid_types::{
     ErrorExt,
 };
 
-use crate::error::{EndpointErrorDetails, Error};
-
-use self::{
-    detailer::{openvpn_endpoint, wireguard_endpoint},
-    matcher::{filter_matching_bridges, filter_matching_relay_list},
-    parsed_relays::ParsedRelays,
-    query::{BridgeQuery, OpenVpnRelayQuery, RelayQuery, WireguardRelayQuery},
-};
-
 /// [`RETRY_ORDER`] defines an ordered set of relay parameters which the relay selector should
 /// prioritize on successive connection attempts. Note that these will *never* override user
 /// preferences. See [the documentation on `RelayQuery`][RelayQuery] for further details.
@@ -221,49 +226,6 @@ pub enum GetRelay {
     Custom(CustomTunnelEndpoint),
 }
 
-/// This struct defines the different Wireguard relays the the relay selector can end up selecting
-/// for an arbitrary Wireguard [`query`].
-///
-/// - [`WireguardConfig::Singlehop`]; A normal wireguard relay where VPN traffic enters and exits
-///   through this sole relay.
-/// - [`WireguardConfig::Multihop`]; Two wireguard relays to be used in a multihop circuit. VPN
-///   traffic will enter through `entry` and eventually come out from `exit` before the traffic will
-///   actually be routed to the broader internet.
-#[derive(Clone, Debug)]
-pub enum WireguardConfig {
-    /// Strongly prefer to instantiate this variant using [`WireguardConfig::singlehop`] as that
-    /// will assert that the relay is of the expected type.
-    Singlehop { exit: Relay },
-    /// Strongly prefer to instantiate this variant using [`WireguardConfig::multihop`] as that
-    /// will assert that the entry & exit relays are of the expected type.
-    Multihop { exit: Relay, entry: Relay },
-}
-
-impl WireguardConfig {
-    const fn singlehop(exit: Relay) -> Self {
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            exit.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        Self::Singlehop { exit }
-    }
-
-    const fn multihop(exit: Relay, entry: Relay) -> Self {
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            exit.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            entry.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        Self::Multihop { exit, entry }
-    }
-}
-
 #[derive(Clone, Debug)]
 pub enum SelectedBridge {
     Normal { settings: CustomProxy, relay: Relay },
@@ -675,7 +637,7 @@ impl RelaySelector {
         Self::get_wireguard_relay(&query, custom_lists, parsed_relays)
     }
 
-    /// Derive a valid Wireguard relay configuration from `query`.
+    /// Derive a valid relay configuration from `query`.
     ///
     /// # Note
     /// This function should *only* be called with a Wireguard query.
@@ -697,11 +659,7 @@ impl RelaySelector {
             query.tunnel_protocol(),
             Constraint::Only(TunnelType::Wireguard)
         );
-        let inner = if !query.wireguard_constraints().multihop() {
-            Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays)?
-        } else {
-            Self::get_wireguard_multihop_config(query, custom_lists, parsed_relays)?
-        };
+        let inner = Self::get_wireguard_relay_config(query, custom_lists, parsed_relays)?;
         let endpoint = Self::get_wireguard_endpoint(query, parsed_relays, &inner)?;
         let obfuscator =
             Self::get_wireguard_obfuscator(query, inner.clone(), &endpoint, parsed_relays)?;
@@ -713,67 +671,85 @@ impl RelaySelector {
         })
     }
 
-    /// Select a valid Wireguard exit relay.
+    /// Derive a valid Wireguard relay configuration from `query`.
     ///
     /// # Returns
     /// * An `Err` if no exit relay can be chosen
-    /// * `Ok(WireguardInner::Singlehop)` otherwise
-    fn get_wireguard_singlehop_config(
+    /// * An `Err` if no entry relay can be chosen (if multihop is enabled on `query`)
+    /// * `Ok(WireguardConfig)` otherwise
+    fn get_wireguard_relay_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<WireguardConfig, Error> {
-        let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
-
-        // are we using daita?
-        let using_daita = || query.wireguard_constraints().daita == Constraint::Only(true);
-
-        // is the `candidates` list empty because DAITA is enabled?
-        let no_relay_because_daita = || {
-            let mut query = query.clone();
-            let mut wireguard_constraints = query.wireguard_constraints().clone();
-            wireguard_constraints.daita = Constraint::Any;
-            query.set_wireguard_constraints(wireguard_constraints)?;
-            let candidates = filter_matching_relay_list(&query, parsed_relays, custom_lists);
-            Result::<_, Error>::Ok(!candidates.is_empty())
-        };
-
-        // is `use_multihop_if_necessary` enabled?
-        let use_multihop_if_necessary = || {
-            query
-                .wireguard_constraints()
-                .daita_use_multihop_if_necessary
-                .intersection(Constraint::Only(true))
-                .is_some()
+        // TODO: Remove when Android gets support for multihop.
+        if cfg!(target_os = "android") {
+            let relay = Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays)
+                .ok_or(Error::NoRelay)?;
+            return Ok(WireguardConfig::from(relay));
+        }
+        let inner = if query.singlehop() {
+            match Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays) {
+                Some(exit) => WireguardConfig::from(exit),
+                None => {
+                    // If we found no matching relays because DAITA was enabled, and
+                    // `use_multihop_if_necessary` is enabled, try enabling
+                    // multihop and connecting using an automatically selected
+                    // entry relay.
+                    if query.using_daita() && query.use_multihop_if_necessary() {
+                        let multihop = Self::get_wireguard_auto_multihop_config(
+                            query,
+                            custom_lists,
+                            parsed_relays,
+                        )?;
+                        WireguardConfig::from(multihop)
+                    } else {
+                        return Err(Error::NoRelay);
+                    }
+                }
+            }
+        } else {
+            // A DAITA compatible entry should be used even when the exit is DAITA compatible.
+            // This only makes sense in context: The user is no longer able to explicitly choose an
+            // entry relay with smarting routing enabled, even if multihop is turned on
+            // Also implied: Multihop is enabled.
+            let multihop = if query.using_daita() && query.use_multihop_if_necessary() {
+                Self::get_wireguard_auto_multihop_config(query, custom_lists, parsed_relays)?
+            } else {
+                Self::get_wireguard_multihop_config(query, custom_lists, parsed_relays)?
+            };
+            WireguardConfig::from(multihop)
         };
 
-        // if we found no matching relays because DAITA was enabled, and `use_multihop_if_necessary`
-        // is enabled, try enabling multihop and connecting using an automatically selected
-        // entry relay.
-        if candidates.is_empty()
-            && using_daita()
-            && no_relay_because_daita()?
-            && use_multihop_if_necessary()
-        {
-            return Self::get_wireguard_auto_multihop_config(query, custom_lists, parsed_relays);
-        }
+        Ok(inner)
+    }
 
+    /// Select a valid Wireguard exit relay.
+    ///
+    /// # Returns
+    /// * `Ok(Singlehop)` if an exit relay was selected
+    /// * `None` otherwise
+    fn get_wireguard_singlehop_config(
+        query: &RelayQuery,
+        custom_lists: &CustomListsSettings,
+        parsed_relays: &ParsedRelays,
+    ) -> Option<Singlehop> {
+        let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
         helpers::pick_random_relay(&candidates)
             .cloned()
-            .map(WireguardConfig::singlehop)
-            .ok_or(Error::NoRelay)
+            .map(Singlehop::new)
     }
 
     /// Select a valid Wireguard exit relay, together with with an automatically chosen entry relay.
     ///
     /// # Returns
     /// * An `Err` if no entry/exit relay can be chosen
-    /// * `Ok(WireguardInner::Multihop)` otherwise
+    /// * `Ok(Multihop)` otherwise
     fn get_wireguard_auto_multihop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
-    ) -> Result<WireguardConfig, Error> {
+    ) -> Result<Multihop, Error> {
         let mut exit_relay_query = query.clone();
 
         // DAITA should only be enabled for the entry relay
@@ -810,7 +786,7 @@ impl RelaySelector {
         let entry =
             helpers::pick_random_relay_excluding(&entry_candidates, exit).ok_or(Error::NoRelay)?;
 
-        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
+        Ok(Multihop::new(entry.clone(), exit.clone()))
     }
 
     /// This function selects a valid entry and exit relay to be used in a multihop configuration.
@@ -819,14 +795,14 @@ impl RelaySelector {
     /// * An `Err` if no exit relay can be chosen
     /// * An `Err` if no entry relay can be chosen
     /// * An `Err` if the chosen entry and exit relays are the same
-    /// * `Ok(WireguardInner::Multihop)` otherwise
+    /// * `Ok(WireguardConfig::Multihop)` otherwise
     fn get_wireguard_multihop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
-    ) -> Result<WireguardConfig, Error> {
+    ) -> Result<Multihop, Error> {
         // Here, we modify the original query just a bit.
-        // The actual query for an exit relay is identical as for an exit relay, with the
+        // The actual query for an entry relay is identical as for an exit relay, with the
         // exception that the location is different. It is simply the location as dictated by
         // the query's multihop constraint.
         let mut entry_relay_query = query.clone();
@@ -864,7 +840,7 @@ impl RelaySelector {
         }
         .ok_or(Error::NoRelay)?;
 
-        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
+        Ok(Multihop::new(entry.clone(), exit.clone()))
     }
 
     /// Constructs a [`MullvadEndpoint`] with details for how to connect to `relay`.
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index 570a2212cb..162865368a 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -913,6 +913,34 @@ pub mod builder {
     }
 }
 
+/// This trait defines a bunch of helper methods on [`RelayQuery`].
+pub trait RelayQueryExt {
+    /// Are we using daita?
+    fn using_daita(&self) -> bool;
+    /// is `use_multihop_if_necessary` enabled? In other words, is `Direct only` disabled?
+    fn use_multihop_if_necessary(&self) -> bool;
+    /// Are we using singlehop? I.e. is multihop *not* explicitly enabled?
+    fn singlehop(&self) -> bool;
+}
+
+impl RelayQueryExt for RelayQuery {
+    fn using_daita(&self) -> bool {
+        self.wireguard_constraints()
+            .daita
+            .is_only_and(|enabled| enabled)
+    }
+    fn use_multihop_if_necessary(&self) -> bool {
+        self.wireguard_constraints()
+            .daita_use_multihop_if_necessary
+            // 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()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use mullvad_types::{
diff --git a/mullvad-relay-selector/src/relay_selector/relays.rs b/mullvad-relay-selector/src/relay_selector/relays.rs
new file mode 100644
index 0000000000..3cb8ff3c5d
--- /dev/null
+++ b/mullvad-relay-selector/src/relay_selector/relays.rs
@@ -0,0 +1,81 @@
+//! This module defines wrapper types around [`Relay`], often to provide certain runtime guarantees
+//! or disambiguate the type of relay which is used in the relay selector's internal APIs.
+
+use mullvad_types::relay_list::{Relay, RelayEndpointData};
+
+/// - [`WireguardConfig::Singlehop`]: A wireguard relay where VPN traffic enters and exits.
+/// - [`WireguardConfig::Multihop`]: Two wireguard relays to be used in a multihop circuit. VPN
+///   traffic will enter through `entry` and eventually exit through `exit` before the traffic will
+///   actually be routed to the internet.
+#[derive(Clone, Debug)]
+pub enum WireguardConfig {
+    /// An exit relay.
+    Singlehop { exit: Relay },
+    /// An entry and an exit relay.
+    Multihop { exit: Relay, entry: Relay },
+}
+
+/// A type representing single Wireguard relay.
+///
+/// Before you can read any data out of a [`Singlehop`] value uou need to convert it to
+/// [`WireguardConfig`]. This is easy since [`Singlehop`] implements [`Into<WireguardConfig>`].
+///
+/// # Why not simply use [`Relay`]?
+/// The only way to construct a [`Singlehop`] value is with [`Singlehop::new`] which performs
+/// additional validation which guarantees that the relay actually is a Wireguard relay, while
+/// [`Relay`] is not guaranteed to be a Wireguard relay.
+pub struct Singlehop(Relay);
+/// A type representing two Wireguard relay - an entry and an exit.
+///
+/// Before you can read any data out of a [`Multihop`] value uou need to convert it to
+/// [`WireguardConfig`]. This is easy since [`Multihop`] implements [`Into<WireguardConfig>`].
+///
+/// # Why not simply use [`Relay`]?
+/// The same rationale as for [`Singlehop`] applies - [`Multihop::new`] performs additional
+/// validation on the entry and exit relays.
+pub struct Multihop {
+    entry: Relay,
+    exit: Relay,
+}
+
+impl From<Singlehop> for WireguardConfig {
+    fn from(relay: Singlehop) -> Self {
+        Self::Singlehop { exit: relay.0 }
+    }
+}
+
+impl From<Multihop> for WireguardConfig {
+    fn from(relay: Multihop) -> Self {
+        WireguardConfig::Multihop {
+            exit: relay.exit,
+            entry: relay.entry,
+        }
+    }
+}
+
+impl Singlehop {
+    pub const fn new(exit: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Self(exit)
+    }
+}
+
+impl Multihop {
+    pub const fn new(entry: Relay, exit: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            entry.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Multihop { exit, entry }
+    }
+}
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 1ce69b0591..6908d468de 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -96,6 +96,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(),
                     ipv6_addr_in: None,
@@ -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 disregards 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]