Merge branch 'relay-selector-bug'

2 files changed, 88 insertions, 77 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 3f92bd344b..843eea9a20 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -472,7 +472,7 @@ impl RelaySelector {
         let near_location = match specialized_config {
             SpecializedSelectorConfig::Normal(config) => {
                 let user_preferences = RelayQuery::from(config.clone());
-                Self::get_relay_midpoint(&user_preferences, parsed_relays, &config)
+                Self::get_relay_midpoint(&user_preferences, parsed_relays, config.custom_lists)
             }
             SpecializedSelectorConfig::Custom(_) => None,
         };
@@ -508,9 +508,9 @@ impl RelaySelector {
             SpecializedSelectorConfig::Custom(custom_config) => {
                 Ok(GetRelay::Custom(custom_config.clone()))
             }
-            SpecializedSelectorConfig::Normal(pure_config) => {
+            SpecializedSelectorConfig::Normal(normal_config) => {
                 let parsed_relays = &self.parsed_relays.lock().unwrap();
-                Self::get_relay_inner(query, parsed_relays, &pure_config)
+                Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
             }
         }
     }
@@ -527,30 +527,6 @@ impl RelaySelector {
         self.get_relay_with_custom_params(retry_attempt, &RETRY_ORDER, runtime_params)
     }
 
-    /// Peek at which [`TunnelType`] that would be returned for a certain connection attempt for a
-    /// given [`SelectorConfig`]. Returns [`Option::None`] if the given config would return a
-    /// custom tunnel endpoint.
-    ///
-    /// # Note
-    /// This function is only really useful for testing-purposes. It is exposed to ease testing of
-    /// other mullvad crates which depend on the retry behaviour of [`RelaySelector`].
-    pub fn would_return(connection_attempt: usize, config: &SelectorConfig) -> Option<TunnelType> {
-        match SpecializedSelectorConfig::from(config) {
-            // This case is not really interesting
-            SpecializedSelectorConfig::Custom(_) => None,
-            SpecializedSelectorConfig::Normal(config) => Some(
-                Self::pick_and_merge_query(
-                    connection_attempt,
-                    &RETRY_ORDER,
-                    RuntimeParameters::default(),
-                    RelayQuery::from(config),
-                )
-                .tunnel_protocol
-                .unwrap_or(TunnelType::Wireguard),
-            ),
-        }
-    }
-
     /// Returns a random relay and relay endpoint matching the current constraints defined by
     /// `retry_order` corresponding to `retry_attempt`.
     pub fn get_relay_with_custom_params(
@@ -571,14 +547,14 @@ impl RelaySelector {
             SpecializedSelectorConfig::Normal(normal_config) => {
                 let parsed_relays = &self.parsed_relays.lock().unwrap();
                 // Merge user preferences with the relay selector's default preferences.
-                let user_preferences = RelayQuery::from(normal_config.clone());
                 let query = Self::pick_and_merge_query(
                     retry_attempt,
                     retry_order,
                     runtime_params,
-                    user_preferences,
-                );
-                Self::get_relay_inner(query, parsed_relays, &normal_config)
+                    &normal_config,
+                    parsed_relays,
+                )?;
+                Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
             }
         }
     }
@@ -593,22 +569,29 @@ impl RelaySelector {
     /// Runtime parameters may affect which of the default queries that are considered. For example,
     /// queries which rely on IPv6 will not be considered if working IPv6 is not available at
     /// runtime.
+    ///
+    /// Returns an error iff the intersection between the user's preferences and every default retry
+    /// attempt-query yields queries with no matching relays. I.e., no retry attempt could ever
+    /// resolve to a relay.
     fn pick_and_merge_query(
         retry_attempt: usize,
         retry_order: &[RelayQuery],
         runtime_params: RuntimeParameters,
-        user_preferences: RelayQuery,
-    ) -> RelayQuery {
-        log::trace!("Merging user preferences {user_preferences:?} with default retry strategy");
+        user_config: &NormalSelectorConfig<'_>,
+        parsed_relays: &ParsedRelays,
+    ) -> Result<RelayQuery, Error> {
+        let user_query = RelayQuery::from(user_config.clone());
+        log::trace!("Merging user preferences {user_query:?} with default retry strategy");
         retry_order
             .iter()
             // Remove candidate queries based on runtime parameters before trying to merge user
             // settings
             .filter(|query| runtime_params.compatible(query))
-            .cycle()
-            .filter_map(|query| query.clone().intersection(user_preferences.clone()))
+            .filter_map(|query| query.clone().intersection(user_query.clone()))
+            .filter(|query| Self::get_relay_inner(query, parsed_relays, user_config.custom_lists).is_ok())
+            .cycle() // If the above filters remove all relays, cycle will also return an empty iterator
             .nth(retry_attempt)
-            .unwrap_or(user_preferences)
+            .ok_or(Error::NoRelay)
     }
 
     /// "Execute" the given query, yielding a final set of relays and/or bridges which the VPN
@@ -629,16 +612,16 @@ impl RelaySelector {
     /// * An `Err` if no suitable bridge is found
     #[cfg(not(target_os = "android"))]
     fn get_relay_inner(
-        query: RelayQuery,
+        query: &RelayQuery,
         parsed_relays: &ParsedRelays,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
     ) -> Result<GetRelay, Error> {
         match query.tunnel_protocol {
             Constraint::Only(TunnelType::Wireguard) => {
-                Self::get_wireguard_relay(query, config, parsed_relays)
+                Self::get_wireguard_relay(query, custom_lists, parsed_relays)
             }
             Constraint::Only(TunnelType::OpenVpn) => {
-                Self::get_openvpn_relay(query, config, parsed_relays)
+                Self::get_openvpn_relay(query, custom_lists, parsed_relays)
             }
             Constraint::Any => {
                 // Try Wireguard, then OpenVPN, then fail
@@ -646,7 +629,9 @@ impl RelaySelector {
                     let mut new_query = query.clone();
                     new_query.tunnel_protocol = Constraint::Only(tunnel_type);
                     // If a suitable relay is found, short-circuit and return it
-                    if let Ok(relay) = Self::get_relay_inner(new_query, parsed_relays, config) {
+                    if let Ok(relay) =
+                        Self::get_relay_inner(&new_query, parsed_relays, custom_lists)
+                    {
                         return Ok(relay);
                     }
                 }
@@ -657,16 +642,17 @@ impl RelaySelector {
 
     #[cfg(target_os = "android")]
     fn get_relay_inner(
-        mut query: RelayQuery,
+        query: &RelayQuery,
         parsed_relays: &ParsedRelays,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
     ) -> Result<GetRelay, Error> {
         // FIXME: A bit of defensive programming - calling `get_wiregurad_relay` with a query that
         // doesn't specify Wireguard as the desired tunnel type is not valid and will lead
         // to unwanted behavior. This should be seen as a workaround, and it would be nicer
         // to lift this invariant to be checked by the type system instead.
+        let mut query = query.clone();
         query.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
-        Self::get_wireguard_relay(query, config, parsed_relays)
+        Self::get_wireguard_relay(&query, custom_lists, parsed_relays)
     }
 
     /// Derive a valid Wireguard relay configuration from `query`.
@@ -683,8 +669,8 @@ impl RelaySelector {
     ///
     /// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
     fn get_wireguard_relay(
-        query: RelayQuery,
-        config: &NormalSelectorConfig<'_>,
+        query: &RelayQuery,
+        custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<GetRelay, Error> {
         assert_eq!(
@@ -692,13 +678,13 @@ impl RelaySelector {
             Constraint::Only(TunnelType::Wireguard)
         );
         let inner = if !query.wireguard_constraints.multihop() {
-            Self::get_wireguard_singlehop_config(&query, config, parsed_relays)?
+            Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays)?
         } else {
-            Self::get_wireguard_multihop_config(&query, config, parsed_relays)?
+            Self::get_wireguard_multihop_config(query, custom_lists, parsed_relays)?
         };
-        let endpoint = Self::get_wireguard_endpoint(&query, parsed_relays, &inner)?;
+        let endpoint = Self::get_wireguard_endpoint(query, parsed_relays, &inner)?;
         let obfuscator =
-            Self::get_wireguard_obfuscator(&query, inner.clone(), &endpoint, parsed_relays)?;
+            Self::get_wireguard_obfuscator(query, inner.clone(), &endpoint, parsed_relays)?;
 
         Ok(GetRelay::Wireguard {
             endpoint,
@@ -714,11 +700,10 @@ impl RelaySelector {
     /// * `Ok(WireguardInner::Singlehop)` otherwise
     fn get_wireguard_singlehop_config(
         query: &RelayQuery,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<WireguardConfig, Error> {
-        let candidates =
-            filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists);
+        let candidates = filter_matching_relay_list(query, parsed_relays.relays(), custom_lists);
         helpers::pick_random_relay(&candidates)
             .cloned()
             .map(WireguardConfig::singlehop)
@@ -734,7 +719,7 @@ impl RelaySelector {
     /// * `Ok(WireguardInner::Multihop)` otherwise
     fn get_wireguard_multihop_config(
         query: &RelayQuery,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<WireguardConfig, Error> {
         // Here, we modify the original query just a bit.
@@ -749,16 +734,10 @@ impl RelaySelector {
         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(),
-            config.custom_lists,
-        );
+        let exit_candidates =
+            filter_matching_relay_list(&exit_relay_query, parsed_relays.relays(), custom_lists);
+        let entry_candidates =
+            filter_matching_relay_list(&entry_relay_query, parsed_relays.relays(), custom_lists);
 
         fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
             list.iter()
@@ -846,16 +825,21 @@ impl RelaySelector {
     /// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
     #[cfg(not(target_os = "android"))]
     fn get_openvpn_relay(
-        query: RelayQuery,
-        config: &NormalSelectorConfig<'_>,
+        query: &RelayQuery,
+        custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<GetRelay, Error> {
         assert_eq!(query.tunnel_protocol, Constraint::Only(TunnelType::OpenVpn));
         let exit =
-            Self::choose_openvpn_relay(&query, config, parsed_relays).ok_or(Error::NoRelay)?;
-        let endpoint = Self::get_openvpn_endpoint(&query, &exit, parsed_relays)?;
-        let bridge =
-            Self::get_openvpn_bridge(&query, &exit, &endpoint.protocol, parsed_relays, config)?;
+            Self::choose_openvpn_relay(query, custom_lists, parsed_relays).ok_or(Error::NoRelay)?;
+        let endpoint = Self::get_openvpn_endpoint(query, &exit, parsed_relays)?;
+        let bridge = Self::get_openvpn_bridge(
+            query,
+            &exit,
+            &endpoint.protocol,
+            parsed_relays,
+            custom_lists,
+        )?;
 
         // FIXME: This assert would be better to encode at the type level.
         assert!(matches!(exit.endpoint_data, RelayEndpointData::Openvpn));
@@ -908,13 +892,13 @@ impl RelaySelector {
         relay: &Relay,
         protocol: &TransportProtocol,
         parsed_relays: &ParsedRelays,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
     ) -> Result<Option<SelectedBridge>, Error> {
         if !BridgeQuery::should_use_bridge(&query.openvpn_constraints.bridge_settings) {
             Ok(None)
         } else {
             let bridge_query = &query.openvpn_constraints.bridge_settings.clone().unwrap();
-            let custom_lists = &config.custom_lists;
+            let custom_lists = &custom_lists;
             match protocol {
                 TransportProtocol::Udp => {
                     log::error!("Can not use OpenVPN bridges over UDP");
@@ -1032,7 +1016,7 @@ impl RelaySelector {
     fn get_relay_midpoint(
         query: &RelayQuery,
         parsed_relays: &ParsedRelays,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
     ) -> Option<Coordinates> {
         use std::ops::Not;
         if query.location.is_any() {
@@ -1040,7 +1024,7 @@ impl RelaySelector {
         }
 
         let matching_locations: Vec<Location> =
-            filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists)
+            filter_matching_relay_list(query, parsed_relays.relays(), custom_lists)
                 .into_iter()
                 .filter_map(|relay| relay.location)
                 .unique_by(|location| location.city.clone())
@@ -1058,12 +1042,12 @@ impl RelaySelector {
     #[cfg(not(target_os = "android"))]
     fn choose_openvpn_relay(
         query: &RelayQuery,
-        config: &NormalSelectorConfig<'_>,
+        custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Option<Relay> {
         // Filter among all valid relays
         let relays = parsed_relays.relays();
-        let candidates = filter_matching_relay_list(query, relays, config.custom_lists);
+        let candidates = filter_matching_relay_list(query, relays, custom_lists);
         // Pick one of the valid relays.
         helpers::pick_random_relay(&candidates).cloned()
     }
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index d92e7a72a9..7cf2329e9a 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -1296,3 +1296,30 @@ fn test_daita() {
         ),
     }
 }
+
+/// Check that if  the original user query would yield a relay, the result of running the query
+/// which is the intersection between the user query and any of the default queries shall never
+/// fail.
+#[test]
+fn valid_user_setting_should_yield_relay() {
+    // Make a valid user relay constraint
+    let location = GeographicLocationConstraint::hostname("se", "got", "se9-wireguard");
+    let user_query = RelayQueryBuilder::new().location(location.clone()).build();
+    let user_constraints = RelayQueryBuilder::new()
+        .location(location.clone())
+        .into_constraint();
+
+    let config = SelectorConfig {
+        relay_settings: user_constraints.into(),
+        ..SelectorConfig::default()
+    };
+    let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
+    let user_result = relay_selector.get_relay_by_query(user_query.clone());
+    for retry_attempt in 0..RETRY_ORDER.len() {
+        let post_unification_result =
+            relay_selector.get_relay(retry_attempt, RuntimeParameters::default());
+        if user_result.is_ok() {
+            assert!(post_unification_result.is_ok(), "Expected Post-unification query to be valid because original query {:#?} yielded a connection configuration", user_query)
+        }
+    }
+}