Ignore obfuscation protocol specific constraints when 'auto' is selected

5 files changed, 96 insertions, 50 deletions

diff --git a/mullvad-api/src/lib.rs b/mullvad-api/src/lib.rs
index 833dcf3f67..87b6e3d656 100644
--- a/mullvad-api/src/lib.rs
+++ b/mullvad-api/src/lib.rs
@@ -139,8 +139,7 @@ pub struct ApiEndpoint {
     /// Whether bridges/proxies can be used to access the API or not. This is
     /// useful primarily for testing purposes.
     ///
-    /// * If `force_direct` is `true`, bridges and proxies will not be used to
-    ///   reach the API.
+    /// * If `force_direct` is `true`, bridges and proxies will not be used to reach the API.
     /// * If `force_direct` is `false`, bridges and proxies can be used to reach the API.
     ///
     /// # Note
diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index 48c79f05d1..442b1f596f 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -58,7 +58,7 @@ pub fn pick_random_relay_weighted<RelayType>(
 }
 
 pub fn get_udp2tcp_obfuscator(
-    obfuscation_settings_constraint: &Constraint<Udp2TcpObfuscationSettings>,
+    obfuscation_settings_constraint: &Udp2TcpObfuscationSettings,
     udp2tcp_ports: &[u16],
     relay: Relay,
     endpoint: &MullvadWireguardEndpoint,
@@ -73,17 +73,16 @@ pub fn get_udp2tcp_obfuscator(
 }
 
 pub fn get_udp2tcp_obfuscator_port(
-    obfuscation_settings_constraint: &Constraint<Udp2TcpObfuscationSettings>,
+    obfuscation_settings: &Udp2TcpObfuscationSettings,
     udp2tcp_ports: &[u16],
 ) -> Option<u16> {
-    match obfuscation_settings_constraint {
-        Constraint::Only(obfuscation_settings) if obfuscation_settings.port.is_only() => {
-            udp2tcp_ports
-                .iter()
-                .find(|&candidate| obfuscation_settings.port == Constraint::Only(*candidate))
-                .copied()
-        }
+    if let Constraint::Only(desired_port) = obfuscation_settings.port {
+        udp2tcp_ports
+            .iter()
+            .find(|&candidate| desired_port == *candidate)
+            .copied()
+    } else {
         // There are no specific obfuscation settings to take into consideration in this case.
-        Constraint::Any | Constraint::Only(_) => udp2tcp_ports.choose(&mut thread_rng()).copied(),
+        udp2tcp_ports.choose(&mut thread_rng()).copied()
     }
 }
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 5093ffdeca..7e6acab43a 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -9,6 +9,7 @@ pub mod query;
 use chrono::{DateTime, Local};
 use itertools::Itertools;
 use once_cell::sync::Lazy;
+use query::ObfuscationQuery;
 use rand::{seq::IteratorRandom, thread_rng};
 use std::{
     path::Path,
@@ -24,7 +25,7 @@ use mullvad_types::{
     relay_constraints::{
         BridgeSettings, BridgeState, InternalBridgeConstraints, ObfuscationSettings,
         OpenVpnConstraints, RelayConstraints, RelayOverride, RelaySettings, ResolvedBridgeSettings,
-        SelectedObfuscation, WireguardConstraints,
+        WireguardConstraints,
     },
     relay_list::{Relay, RelayEndpointData, RelayList},
     settings::Settings,
@@ -342,8 +343,7 @@ impl<'a> From<NormalSelectorConfig<'a>> for RelayQuery {
                 ip_version,
                 use_multihop: Constraint::Only(use_multihop),
                 entry_location,
-                obfuscation: obfuscation_settings.selected_obfuscation,
-                udp2tcp_port: Constraint::Only(obfuscation_settings.udp2tcp.clone()),
+                obfuscation: ObfuscationQuery::from(obfuscation_settings),
                 daita: Constraint::Only(daita),
             }
         }
@@ -790,23 +790,18 @@ impl RelaySelector {
         endpoint: &MullvadWireguardEndpoint,
         parsed_relays: &ParsedRelays,
     ) -> Result<Option<SelectedObfuscator>, Error> {
-        match query.wireguard_constraints.obfuscation {
-            SelectedObfuscation::Off | SelectedObfuscation::Auto => Ok(None),
-            SelectedObfuscation::Udp2Tcp => {
+        match &query.wireguard_constraints.obfuscation {
+            ObfuscationQuery::Off | ObfuscationQuery::Auto => Ok(None),
+            ObfuscationQuery::Udp2tcp { port } => {
                 let obfuscator_relay = match relay {
                     WireguardConfig::Singlehop { exit } => exit,
                     WireguardConfig::Multihop { entry, .. } => entry,
                 };
                 let udp2tcp_ports = &parsed_relays.parsed_list().wireguard.udp2tcp_ports;
 
-                helpers::get_udp2tcp_obfuscator(
-                    &query.wireguard_constraints.udp2tcp_port,
-                    udp2tcp_ports,
-                    obfuscator_relay,
-                    endpoint,
-                )
-                .map(Some)
-                .ok_or(Error::NoObfuscator)
+                helpers::get_udp2tcp_obfuscator(port, udp2tcp_ports, obfuscator_relay, endpoint)
+                    .map(Some)
+                    .ok_or(Error::NoObfuscator)
             }
         }
     }
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index 220f232303..48ad5bd1dc 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -32,8 +32,8 @@ use crate::AdditionalWireguardConstraints;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
-        BridgeConstraints, LocationConstraint, OpenVpnConstraints, Ownership, Providers,
-        RelayConstraints, RelaySettings, SelectedObfuscation, TransportPort,
+        BridgeConstraints, LocationConstraint, ObfuscationSettings, OpenVpnConstraints, Ownership,
+        Providers, RelayConstraints, RelaySettings, SelectedObfuscation, TransportPort,
         Udp2TcpObfuscationSettings, WireguardConstraints,
     },
     Intersection,
@@ -150,11 +150,50 @@ pub struct WireguardRelayQuery {
     pub ip_version: Constraint<IpVersion>,
     pub use_multihop: Constraint<bool>,
     pub entry_location: Constraint<LocationConstraint>,
-    pub obfuscation: SelectedObfuscation,
-    pub udp2tcp_port: Constraint<Udp2TcpObfuscationSettings>,
+    pub obfuscation: ObfuscationQuery,
     pub daita: Constraint<bool>,
 }
 
+#[derive(Default, Debug, Clone, Eq, PartialEq)]
+pub enum ObfuscationQuery {
+    Off,
+    #[default]
+    Auto,
+    Udp2tcp {
+        port: Udp2TcpObfuscationSettings,
+    },
+}
+
+impl From<ObfuscationSettings> for ObfuscationQuery {
+    /// A query for obfuscation settings.
+    ///
+    /// Note that this drops obfuscation protocol specific constraints from [`ObfuscationSettings`]
+    /// when the selected obfuscation type is auto.
+    fn from(obfuscation: ObfuscationSettings) -> Self {
+        match obfuscation.selected_obfuscation {
+            SelectedObfuscation::Off => ObfuscationQuery::Off,
+            SelectedObfuscation::Auto => ObfuscationQuery::Auto,
+            SelectedObfuscation::Udp2Tcp => ObfuscationQuery::Udp2tcp {
+                port: obfuscation.udp2tcp,
+            },
+        }
+    }
+}
+
+impl Intersection for ObfuscationQuery {
+    fn intersection(self, other: Self) -> Option<Self> {
+        match (self, other) {
+            (ObfuscationQuery::Off, _) | (_, ObfuscationQuery::Off) => Some(ObfuscationQuery::Off),
+            (ObfuscationQuery::Auto, other) | (other, ObfuscationQuery::Auto) => Some(other),
+            (ObfuscationQuery::Udp2tcp { port: a }, ObfuscationQuery::Udp2tcp { port: b }) => {
+                Some(ObfuscationQuery::Udp2tcp {
+                    port: a.intersection(b)?,
+                })
+            }
+        }
+    }
+}
+
 impl WireguardRelayQuery {
     pub fn multihop(&self) -> bool {
         matches!(self.use_multihop, Constraint::Only(true))
@@ -168,8 +207,7 @@ impl WireguardRelayQuery {
             ip_version: Constraint::Any,
             use_multihop: Constraint::Any,
             entry_location: Constraint::Any,
-            obfuscation: SelectedObfuscation::Auto,
-            udp2tcp_port: Constraint::Any,
+            obfuscation: ObfuscationQuery::Auto,
             daita: Constraint::Any,
         }
     }
@@ -310,7 +348,7 @@ pub mod builder {
     };
     use talpid_types::net::TunnelType;
 
-    use super::{BridgeQuery, RelayQuery};
+    use super::{BridgeQuery, ObfuscationQuery, RelayQuery};
 
     // Re-exports
     pub use mullvad_types::relay_constraints::{
@@ -505,8 +543,8 @@ pub mod builder {
                 obfuscation: obfuscation.clone(),
                 daita: self.protocol.daita,
             };
-            self.query.wireguard_constraints.udp2tcp_port = Constraint::Only(obfuscation);
-            self.query.wireguard_constraints.obfuscation = SelectedObfuscation::Udp2Tcp;
+            self.query.wireguard_constraints.obfuscation =
+                ObfuscationQuery::Udp2tcp { port: obfuscation };
             RelayQueryBuilder {
                 query: self.query,
                 protocol,
@@ -519,8 +557,9 @@ pub mod builder {
         /// protocol should use to connect to a relay.
         pub fn udp2tcp_port(mut self, port: u16) -> Self {
             self.protocol.obfuscation.port = Constraint::Only(port);
-            self.query.wireguard_constraints.udp2tcp_port =
-                Constraint::Only(self.protocol.obfuscation.clone());
+            self.query.wireguard_constraints.obfuscation = ObfuscationQuery::Udp2tcp {
+                port: self.protocol.obfuscation.clone(),
+            };
             self
         }
     }
@@ -639,10 +678,13 @@ pub mod builder {
 
 #[cfg(test)]
 mod test {
-    use mullvad_types::constraints::Constraint;
+    use mullvad_types::{
+        constraints::Constraint,
+        relay_constraints::{ObfuscationSettings, SelectedObfuscation, Udp2TcpObfuscationSettings},
+    };
     use proptest::prelude::*;
 
-    use super::Intersection;
+    use super::{Intersection, ObfuscationQuery};
 
     // Define proptest combinators for the `Constraint` type.
 
@@ -719,5 +761,18 @@ mod test {
             };
             prop_assert_eq!(left, right);
         }
+
+        /// When obfuscation is set to automatic in [`ObfuscationSettings`], the query should not
+        /// contain any specific obfuscation protocol settings.
+        #[test]
+        fn test_auto_obfuscation_settings(port in constraint(proptest::arbitrary::any::<u16>())) {
+            let query = ObfuscationQuery::from(ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Auto,
+                udp2tcp: Udp2TcpObfuscationSettings {
+                    port,
+                },
+            });
+            assert_eq!(query, ObfuscationQuery::Auto);
+        }
     }
 }
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 7cf2329e9a..053de20c1a 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -11,7 +11,7 @@ use talpid_types::net::{
 };
 
 use mullvad_relay_selector::{
-    query::{builder::RelayQueryBuilder, BridgeQuery, OpenVpnRelayQuery},
+    query::{builder::RelayQueryBuilder, BridgeQuery, ObfuscationQuery, OpenVpnRelayQuery},
     Error, GetRelay, RelaySelector, RuntimeParameters, SelectorConfig, WireguardConfig,
     RETRY_ORDER,
 };
@@ -20,7 +20,7 @@ use mullvad_types::{
     endpoint::MullvadEndpoint,
     relay_constraints::{
         BridgeConstraints, BridgeState, GeographicLocationConstraint, Ownership, Providers,
-        SelectedObfuscation, TransportPort,
+        TransportPort,
     },
     relay_list::{
         BridgeEndpointData, OpenVpnEndpoint, OpenVpnEndpointData, Relay, RelayEndpointData,
@@ -311,10 +311,10 @@ fn test_retry_order() {
                     .wireguard_constraints
                     .port
                     .matches_eq(&endpoint.peer.endpoint.port()));
-                assert!(match query.wireguard_constraints.obfuscation {
-                    SelectedObfuscation::Auto => true,
-                    SelectedObfuscation::Off => obfuscator.is_none(),
-                    SelectedObfuscation::Udp2Tcp => obfuscator.is_some(),
+                assert!(match &query.wireguard_constraints.obfuscation {
+                    ObfuscationQuery::Auto => true,
+                    ObfuscationQuery::Off => obfuscator.is_none(),
+                    ObfuscationQuery::Udp2tcp { .. } => obfuscator.is_some(),
                 });
             }
             GetRelay::OpenVpn {
@@ -737,15 +737,13 @@ fn test_selecting_wireguard_endpoint_with_udp2tcp_obfuscation() {
 /// multihop is explicitly turned off. Assert that the relay selector does *not* return an
 /// obfuscator config.
 ///
-/// # Note
-/// This is a highly specific test which details how the relay selector should behave at the time of
-/// writing this test. The cost (in latency primarily) of using obfuscation is deemed to be too high
-/// to enable it as an auto-configuration.
+/// [`RelaySelector::get_relay`] may still enable obfuscation if it is present in [`RETRY_ORDER`].
 #[test]
 fn test_selecting_wireguard_endpoint_with_auto_obfuscation() {
     let relay_selector = default_relay_selector();
+
     let mut query = RelayQueryBuilder::new().wireguard().build();
-    query.wireguard_constraints.obfuscation = SelectedObfuscation::Auto;
+    query.wireguard_constraints.obfuscation = ObfuscationQuery::Auto;
 
     for _ in 0..100 {
         let relay = relay_selector.get_relay_by_query(query.clone()).unwrap();