Merge branch 'select-random-bridge'

2 files changed, 34 insertions, 7 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6779cb27eb..850e8ad35a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@ Line wrap the file at 100 chars.                                              Th
 - Change behavior of escape key in the desktop app. It now navigates backwards one step instead of
   to the main view. To navigate back to the main view Shift+Escape can be used.
 - Update Electron from 16.0.4 to 17.1.0.
+- Randomize bridge selection with a bias in favor of close bridges.
 
 ### Fixed
 - Fix the sometimes incorrect time added text after adding time to the account.
diff --git a/mullvad-daemon/src/relays/mod.rs b/mullvad-daemon/src/relays/mod.rs
index 8ec15acf4d..ceaea8e135 100644
--- a/mullvad-daemon/src/relays/mod.rs
+++ b/mullvad-daemon/src/relays/mod.rs
@@ -53,6 +53,11 @@ const WIREGUARD_EXIT_CONSTRAINTS: WireguardMatcher = WireguardMatcher {
 
 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;
+
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
 pub enum Error {
@@ -693,10 +698,15 @@ impl RelaySelector {
 
         let relay = if let Some(location) = location {
             let location = location.into();
-            matching_relays.sort_by_cached_key(|relay| {
-                (relay.location.as_ref().unwrap().distance_from(&location) * 1000.0) as i64
+            matching_relays.sort_by_cached_key(|relay: &Relay| {
+                relay.location.as_ref().unwrap().distance_from(&location) as usize
             });
-            matching_relays.get(0)
+            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)
+            };
+            self.pick_random_relay_fn(&matching_relays, weight_fn)
         } else {
             self.pick_random_relay(&matching_relays)
         };
@@ -927,12 +937,26 @@ impl RelaySelector {
         Some(filtered_relay)
     }
 
+    /// 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)
+    }
+
     /// 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<'a>(&self, relays: &'a [Relay]) -> Option<&'a Relay> {
-        let total_weight: u64 = relays.iter().map(|relay| relay.weight).sum();
+    fn pick_random_relay_fn<'a>(
+        &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();
         let mut rng = rand::thread_rng();
         if total_weight == 0 {
             relays.choose(&mut rng)
@@ -943,10 +967,12 @@ impl RelaySelector {
             Some(
                 relays
                     .iter()
-                    .find(|relay| {
-                        i = i.saturating_sub(relay.weight);
+                    .enumerate()
+                    .find(|(index, relay)| {
+                        i = i.saturating_sub(weight_fn(*index, relay));
                         i == 0
                     })
+                    .map(|(_, relay)| relay)
                     .expect("At least one relay must've had a weight above 0"),
             )
         }