Merge branch 'adjust-bridge-selection'

2 files changed, 55 insertions, 29 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce4e25db79..ee96fe7c10 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,7 @@ Line wrap the file at 100 chars.                                              Th
 - Reject invalid WireGuard ports in the CLI.
 - Reorganize settings into more logical categories.
 - Upgrade wireguard-go to 20220703234212 (Windows: v0.5.3).
+- Prune bridges far away from the selected relay.
 
 #### Windows
 - Remove dependency on `ipconfig.exe`. Call `DnsFlushResolverCache` to flush the DNS cache.
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 961357abf2..bed781eff7 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -45,10 +45,11 @@ const WIREGUARD_EXIT_IP_VERSION: Constraint<IpVersion> = Constraint::Only(IpVers
 
 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;
+/// 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 +813,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 +828,49 @@ 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
-            });
-            let max_weight = matching_relays.len();
-            let weight_fn = |index, _relay: &Relay| {
-                let w = (max_weight - index) as u64;
-                w.saturating_pow(BRIDGE_PROXIMITY_BIAS)
-            };
+
+            #[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);
+
+            let mut greatest_distance = 0f64;
+            matching_relays = matching_relays
+                .into_iter()
+                .enumerate()
+                .filter_map(|(i, relay)| {
+                    if i < MIN_BRIDGE_COUNT || relay.distance <= MAX_BRIDGE_DISTANCE {
+                        if relay.distance > greatest_distance {
+                            greatest_distance = relay.distance;
+                        }
+                        return Some(relay);
+                    }
+                    None
+                })
+                .collect();
+
+            let weight_fn =
+                |relay: &RelayWithDistance| 1 + (greatest_distance - relay.distance) as u64;
+
             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()))
         })
     }
@@ -1074,23 +1104,19 @@ impl RelaySelector {
     /// 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> {
-        self.pick_random_relay_fn(relays, |_index, relay| relay.weight)
+        self.pick_random_relay_fn(relays, |relay| relay.weight)
     }
 
     /// Pick a random relay from the given slice. Will return `None` if the given slice is empty.
     /// 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> {
-        let total_weight: u64 = relays
-            .iter()
-            .enumerate()
-            .map(|(index, relay)| weight_fn(index, relay))
-            .sum();
+        relays: &'a [RelayType],
+        weight_fn: impl Fn(&RelayType) -> u64,
+    ) -> Option<&'a RelayType> {
+        let total_weight: u64 = relays.iter().map(|relay| weight_fn(relay)).sum();
         let mut rng = rand::thread_rng();
         if total_weight == 0 {
             relays.choose(&mut rng)
@@ -1101,12 +1127,11 @@ impl RelaySelector {
             Some(
                 relays
                     .iter()
-                    .enumerate()
-                    .find(|(index, relay)| {
-                        i = i.saturating_sub(weight_fn(*index, relay));
+                    .find(|relay| {
+                        i = i.saturating_sub(weight_fn(relay));
                         i == 0
                     })
-                    .map(|(_, relay)| relay)
+                    .map(|relay| relay)
                     .expect("At least one relay must've had a weight above 0"),
             )
         }