Filter out 'include_in_country' correctly for bridge relays

2 files changed, 72 insertions, 60 deletions

diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 4435c55d66..33df291903 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -32,7 +32,7 @@ use talpid_types::{
     ErrorExt,
 };
 
-use matcher::{OpenVpnMatcher, RelayMatcher, TunnelMatcher, WireguardMatcher};
+use matcher::{BridgeMatcher, EndpointMatcher, OpenVpnMatcher, RelayMatcher, WireguardMatcher};
 
 mod matcher;
 pub mod updater;
@@ -351,14 +351,16 @@ impl RelaySelector {
             location: location.clone(),
             providers: providers.clone(),
             ownership: *ownership,
-            tunnel: OpenVpnMatcher::new(
+            endpoint_matcher: OpenVpnMatcher::new(
                 openvpn_constraints,
                 self.parsed_relays.lock().locations.openvpn.clone(),
             ),
         };
 
-        if relay_matcher.tunnel.constraints.port.is_any() && bridge_state == BridgeState::On {
-            relay_matcher.tunnel.constraints.port = Constraint::Only(TransportPort {
+        if relay_matcher.endpoint_matcher.constraints.port.is_any()
+            && bridge_state == BridgeState::On
+        {
+            relay_matcher.endpoint_matcher.constraints.port = Constraint::Only(TransportPort {
                 protocol: TransportProtocol::Tcp,
                 port: Constraint::Any,
             });
@@ -370,23 +372,24 @@ impl RelaySelector {
 
         let (preferred_port, preferred_protocol) =
             Self::preferred_openvpn_constraints(retry_attempt);
-        let should_try_preferred = match &mut preferred_relay_matcher.tunnel.constraints.port {
-            any @ Constraint::Any => {
-                *any = Constraint::Only(TransportPort {
-                    protocol: preferred_protocol,
-                    port: preferred_port,
-                });
-                true
-            }
-            Constraint::Only(ref mut port_constraints)
-                if port_constraints.protocol == preferred_protocol
-                    && port_constraints.port.is_any() =>
-            {
-                port_constraints.port = preferred_port;
-                true
-            }
-            _ => false,
-        };
+        let should_try_preferred =
+            match &mut preferred_relay_matcher.endpoint_matcher.constraints.port {
+                any @ Constraint::Any => {
+                    *any = Constraint::Only(TransportPort {
+                        protocol: preferred_protocol,
+                        port: preferred_port,
+                    });
+                    true
+                }
+                Constraint::Only(ref mut port_constraints)
+                    if port_constraints.protocol == preferred_protocol
+                        && port_constraints.port.is_any() =>
+                {
+                    port_constraints.port = preferred_port;
+                    true
+                }
+                _ => false,
+            };
 
         if should_try_preferred {
             self.get_tunnel_endpoint_internal(&preferred_relay_matcher)
@@ -403,7 +406,7 @@ impl RelaySelector {
     ) -> Result<NormalSelectedRelay, Error> {
         let mut exit_matcher = RelayMatcher {
             location: exit_location,
-            tunnel: self.wireguard_exit_matcher(),
+            endpoint_matcher: self.wireguard_exit_matcher(),
             ..entry_matcher.clone()
         };
 
@@ -462,15 +465,15 @@ impl RelaySelector {
             location: location.clone(),
             providers: providers.clone(),
             ownership: *ownership,
-            tunnel: WireguardMatcher::new(
+            endpoint_matcher: WireguardMatcher::new(
                 wireguard_constraints.clone(),
                 self.parsed_relays.lock().locations.wireguard.clone(),
             ),
         };
 
         let mut preferred_matcher: RelayMatcher<WireguardMatcher> = entry_relay_matcher.clone();
-        preferred_matcher.tunnel.port = preferred_matcher
-            .tunnel
+        preferred_matcher.endpoint_matcher.port = preferred_matcher
+            .endpoint_matcher
             .port
             .or(Self::preferred_wireguard_port(retry_attempt));
 
@@ -481,8 +484,8 @@ impl RelaySelector {
         }
 
         entry_relay_matcher.location = wireguard_constraints.entry_location.clone();
-        entry_relay_matcher.tunnel.port = entry_relay_matcher
-            .tunnel
+        entry_relay_matcher.endpoint_matcher.port = entry_relay_matcher
+            .endpoint_matcher
             .port
             .or(Self::preferred_wireguard_port(retry_attempt));
         self.get_wireguard_multi_hop_endpoint(entry_relay_matcher, location.clone())
@@ -516,14 +519,14 @@ impl RelaySelector {
 
         // Pick the entry relay first if its location constraint is a subset of the exit location.
         if relay_constraints.wireguard_constraints.use_multihop {
-            matcher.tunnel.wireguard = self.wireguard_exit_matcher();
+            matcher.endpoint_matcher.wireguard = self.wireguard_exit_matcher();
             if relay_constraints
                 .wireguard_constraints
                 .entry_location
                 .is_subset(&matcher.location)
             {
                 if let Ok((entry_relay, entry_endpoint)) = self.get_entry_endpoint(&entry_matcher) {
-                    matcher.tunnel.wireguard.peer = Some(entry_relay.clone());
+                    matcher.endpoint_matcher.wireguard.peer = Some(entry_relay.clone());
                     selected_entry_relay = Some(entry_relay);
                     selected_entry_endpoint = Some(entry_endpoint);
                 }
@@ -542,7 +545,7 @@ impl RelaySelector {
                 .entry_location
                 .is_subset(&matcher.location)
             {
-                entry_matcher.tunnel.peer = Some(selected_relay.exit_relay.clone());
+                entry_matcher.endpoint_matcher.peer = Some(selected_relay.exit_relay.clone());
                 if let Ok((entry_relay, entry_endpoint)) = self.get_entry_endpoint(&entry_matcher) {
                     selected_entry_relay = Some(entry_relay);
                     selected_entry_endpoint = Some(entry_endpoint);
@@ -804,14 +807,14 @@ impl RelaySelector {
         constraints: &InternalBridgeConstraints,
         location: Option<T>,
     ) -> Option<(ProxySettings, Relay)> {
-        let matching_relays: Vec<Relay> = self
-            .parsed_relays
-            .lock()
-            .relays()
-            .iter()
-            .filter(|relay| relay.active && Self::matching_bridge_relay(relay, constraints))
-            .cloned()
-            .collect();
+        let matcher = RelayMatcher {
+            location: constraints.location.clone(),
+            providers: constraints.providers.clone(),
+            ownership: constraints.ownership,
+            endpoint_matcher: BridgeMatcher(()),
+        };
+        let matching_relays: Vec<Relay> =
+            matcher.filter_matching_relay_list(self.parsed_relays.lock().relays());
 
         if matching_relays.is_empty() {
             return None;
@@ -1059,7 +1062,7 @@ impl RelaySelector {
     }
 
     /// Returns a random relay endpoint if any is matching the given constraints.
-    fn get_tunnel_endpoint_internal<T: TunnelMatcher>(
+    fn get_tunnel_endpoint_internal<T: EndpointMatcher>(
         &self,
         matcher: &RelayMatcher<T>,
     ) -> Result<NormalSelectedRelay, Error> {
@@ -1081,13 +1084,6 @@ impl RelaySelector {
             .ok_or(Error::NoRelay)
     }
 
-    fn matching_bridge_relay(relay: &Relay, constraints: &InternalBridgeConstraints) -> bool {
-        constraints.location.matches(relay)
-            && constraints.providers.matches(relay)
-            && constraints.ownership.matches(relay)
-            && relay.endpoint_data == RelayEndpointData::Bridge
-    }
-
     /// Picks a relay using [Self::pick_random_relay_fn], using the `weight` member of each relay
     /// as the weight function.
     fn pick_random_relay<'a>(&self, relays: &'a [Relay]) -> Option<&'a Relay> {
diff --git a/mullvad-relay-selector/src/matcher.rs b/mullvad-relay-selector/src/matcher.rs
index fb75ea3882..e414a111b9 100644
--- a/mullvad-relay-selector/src/matcher.rs
+++ b/mullvad-relay-selector/src/matcher.rs
@@ -16,11 +16,11 @@ use std::net::{IpAddr, SocketAddr};
 use talpid_types::net::{all_of_the_internet, wireguard, Endpoint, IpVersion, TunnelType};
 
 #[derive(Clone)]
-pub struct RelayMatcher<T: TunnelMatcher> {
+pub struct RelayMatcher<T: EndpointMatcher> {
     pub location: Constraint<LocationConstraint>,
     pub providers: Constraint<Providers>,
     pub ownership: Constraint<Ownership>,
-    pub tunnel: T,
+    pub endpoint_matcher: T,
 }
 
 impl RelayMatcher<AnyTunnelMatcher> {
@@ -33,7 +33,7 @@ impl RelayMatcher<AnyTunnelMatcher> {
             location: constraints.location,
             providers: constraints.providers,
             ownership: constraints.ownership,
-            tunnel: AnyTunnelMatcher {
+            endpoint_matcher: AnyTunnelMatcher {
                 wireguard: WireguardMatcher::new(constraints.wireguard_constraints, wireguard_data),
                 openvpn: OpenVpnMatcher::new(constraints.openvpn_constraints, openvpn_data),
                 tunnel_type: constraints.tunnel_protocol,
@@ -43,7 +43,7 @@ impl RelayMatcher<AnyTunnelMatcher> {
 
     pub fn into_wireguard_matcher(self) -> RelayMatcher<WireguardMatcher> {
         RelayMatcher {
-            tunnel: self.tunnel.wireguard,
+            endpoint_matcher: self.endpoint_matcher.wireguard,
             location: self.location,
             providers: self.providers,
             ownership: self.ownership,
@@ -53,11 +53,11 @@ impl RelayMatcher<AnyTunnelMatcher> {
 
 impl RelayMatcher<WireguardMatcher> {
     pub fn set_peer(&mut self, peer: Relay) {
-        self.tunnel.peer = Some(peer);
+        self.endpoint_matcher.peer = Some(peer);
     }
 }
 
-impl<T: TunnelMatcher> RelayMatcher<T> {
+impl<T: EndpointMatcher> RelayMatcher<T> {
     /// 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(&self, relays: &[Relay]) -> Vec<Relay> {
@@ -83,7 +83,7 @@ impl<T: TunnelMatcher> RelayMatcher<T> {
             return None;
         }
 
-        self.tunnel.filter_matching_endpoints(relay)
+        self.endpoint_matcher.filter_matching_endpoints(relay)
     }
 
     /// Filter a relay and its endpoints based on constraints, 2nd pass.
@@ -103,14 +103,14 @@ impl<T: TunnelMatcher> RelayMatcher<T> {
     }
 
     pub fn mullvad_endpoint(&self, relay: &Relay) -> Option<MullvadEndpoint> {
-        self.tunnel.mullvad_endpoint(relay)
+        self.endpoint_matcher.mullvad_endpoint(relay)
     }
 }
 
-/// TunnelMatcher allows to abstract over different tunnel-specific constraints,
-/// as to not have false dependencies on OpenVpn specific constraints when
+/// EndpointMatcher allows to abstract over different tunnel-specific or bridge constraints.
+/// This enables one to not have false dependencies on OpenVpn specific constraints when
 /// selecting only WireGuard tunnels.
-pub trait TunnelMatcher: Clone {
+pub trait EndpointMatcher: Clone {
     /// Filter a relay and its endpoints based on constraints.
     /// Only matching endpoints are included in the returned Relay.
     fn filter_matching_endpoints(&self, relay: &Relay) -> Option<Relay>;
@@ -119,7 +119,7 @@ pub trait TunnelMatcher: Clone {
     fn mullvad_endpoint(&self, relay: &Relay) -> Option<MullvadEndpoint>;
 }
 
-impl TunnelMatcher for OpenVpnMatcher {
+impl EndpointMatcher for OpenVpnMatcher {
     fn filter_matching_endpoints(&self, relay: &Relay) -> Option<Relay> {
         if !self.matches(&self.data) || !matches!(relay.endpoint_data, RelayEndpointData::Openvpn) {
             return None;
@@ -192,7 +192,7 @@ pub struct AnyTunnelMatcher {
     pub tunnel_type: Constraint<TunnelType>,
 }
 
-impl TunnelMatcher for AnyTunnelMatcher {
+impl EndpointMatcher for AnyTunnelMatcher {
     fn filter_matching_endpoints(&self, relay: &Relay) -> Option<Relay> {
         match self.tunnel_type {
             Constraint::Any => {
@@ -328,7 +328,7 @@ impl WireguardMatcher {
     }
 }
 
-impl TunnelMatcher for WireguardMatcher {
+impl EndpointMatcher for WireguardMatcher {
     fn filter_matching_endpoints(&self, relay: &Relay) -> Option<Relay> {
         if self
             .peer
@@ -348,3 +348,19 @@ impl TunnelMatcher for WireguardMatcher {
         self.wg_data_to_endpoint(relay, &self.data)
     }
 }
+
+#[derive(Clone)]
+pub struct BridgeMatcher(pub ());
+
+impl EndpointMatcher for BridgeMatcher {
+    fn filter_matching_endpoints(&self, relay: &Relay) -> Option<Relay> {
+        if !matches!(relay.endpoint_data, RelayEndpointData::Bridge) {
+            return None;
+        }
+        Some(relay.clone())
+    }
+
+    fn mullvad_endpoint(&self, _relay: &Relay) -> Option<MullvadEndpoint> {
+        None
+    }
+}