diff options
| -rw-r--r-- | mullvad-relay-selector/src/lib.rs | 41 | ||||
| -rw-r--r-- | mullvad-relay-selector/src/matcher.rs | 37 | ||||
| -rw-r--r-- | mullvad-types/src/relay_constraints.rs | 23 |
3 files changed, 68 insertions, 33 deletions
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs index 94725e767e..4435c55d66 100644 --- a/mullvad-relay-selector/src/lib.rs +++ b/mullvad-relay-selector/src/lib.rs @@ -321,17 +321,12 @@ impl RelaySelector { }; let matcher = RelayMatcher::new(relay_constraints.clone(), openvpn_data, wireguard_data); - let mut matching_locations: Vec<Location> = self - .parsed_relays - .lock() - .relays() - .iter() - .filter(|relay| relay.active) - .filter_map(|relay| { - matcher - .filter_matching_relay(relay) - .and_then(|relay| relay.location) - }) + + let parsed_relays = self.parsed_relays.lock(); + let mut matching_locations: Vec<Location> = matcher + .filter_matching_relay_list(parsed_relays.relays()) + .into_iter() + .filter_map(|relay| relay.location) .collect(); matching_locations.dedup_by(|a, b| a.has_same_city(b)); @@ -697,13 +692,9 @@ impl RelaySelector { &self, matcher: &RelayMatcher<WireguardMatcher>, ) -> Result<(Relay, MullvadWireguardEndpoint), Error> { - let matching_relays: Vec<Relay> = self - .parsed_relays - .lock() - .relays() - .iter() - .filter(|relay| relay.active) - .filter_map(|relay| matcher.filter_matching_relay(relay)) + let matching_relays: Vec<Relay> = matcher + .filter_matching_relay_list(self.parsed_relays.lock().relays()) + .into_iter() .collect(); let relay = self @@ -986,7 +977,7 @@ impl RelaySelector { self.parsed_relays.lock().relays().iter().any(|relay| { relay.active && relay.endpoint_data == RelayEndpointData::Openvpn - && location_constraint.matches(relay) + && location_constraint.matches_with_opts(relay, true) && providers_constraint.matches(relay) && ownership_constraint.matches(relay) }); @@ -1002,7 +993,7 @@ impl RelaySelector { self.parsed_relays.lock().relays().iter().any(|relay| { relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_)) - && location_constraint.matches(relay) + && location_constraint.matches_with_opts(relay, true) && providers_constraint.matches(relay) && ownership_constraint.matches(relay) }); @@ -1072,13 +1063,9 @@ impl RelaySelector { &self, matcher: &RelayMatcher<T>, ) -> Result<NormalSelectedRelay, Error> { - let matching_relays: Vec<Relay> = self - .parsed_relays - .lock() - .relays() - .iter() - .filter(|relay| relay.active) - .filter_map(|relay| matcher.filter_matching_relay(relay)) + let matching_relays: Vec<Relay> = matcher + .filter_matching_relay_list(self.parsed_relays.lock().relays()) + .into_iter() .collect(); self.pick_random_relay(&matching_relays) diff --git a/mullvad-relay-selector/src/matcher.rs b/mullvad-relay-selector/src/matcher.rs index ac3d573832..fb75ea3882 100644 --- a/mullvad-relay-selector/src/matcher.rs +++ b/mullvad-relay-selector/src/matcher.rs @@ -58,12 +58,27 @@ impl RelayMatcher<WireguardMatcher> { } impl<T: TunnelMatcher> RelayMatcher<T> { - /// Filter a relay and its endpoints based on constraints. + /// 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> { + let matches = relays + .iter() + .filter_map(|relay| self.pre_filter_matching_relay(relay)); + + let ignore_include_in_country = !matches.clone().any(|relay| relay.include_in_country); + + matches + .filter_map(|relay| self.post_filter_matching_relay(relay, ignore_include_in_country)) + .collect() + } + + /// Filter a relay and its endpoints based on constraints, 1st pass. /// Only matching endpoints are included in the returned Relay. - pub fn filter_matching_relay(&self, relay: &Relay) -> Option<Relay> { - if !self.location.matches(relay) + fn pre_filter_matching_relay(&self, relay: &Relay) -> Option<Relay> { + if !relay.active || !self.providers.matches(relay) || !self.ownership.matches(relay) + || !self.location.matches_with_opts(relay, true) { return None; } @@ -71,6 +86,22 @@ impl<T: TunnelMatcher> RelayMatcher<T> { self.tunnel.filter_matching_endpoints(relay) } + /// Filter a relay and its endpoints based on constraints, 2nd pass. + /// Only matching endpoints are included in the returned Relay. + fn post_filter_matching_relay( + &self, + relay: Relay, + ignore_include_in_country: bool, + ) -> Option<Relay> { + if !self + .location + .matches_with_opts(&relay, ignore_include_in_country) + { + return None; + } + Some(relay) + } + pub fn mullvad_endpoint(&self, relay: &Relay) -> Option<MullvadEndpoint> { self.tunnel.mullvad_endpoint(relay) } diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs index 00dc0e9110..5119bc72b8 100644 --- a/mullvad-types/src/relay_constraints.rs +++ b/mullvad-types/src/relay_constraints.rs @@ -279,15 +279,15 @@ pub enum LocationConstraint { Hostname(CountryCode, CityCode, Hostname), } -impl Match<Relay> for LocationConstraint { - fn matches(&self, relay: &Relay) -> bool { +impl LocationConstraint { + pub fn matches_with_opts(&self, relay: &Relay, ignore_include_in_country: bool) -> bool { match self { LocationConstraint::Country(ref country) => { relay .location .as_ref() .map_or(false, |loc| loc.country_code == *country) - && relay.include_in_country + && (ignore_include_in_country || relay.include_in_country) } LocationConstraint::City(ref country, ref city) => { relay.location.as_ref().map_or(false, |loc| { @@ -305,6 +305,23 @@ impl Match<Relay> for LocationConstraint { } } +impl Constraint<LocationConstraint> { + pub fn matches_with_opts(&self, relay: &Relay, ignore_include_in_country: bool) -> bool { + match self { + Constraint::Only(constraint) => { + constraint.matches_with_opts(relay, ignore_include_in_country) + } + Constraint::Any => true, + } + } +} + +impl Match<Relay> for LocationConstraint { + fn matches(&self, relay: &Relay) -> bool { + self.matches_with_opts(relay, false) + } +} + impl Set<LocationConstraint> for LocationConstraint { /// Returns whether `self` is equal to or a subset of `other`. fn is_subset(&self, other: &Self) -> bool { |