diff options
| -rw-r--r-- | mullvad-relay-selector/src/lib.rs | 58 |
1 files changed, 46 insertions, 12 deletions
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs index 961357abf2..de4d54cc8b 100644 --- a/mullvad-relay-selector/src/lib.rs +++ b/mullvad-relay-selector/src/lib.rs @@ -48,7 +48,13 @@ const UDP2TCP_PORTS: [u16; 3] = [80, 443, 5001]; /// How much to favor bridges that are closer to the selected relay location. Each /// bridge is assigned a base weight based on its rank order proximity to the location. /// Its final weight equals `(base weight) ^ BRIDGE_PROXIMITY_BIAS`. -const BRIDGE_PROXIMITY_BIAS: u32 = 3; +const BRIDGE_PROXIMITY_BIAS: u32 = 1; + +/// Minimum number of bridges to keep for selection when filtering by distance. +const MIN_BRIDGE_COUNT: usize = 5; + +/// Max distance of bridges to consider for selection (km). +const MAX_BRIDGE_DISTANCE: f64 = 1500f64; #[derive(err_derive::Error, Debug)] #[error(no_from)] @@ -812,7 +818,7 @@ impl RelaySelector { constraints: &InternalBridgeConstraints, location: Option<T>, ) -> Option<(ProxySettings, Relay)> { - let mut matching_relays: Vec<Relay> = self + let matching_relays: Vec<Relay> = self .parsed_relays .lock() .relays() @@ -827,20 +833,48 @@ impl RelaySelector { let relay = if let Some(location) = location { let location = location.into(); - matching_relays.sort_by_cached_key(|relay: &Relay| { - relay.location.as_ref().unwrap().distance_from(&location) as usize - }); + + #[derive(Debug, Clone)] + struct RelayWithDistance { + relay: Relay, + distance: f64, + } + + let mut matching_relays: Vec<RelayWithDistance> = matching_relays + .into_iter() + .map(|relay| { + RelayWithDistance { + distance: relay.location.as_ref().unwrap().distance_from(&location), + relay, + } + }) + .collect(); + matching_relays.sort_unstable_by_key(|relay: &RelayWithDistance| relay.distance as usize); + + matching_relays = matching_relays + .into_iter() + .enumerate() + .filter_map(|(i, relay)| { + if i < MIN_BRIDGE_COUNT || relay.distance <= MAX_BRIDGE_DISTANCE { + return Some(relay); + } + None + }) + .collect(); + let max_weight = matching_relays.len(); - let weight_fn = |index, _relay: &Relay| { + let weight_fn = |index, _relay: &RelayWithDistance| { let w = (max_weight - index) as u64; w.saturating_pow(BRIDGE_PROXIMITY_BIAS) }; self.pick_random_relay_fn(&matching_relays, weight_fn) + .cloned() + .map(|relay_with_distance| relay_with_distance.relay) } else { - self.pick_random_relay(&matching_relays) + self.pick_random_relay(&matching_relays).cloned() }; relay.and_then(|relay| { - self.pick_random_bridge(&self.parsed_relays.lock().locations.bridge, relay) + self.pick_random_bridge(&self.parsed_relays.lock().locations.bridge, &relay) .map(|bridge| (bridge, relay.clone())) }) } @@ -1081,11 +1115,11 @@ impl RelaySelector { /// If all of the relays have a weight of 0, one will be picked at random without bias, /// otherwise roulette wheel selection will be used to pick only relays with non-zero /// weights. - fn pick_random_relay_fn<'a>( + fn pick_random_relay_fn<'a, RelayType>( &self, - relays: &'a [Relay], - weight_fn: impl Fn(usize, &Relay) -> u64, - ) -> Option<&'a Relay> { + relays: &'a [RelayType], + weight_fn: impl Fn(usize, &RelayType) -> u64, + ) -> Option<&'a RelayType> { let total_weight: u64 = relays .iter() .enumerate() |