Add DAITA to relay selection

8 files changed, 87 insertions, 40 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 2733482a4c..d0fad0493d 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -37,7 +37,9 @@ use futures::{
     StreamExt,
 };
 use geoip::GeoIpHandler;
-use mullvad_relay_selector::{AdditionalRelayConstraints, AdditionalWireguardConstraints, RelaySelector, SelectorConfig};
+use mullvad_relay_selector::{
+    AdditionalRelayConstraints, AdditionalWireguardConstraints, RelaySelector, SelectorConfig,
+};
 #[cfg(target_os = "android")]
 use mullvad_types::account::{PlayPurchase, PlayPurchasePaymentToken};
 #[cfg(target_os = "windows")]
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 9ac49d0b1f..73c709a844 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -10,6 +10,7 @@ mod relay_selector;
 pub use error::Error;
 pub use relay_selector::detailer;
 pub use relay_selector::{
-    query, GetRelay, RelaySelector, RuntimeParameters, SelectedBridge, SelectedObfuscator,
-    SelectorConfig, WireguardConfig, RETRY_ORDER,
+    query, AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, RelaySelector,
+    RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig, WireguardConfig,
+    RETRY_ORDER,
 };
diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index 2987d8965a..f6c244d3f2 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/mullvad-relay-selector/src/relay_selector/matcher.rs
@@ -8,7 +8,7 @@ use mullvad_types::{
         GeographicLocationConstraint, InternalBridgeConstraints, LocationConstraint, Ownership,
         Providers,
     },
-    relay_list::{Relay, RelayEndpointData},
+    relay_list::{Relay, RelayEndpointData, WireguardRelayEndpointData},
 };
 use talpid_types::net::TunnelType;
 
@@ -32,7 +32,9 @@ pub fn filter_matching_relay_list<'a, R: Iterator<Item = &'a Relay> + Clone>(
             // Filter by ownership
             .filter(|relay| filter_on_ownership(&query.ownership, relay))
             // Filter by providers
-            .filter(|relay| filter_on_providers(&query.providers, relay));
+            .filter(|relay| filter_on_providers(&query.providers, relay))
+            // Filter by DAITA support
+            .filter(|relay| filter_on_daita(&query.wireguard_constraints.daita, relay));
 
     // The last filtering to be done is on the `include_in_country` attribute found on each
     // relay. When the location constraint is based on country, a relay which has
@@ -76,7 +78,7 @@ pub fn filter_matching_bridges<'a, R: Iterator<Item = &'a Relay> + Clone>(
             .filter(|relay| filter_on_location(&locations, relay))
             // Filter by ownership
             .filter(|relay| filter_on_ownership(&constraints.ownership, relay))
-            // Filter by constraints
+            // Filter by providers
             .filter(|relay| filter_on_providers(&constraints.providers, relay))
             .cloned()
             .collect()
@@ -108,6 +110,19 @@ pub fn filter_on_providers(filter: &Constraint<Providers>, relay: &Relay) -> boo
     filter.matches(relay)
 }
 
+/// Returns whether `relay` satisfy the daita constraint posed by `filter`.
+pub fn filter_on_daita(filter: &Constraint<bool>, relay: &Relay) -> bool {
+    match (filter, &relay.endpoint_data) {
+        // Only a subset of relays support DAITA, so filter out ones that don't.
+        (
+            Constraint::Only(true),
+            RelayEndpointData::Wireguard(WireguardRelayEndpointData { daita, .. }),
+        ) => *daita,
+        // If we don't require DAITA, any relay works.
+        _ => true,
+    }
+}
+
 /// Returns whether the relay is an OpenVPN relay.
 pub const fn filter_openvpn(relay: &Relay) -> bool {
     matches!(relay.endpoint_data, RelayEndpointData::Openvpn)
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 7ec2fc02bb..3f92bd344b 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -184,6 +184,7 @@ enum SpecializedSelectorConfig<'a> {
 #[derive(Debug, Clone)]
 struct NormalSelectorConfig<'a> {
     user_preferences: &'a RelayConstraints,
+    additional_preferences: &'a AdditionalRelayConstraints,
     custom_lists: &'a CustomListsSettings,
     // Wireguard specific data
     obfuscation_settings: &'a ObfuscationSettings,
@@ -288,7 +289,7 @@ impl Default for SelectorConfig {
         let default_settings = Settings::default();
         SelectorConfig {
             relay_settings: default_settings.relay_settings,
-            additional_constraints: default_settings.additional_constraints,
+            additional_constraints: AdditionalRelayConstraints::default(),
             bridge_settings: default_settings.bridge_settings,
             obfuscation_settings: default_settings.obfuscation_settings,
             bridge_state: default_settings.bridge_state,
@@ -307,6 +308,7 @@ impl<'a> From<&'a SelectorConfig> for SpecializedSelectorConfig<'a> {
             RelaySettings::Normal(user_preferences) => {
                 SpecializedSelectorConfig::Normal(NormalSelectorConfig {
                     user_preferences,
+                    additional_preferences: &value.additional_constraints,
                     obfuscation_settings: &value.obfuscation_settings,
                     bridge_state: &value.bridge_state,
                     bridge_settings: &value.bridge_settings,
@@ -323,6 +325,7 @@ impl<'a> From<NormalSelectorConfig<'a>> for RelayQuery {
         /// Map the Wireguard-specific bits of `value` to [`WireguradRelayQuery`]
         fn wireguard_constraints(
             wireguard_constraints: WireguardConstraints,
+            additional_constraints: AdditionalWireguardConstraints,
             obfuscation_settings: ObfuscationSettings,
         ) -> WireguardRelayQuery {
             let WireguardConstraints {
@@ -331,6 +334,7 @@ impl<'a> From<NormalSelectorConfig<'a>> for RelayQuery {
                 use_multihop,
                 entry_location,
             } = wireguard_constraints;
+            let AdditionalWireguardConstraints { daita } = additional_constraints;
             WireguardRelayQuery {
                 port,
                 ip_version,
@@ -338,6 +342,7 @@ impl<'a> From<NormalSelectorConfig<'a>> for RelayQuery {
                 entry_location,
                 obfuscation: obfuscation_settings.selected_obfuscation,
                 udp2tcp_port: Constraint::Only(obfuscation_settings.udp2tcp.clone()),
+                daita: Constraint::Only(daita),
             }
         }
 
@@ -366,6 +371,7 @@ impl<'a> From<NormalSelectorConfig<'a>> for RelayQuery {
 
         let wireguard_constraints = wireguard_constraints(
             value.user_preferences.wireguard_constraints.clone(),
+            value.additional_preferences.wireguard.clone(),
             value.obfuscation_settings.clone(),
         );
         let openvpn_constraints = openvpn_constraints(
@@ -740,8 +746,14 @@ impl RelaySelector {
         // After we have our two queries (one for the exit relay & one for the entry relay),
         // we can query for all exit & entry candidates! All candidates are needed for the next
         // step.
-        let exit_candidates =
-            filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists);
+        let mut exit_relay_query = query.clone();
+        // DAITA should only be enabled for the entry relay
+        exit_relay_query.wireguard_constraints.daita = Constraint::Only(false);
+        let exit_candidates = filter_matching_relay_list(
+            &exit_relay_query,
+            parsed_relays.relays(),
+            config.custom_lists,
+        );
         let entry_candidates = filter_matching_relay_list(
             &entry_relay_query,
             parsed_relays.relays(),
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index f112e83968..ff68407051 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -28,6 +28,7 @@
 //! queries and ensure that queries are built in a type-safe manner, reducing the risk
 //! of runtime errors and improving code readability.
 
+use crate::AdditionalWireguardConstraints;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
@@ -139,6 +140,7 @@ pub struct WireguardRelayQuery {
     pub entry_location: Constraint<LocationConstraint>,
     pub obfuscation: SelectedObfuscation,
     pub udp2tcp_port: Constraint<Udp2TcpObfuscationSettings>,
+    pub daita: Constraint<bool>,
 }
 
 impl WireguardRelayQuery {
@@ -156,6 +158,7 @@ impl WireguardRelayQuery {
             entry_location: Constraint::Any,
             obfuscation: SelectedObfuscation::Auto,
             udp2tcp_port: Constraint::Any,
+            daita: Constraint::Any,
         }
     }
 }
@@ -172,6 +175,17 @@ impl From<WireguardRelayQuery> for WireguardConstraints {
     }
 }
 
+impl From<WireguardRelayQuery> for AdditionalWireguardConstraints {
+    /// The mapping from [`WireguardRelayQuery`] to [`AdditionalWireguardConstraints`].
+    fn from(value: WireguardRelayQuery) -> Self {
+        AdditionalWireguardConstraints {
+            daita: value
+                .daita
+                .unwrap_or(AdditionalWireguardConstraints::default().daita),
+        }
+    }
+}
+
 /// A query for a relay with OpenVPN-specific properties, such as `bridge_settings`.
 ///
 /// This struct may look a lot like [`OpenVpnConstraints`], and that is the point!
@@ -344,10 +358,11 @@ pub mod builder {
             }
         }
         /// Set the VPN protocol for this [`RelayQueryBuilder`] to Wireguard.
-        pub fn wireguard(mut self) -> RelayQueryBuilder<Wireguard<Any, Any>> {
+        pub fn wireguard(mut self) -> RelayQueryBuilder<Wireguard<Any, Any, Any>> {
             let protocol = Wireguard {
                 multihop: Any,
                 obfuscation: Any,
+                daita: Any,
             };
             self.query.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
             // Update the type state
@@ -381,13 +396,14 @@ pub mod builder {
     /// select entry point.
     ///
     /// [`WireguardRelayQuery`]: super::WireguardRelayQuery
-    pub struct Wireguard<Multihop, Obfuscation> {
+    pub struct Wireguard<Multihop, Obfuscation, Daita> {
         multihop: Multihop,
         obfuscation: Obfuscation,
+        daita: Daita,
     }
 
     // This impl-block is quantified over all configurations
-    impl<Multihop, Obfuscation> RelayQueryBuilder<Wireguard<Multihop, Obfuscation>> {
+    impl<Multihop, Obfuscation, Daita> RelayQueryBuilder<Wireguard<Multihop, Obfuscation, Daita>> {
         /// Specify the port to ues when connecting to the selected
         /// Wireguard relay.
         pub const fn port(mut self, port: u16) -> Self {
@@ -403,11 +419,27 @@ pub mod builder {
         }
     }
 
-    impl<Obfuscation> RelayQueryBuilder<Wireguard<Any, Obfuscation>> {
+    impl<Multihop, Obfuscation> RelayQueryBuilder<Wireguard<Multihop, Obfuscation, Any>> {
+        /// Enable DAITA support.
+        pub fn daita(mut self) -> RelayQueryBuilder<Wireguard<Multihop, Obfuscation, bool>> {
+            self.query.wireguard_constraints.daita = Constraint::Only(true);
+            // Update the type state
+            RelayQueryBuilder {
+                query: self.query,
+                protocol: Wireguard {
+                    multihop: self.protocol.multihop,
+                    obfuscation: self.protocol.obfuscation,
+                    daita: true,
+                },
+            }
+        }
+    }
+
+    impl<Obfuscation, Daita> RelayQueryBuilder<Wireguard<Any, Obfuscation, Daita>> {
         /// Enable multihop.
         ///
         /// To configure the entry relay, see [`RelayQueryBuilder::entry`].
-        pub fn multihop(mut self) -> RelayQueryBuilder<Wireguard<bool, Obfuscation>> {
+        pub fn multihop(mut self) -> RelayQueryBuilder<Wireguard<bool, Obfuscation, Daita>> {
             self.query.wireguard_constraints.use_multihop = Constraint::Only(true);
             // Update the type state
             RelayQueryBuilder {
@@ -415,12 +447,13 @@ pub mod builder {
                 protocol: Wireguard {
                     multihop: true,
                     obfuscation: self.protocol.obfuscation,
+                    daita: self.protocol.daita,
                 },
             }
         }
     }
 
-    impl<Obfuscation> RelayQueryBuilder<Wireguard<bool, Obfuscation>> {
+    impl<Obfuscation, Daita> RelayQueryBuilder<Wireguard<bool, Obfuscation, Daita>> {
         /// Set the entry location in a multihop configuration. This requires
         /// multihop to be enabled.
         pub fn entry(mut self, location: GeographicLocationConstraint) -> Self {
@@ -430,18 +463,19 @@ pub mod builder {
         }
     }
 
-    impl<Multihop> RelayQueryBuilder<Wireguard<Multihop, Any>> {
+    impl<Multihop, Daita> RelayQueryBuilder<Wireguard<Multihop, Any, Daita>> {
         /// Enable `UDP2TCP` obufscation. This will in turn enable the option to configure the
         /// `UDP2TCP` port.
         pub fn udp2tcp(
             mut self,
-        ) -> RelayQueryBuilder<Wireguard<Multihop, Udp2TcpObfuscationSettings>> {
+        ) -> RelayQueryBuilder<Wireguard<Multihop, Udp2TcpObfuscationSettings, Daita>> {
             let obfuscation = Udp2TcpObfuscationSettings {
                 port: Constraint::Any,
             };
             let protocol = Wireguard {
                 multihop: self.protocol.multihop,
                 obfuscation: obfuscation.clone(),
+                daita: self.protocol.daita,
             };
             self.query.wireguard_constraints.udp2tcp_port = Constraint::Only(obfuscation);
             self.query.wireguard_constraints.obfuscation = SelectedObfuscation::Udp2Tcp;
@@ -452,7 +486,7 @@ pub mod builder {
         }
     }
 
-    impl<Multihop> RelayQueryBuilder<Wireguard<Multihop, Udp2TcpObfuscationSettings>> {
+    impl<Multihop, Daita> RelayQueryBuilder<Wireguard<Multihop, Udp2TcpObfuscationSettings, Daita>> {
         /// Set the `UDP2TCP` port. This is the TCP port which the `UDP2TCP` obfuscation
         /// protocol should use to connect to a relay.
         pub fn udp2tcp_port(mut self, port: u16) -> Self {
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 40efae24e2..2a53afe11a 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -12,15 +12,15 @@ use talpid_types::net::{
 
 use mullvad_relay_selector::{
     query::{builder::RelayQueryBuilder, BridgeQuery, OpenVpnRelayQuery},
-    Error, GetRelay, RelaySelector, RuntimeParameters, SelectorConfig, WireguardConfig,
-    RETRY_ORDER,
+    AdditionalRelayConstraints, AdditionalWireguardConstraints, Error, GetRelay, RelaySelector,
+    RuntimeParameters, SelectorConfig, WireguardConfig, RETRY_ORDER,
 };
 use mullvad_types::{
     constraints::Constraint,
     endpoint::MullvadEndpoint,
     relay_constraints::{
-        BridgeConstraints, BridgeState, GeographicLocationConstraint, Ownership, Providers,
-        SelectedObfuscation, TransportPort,
+        BridgeConstraints, BridgeState, GeographicLocationConstraint, LocationConstraint,
+        Ownership, Providers, SelectedObfuscation, TransportPort,
     },
     relay_list::{
         BridgeEndpointData, OpenVpnEndpoint, OpenVpnEndpointData, Relay, RelayEndpointData,
diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs
index e89b1be42d..e272e794c7 100644
--- a/talpid-tunnel-config-client/src/lib.rs
+++ b/talpid-tunnel-config-client/src/lib.rs
@@ -107,23 +107,6 @@ pub async fn request_ephemeral_peer(
     enable_post_quantum: bool,
     enable_daita: bool,
 ) -> Result<EphemeralPeer, Error> {
-    request_ephemeral_peer_with_opts(
-        service_address,
-        parent_pubkey,
-        ephemeral_pubkey,
-        enable_post_quantum,
-        enable_daita,
-    )
-    .await
-}
-
-pub async fn request_ephemeral_peer_with_opts(
-    service_address: IpAddr,
-    parent_pubkey: PublicKey,
-    ephemeral_pubkey: PublicKey,
-    enable_post_quantum: bool,
-    enable_daita: bool,
-) -> Result<EphemeralPeer, Error> {
     let (pq_request, kem_secrets) = if enable_post_quantum {
         let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
         let kyber_keypair = kyber::keypair(&mut rand::thread_rng());
diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index bdef52343e..7c01538b60 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -726,7 +726,7 @@ impl WireguardMonitor {
 
         let ephemeral = tokio::time::timeout(
             timeout,
-            talpid_tunnel_config_client::request_ephemeral_peer_with_opts(
+            talpid_tunnel_config_client::request_ephemeral_peer(
                 IpAddr::from(config.ipv4_gateway),
                 config.tunnel.private_key.public_key(),
                 wg_psk_pubkey,