Merge branch 'setting-the-entry-and-exit-location-to-sweden-causes-relay-des-2382'

4 files changed, 185 insertions, 38 deletions

diff --git a/.github/workflows/git-commit-message-style.yml b/.github/workflows/git-commit-message-style.yml
index b32229754a..6f93ea9d40 100644
--- a/.github/workflows/git-commit-message-style.yml
+++ b/.github/workflows/git-commit-message-style.yml
@@ -34,4 +34,4 @@ jobs:
           # This action defaults to 50 char subjects, but 72 is fine.
           max-subject-line-length: '72'
           # The action's wordlist is a bit short. Add more accepted verbs
-          additional-verbs: 'tidy, wrap, obfuscate, bias, prohibit, forbid, revert, slim, impl'
+          additional-verbs: 'tidy, wrap, obfuscate, bias, prohibit, forbid, revert, slim, impl, disregard, reproduce'
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`.
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 87a089af62..bab857cf01 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -1728,3 +1728,97 @@ fn test_runtime_ipv4_unavailable() {
         ),
     }
 }
+
+/// Check that the relay selector is able to disregard `include_in_country` flag if necessary.
+///
+/// This test case prevents regressions to the `include_in_country` filtering logic.
+#[test]
+fn include_in_country_with_few_relays() -> Result<(), Error> {
+    let query = RelayQueryBuilder::wireguard()
+        .multihop()
+        .location(GeographicLocationConstraint::country("se"))
+        .entry(GeographicLocationConstraint::country("se"))
+        .build();
+
+    // The relay selector ought to resolve the query to any of the following configurations
+    // {entry: se-sto-wg-009, exit: se-sto-wg-204}
+    // {entry: se-sto-wg-204, exit: se-sto-wg-009}
+    let relays = {
+        let stockholm = Location {
+            country: "Sweden".to_string(),
+            country_code: "se".to_string(),
+            city: "Stockholm".to_string(),
+            city_code: "sto".to_string(),
+            latitude: 59.3289,
+            longitude: 18.0649,
+        };
+        let wireguard = WireguardEndpointData {
+            port_ranges: vec![443..=443],
+            ..Default::default()
+        };
+        RelayList {
+            countries: vec![RelayListCountry {
+                name: "Sweden".to_string(),
+                code: "se".to_string(),
+                cities: vec![RelayListCity {
+                    name: "Stockholm".to_string(),
+                    code: "sto".to_string(),
+                    latitude: 59.3289,
+                    longitude: 18.0649,
+                    relays: vec![
+                        Relay {
+                            hostname: "se-sto-wg-009".to_string(),
+                            ipv4_addr_in: "185.195.233.69".parse().unwrap(),
+                            ipv6_addr_in: "2a03:1b20:4:f011::a09f".parse().ok(),
+                            overridden_ipv4: false,
+                            overridden_ipv6: false,
+                            // This is the important part
+                            include_in_country: false,
+                            active: true,
+                            owned: true,
+                            provider: "31173".to_string(),
+                            weight: 1,
+                            location: stockholm.clone(),
+                            endpoint_data: RelayEndpointData::Wireguard(
+                                WireguardRelayEndpointData::new(
+                                    PublicKey::from_base64(
+                                        "t1XlQD7rER0JUPrmh3R5IpxjUP9YOqodJAwfRorNxl4=",
+                                    )
+                                    .unwrap(),
+                                ),
+                            ),
+                        },
+                        Relay {
+                            hostname: "se-sto-wg-204".to_string(),
+                            ipv4_addr_in: "89.37.63.190".parse().unwrap(),
+                            ipv6_addr_in: "2a02:6ea0:1508:4::f001".parse().ok(),
+                            overridden_ipv4: false,
+                            overridden_ipv6: false,
+                            // This is the important part
+                            include_in_country: true,
+                            active: true,
+                            owned: false,
+                            provider: "DataPacket".to_string(),
+                            weight: 200,
+                            location: stockholm,
+                            endpoint_data: RelayEndpointData::Wireguard(
+                                WireguardRelayEndpointData::new(
+                                    PublicKey::from_base64(
+                                        "cPhM7ShRWQmKiJtD9Wd1vDh0GwIlaMvFb/WPrP58FH8=",
+                                    )
+                                    .unwrap(),
+                                ),
+                            ),
+                        },
+                    ],
+                }],
+            }],
+            wireguard,
+            ..Default::default()
+        }
+    };
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relays);
+
+    relay_selector.get_relay_by_query(query)?;
+    Ok(())
+}