Disregard `include_in_country` flag when needed

In an edge case where there are very few relays to choose from and the
location constraint is a country for both the exit and entry, the
`include_in_country` could filter out too many relays to be able to pick
a working multihop circuit.

Opportunistically filter based on the `include_in_country` flag. If the
relay selector fails to select a valid multihop circuit, try to filter
relays without considering the `include_in_country` flag.

2 files changed, 90 insertions, 37 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index 6b76eecef8..b65f6f623a 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/mullvad-relay-selector/src/relay_selector/matcher.rs
@@ -16,15 +16,28 @@ use super::query::{ObfuscationQuery, RelayQuery, WireguardRelayQuery};
 
 /// Filter a list of relays and their endpoints based on constraints.
 /// Only relays with (and including) matching endpoints are returned.
+///
+/// This function filter relays on the `include_in_country` flag, as opposed to [filter_matching_relay_by_query].
 pub fn filter_matching_relay_list(
     query: &RelayQuery,
     relay_list: &RelayList,
     custom_lists: &CustomListsSettings,
 ) -> Vec<Relay> {
-    let relays = relay_list.relays();
+    let relays = filter_matching_relay_list_include_all(query, relay_list, custom_lists);
+    let locations = ResolvedLocationConstraint::from_constraint(query.location(), custom_lists);
+    filter_on_include_in_country(locations, relays)
+}
 
+/// Filter a list of relays and their endpoints based on constraints.
+/// Only relays with (and including) matching endpoints are returned.
+pub fn filter_matching_relay_list_include_all(
+    query: &RelayQuery,
+    relay_list: &RelayList,
+    custom_lists: &CustomListsSettings,
+) -> Vec<Relay> {
+    let relays = relay_list.relays();
     let locations = ResolvedLocationConstraint::from_constraint(query.location(), custom_lists);
-    let shortlist = relays
+    relays
             // Filter on tunnel type
             .filter(|relay| filter_tunnel_type(&query.tunnel_protocol(), relay))
             // Filter on active relays
@@ -38,32 +51,7 @@ pub fn filter_matching_relay_list(
             // Filter by DAITA support
             .filter(|relay| filter_on_daita(&query.wireguard_constraints().daita, relay))
             // Filter by obfuscation support
-            .filter(|relay| filter_on_obfuscation(query.wireguard_constraints(), relay_list, 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
-    // `include_in_country` set to true should always be prioritized over relays which has this
-    // flag set to false. We should only consider relays with `include_in_country` set to false
-    // if there are no other candidates left.
-    match &locations {
-        Constraint::Any => shortlist.cloned().collect(),
-        Constraint::Only(locations) => {
-            let mut included = HashSet::new();
-            let mut excluded = HashSet::new();
-            for location in locations {
-                let (included_in_country, not_included_in_country): (Vec<_>, Vec<_>) = shortlist
-                    .clone()
-                    .partition(|relay| location.is_country() && relay.include_in_country);
-                included.extend(included_in_country);
-                excluded.extend(not_included_in_country);
-            }
-            if included.is_empty() {
-                excluded.into_iter().cloned().collect()
-            } else {
-                included.into_iter().cloned().collect()
-            }
-        }
-    }
+            .filter(|relay| filter_on_obfuscation(query.wireguard_constraints(), relay_list, relay)).cloned().collect()
 }
 
 pub fn filter_matching_bridges<'a, R: Iterator<Item = &'a Relay> + Clone>(
@@ -191,6 +179,35 @@ fn filter_on_shadowsocks(
     }
 }
 
+/// When the location constraint is based on country, a relay which has
+/// `include_in_country` set to true should always be prioritized over relays which has this
+/// flag set to false. We should only consider relays with `include_in_country` set to false
+/// if there are no other candidates left.
+fn filter_on_include_in_country(
+    locations: Constraint<ResolvedLocationConstraint<'_>>,
+    relays: Vec<Relay>,
+) -> Vec<Relay> {
+    match locations {
+        Constraint::Any => relays,
+        Constraint::Only(locations) => {
+            let mut included = HashSet::new();
+            let mut excluded = HashSet::new();
+            for location in &locations {
+                let (included_in_country, not_included_in_country): (Vec<_>, Vec<_>) = relays
+                    .iter()
+                    .partition(|relay| location.is_country() && relay.include_in_country);
+                included.extend(included_in_country);
+                excluded.extend(not_included_in_country);
+            }
+            if included.is_empty() {
+                excluded.into_iter().cloned().collect()
+            } else {
+                included.into_iter().cloned().collect()
+            }
+        }
+    }
+}
+
 /// 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 8c5f729d56..118cda99d4 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -8,7 +8,9 @@ pub mod query;
 pub mod relays;
 
 use detailer::resolve_ip_version;
-use matcher::{filter_matching_bridges, filter_matching_relay_list};
+use matcher::{
+    filter_matching_bridges, filter_matching_relay_list, filter_matching_relay_list_include_all,
+};
 use parsed_relays::ParsedRelays;
 use relays::{Multihop, Singlehop, WireguardConfig};
 
@@ -781,9 +783,10 @@ impl RelaySelector {
     ) -> Result<Multihop, Error> {
         let mut exit_relay_query = query.clone();
 
-        // DAITA should only be enabled for the entry relay
+        // DAITA & obfuscation should only be enabled for the entry relay
         let mut wireguard_constraints = exit_relay_query.wireguard_constraints().clone();
         wireguard_constraints.daita = Constraint::Only(false);
+        wireguard_constraints.obfuscation = ObfuscationQuery::Off;
         exit_relay_query.set_wireguard_constraints(wireguard_constraints)?;
 
         let exit_candidates =
@@ -840,20 +843,56 @@ impl RelaySelector {
         // we can query for all exit & entry candidates! All candidates are needed for the next
         // step.
         let mut exit_relay_query = query.clone();
-        // DAITA should only be enabled for the entry relay
 
+        // DAITA & Obfuscation should only be enabled for the entry relay
         let mut wg_constraints = exit_relay_query.wireguard_constraints().clone();
         wg_constraints.daita = Constraint::Only(false);
+        wg_constraints.obfuscation = ObfuscationQuery::Off;
         exit_relay_query.set_wireguard_constraints(wg_constraints)?;
 
+        // Opportunistically filter on `include_in_country`.
         let exit_candidates =
             filter_matching_relay_list(&exit_relay_query, parsed_relays, custom_lists);
         let entry_candidates =
             filter_matching_relay_list(&entry_relay_query, parsed_relays, custom_lists);
 
-        // We avoid picking the same relay for entry and exit by choosing one and excluding it when
-        // choosing the other.
-        let (exit, entry) = match (exit_candidates.as_slice(), entry_candidates.as_slice()) {
+        match Self::pick_working_entry_exit_combo(
+            exit_candidates.as_slice(),
+            entry_candidates.as_slice(),
+        ) {
+            Some((exit, entry)) => Ok(Multihop::new(entry.clone(), exit.clone())),
+            None => {
+                // Sometimes, the set of relays is too small to consider the `include_in_country`
+                // flag. It might just be that if we disregard the `include_in_country` flag, we
+                // manage to find candidate relays. This is rather unlikely, but it might just
+                // happen.
+                let exit_candidates = filter_matching_relay_list_include_all(
+                    &exit_relay_query,
+                    parsed_relays,
+                    custom_lists,
+                );
+                let entry_candidates = filter_matching_relay_list_include_all(
+                    &entry_relay_query,
+                    parsed_relays,
+                    custom_lists,
+                );
+                let (exit, entry) = Self::pick_working_entry_exit_combo(
+                    exit_candidates.as_slice(),
+                    entry_candidates.as_slice(),
+                )
+                .ok_or(Error::NoRelay)?;
+                Ok(Multihop::new(entry.clone(), exit.clone()))
+            }
+        }
+    }
+
+    /// Avoid picking the same relay for entry and exit by choosing one and excluding it when
+    /// choosing the other.
+    fn pick_working_entry_exit_combo<'a>(
+        exit_candidates: &'a [Relay],
+        entry_candidates: &'a [Relay],
+    ) -> Option<(&'a Relay, &'a Relay)> {
+        match (exit_candidates, entry_candidates) {
             // In the case where there is only one entry to choose from, we have to pick it before
             // the exit
             (exits, [entry]) if exits.contains(entry) => {
@@ -867,9 +906,6 @@ impl RelaySelector {
                 helpers::pick_random_relay_excluding(entries, exit).map(|entry| (exit, entry))
             }),
         }
-        .ok_or(Error::NoRelay)?;
-
-        Ok(Multihop::new(entry.clone(), exit.clone()))
     }
 
     /// Constructs a [`MullvadEndpoint`] with details for how to connect to `relay`.