Filter out relays when Shadowsocks is enabled and specific ports are set

Co-authored-by: Sebastian Holmin <sebastian.holmin@mullvad.net>

2 files changed, 62 insertions, 11 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index f6c244d3f2..1bd96a0cc4 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/mullvad-relay-selector/src/relay_selector/matcher.rs
@@ -6,21 +6,26 @@ use mullvad_types::{
     custom_list::CustomListsSettings,
     relay_constraints::{
         GeographicLocationConstraint, InternalBridgeConstraints, LocationConstraint, Ownership,
-        Providers,
+        Providers, SelectedObfuscation, ShadowsocksSettings,
     },
     relay_list::{Relay, RelayEndpointData, WireguardRelayEndpointData},
 };
 use talpid_types::net::TunnelType;
 
-use super::query::RelayQuery;
+use super::{
+    parsed_relays::ParsedRelays,
+    query::{RelayQuery, WireguardRelayQuery},
+};
 
 /// 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<'a, R: Iterator<Item = &'a Relay> + Clone>(
+pub fn filter_matching_relay_list(
     query: &RelayQuery,
-    relays: R,
+    relay_list: &ParsedRelays,
     custom_lists: &CustomListsSettings,
 ) -> Vec<Relay> {
+    let relays = relay_list.relays();
+
     let locations = ResolvedLocationConstraint::from_constraint(&query.location, custom_lists);
     let shortlist = relays
             // Filter on tunnel type
@@ -34,7 +39,9 @@ pub fn filter_matching_relay_list<'a, R: Iterator<Item = &'a Relay> + Clone>(
             // Filter by providers
             .filter(|relay| filter_on_providers(&query.providers, relay))
             // Filter by DAITA support
-            .filter(|relay| filter_on_daita(&query.wireguard_constraints.daita, relay));
+            .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
@@ -123,6 +130,51 @@ pub fn filter_on_daita(filter: &Constraint<bool>, relay: &Relay) -> bool {
     }
 }
 
+/// Returns whether `relay` satisfies the obfuscation settings.
+fn filter_on_obfuscation(
+    query: &WireguardRelayQuery,
+    relay_list: &ParsedRelays,
+    relay: &Relay,
+) -> bool {
+    match query.obfuscation {
+        // Shadowsocks has relay-specific constraints
+        SelectedObfuscation::Shadowsocks => {
+            let wg_data = &relay_list.parsed_list().wireguard;
+            filter_on_shadowsocks(
+                &wg_data.shadowsocks_port_ranges,
+                &query.shadowsocks_port,
+                relay,
+            )
+        }
+
+        // If Shadowsocks is not a requirement, then there are no relay-specific constraints
+        _ => true,
+    }
+}
+
+/// Returns whether `relay` satisfies the Shadowsocks filter posed by `port`.
+fn filter_on_shadowsocks(
+    port_ranges: &[(u16, u16)],
+    settings: &Constraint<ShadowsocksSettings>,
+    relay: &Relay,
+) -> bool {
+    match (&settings, &relay.endpoint_data) {
+        // If Shadowsocks is specifically asked for, we must check if the specific relay supports our port.
+        // If there are extra addresses, then all ports are available, so we do not need to do this.
+        (
+            Constraint::Only(ShadowsocksSettings {
+                port: Constraint::Only(desired_port),
+            }),
+            RelayEndpointData::Wireguard(wg_data),
+        ) if wg_data.shadowsocks_extra_addr_in.is_empty() => port_ranges
+            .iter()
+            .any(|(begin, end)| (*begin..=*end).contains(&desired_port)),
+
+        // Otherwise, 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 5c5b324081..62f8d31ab4 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -704,7 +704,7 @@ impl RelaySelector {
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<WireguardConfig, Error> {
-        let candidates = filter_matching_relay_list(query, parsed_relays.relays(), custom_lists);
+        let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
         helpers::pick_random_relay(&candidates)
             .cloned()
             .map(WireguardConfig::singlehop)
@@ -736,9 +736,9 @@ impl RelaySelector {
         // 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(), custom_lists);
+            filter_matching_relay_list(&exit_relay_query, parsed_relays, custom_lists);
         let entry_candidates =
-            filter_matching_relay_list(&entry_relay_query, parsed_relays.relays(), custom_lists);
+            filter_matching_relay_list(&entry_relay_query, parsed_relays, custom_lists);
 
         fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
             list.iter()
@@ -1035,7 +1035,7 @@ impl RelaySelector {
         }
 
         let matching_locations: Vec<Location> =
-            filter_matching_relay_list(query, parsed_relays.relays(), custom_lists)
+            filter_matching_relay_list(query, parsed_relays, custom_lists)
                 .into_iter()
                 .filter_map(|relay| relay.location)
                 .unique_by(|location| location.city.clone())
@@ -1057,8 +1057,7 @@ impl RelaySelector {
         parsed_relays: &ParsedRelays,
     ) -> Option<Relay> {
         // Filter among all valid relays
-        let relays = parsed_relays.relays();
-        let candidates = filter_matching_relay_list(query, relays, custom_lists);
+        let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
         // Pick one of the valid relays.
         helpers::pick_random_relay(&candidates).cloned()
     }