From 5e1558d4b70456b59c950c5471d6e8073b064797 Mon Sep 17 00:00:00 2001
From: Sebastian Holmin <sebastian.holmin@mullvad.net>
Date: Thu, 28 Mar 2024 14:15:21 +0100
Subject: [PATCH] Refactor bridge entry/exit picking algorithm

---
 .../src/relay_selector/helpers.rs             | 35 ++++++-------------
 .../src/relay_selector/mod.rs                 | 31 ++++++++--------
 2 files changed, 27 insertions(+), 39 deletions(-)

diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index 5ad5bf7ab471..263034d83659 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -8,24 +8,11 @@ use mullvad_types::{
     relay_constraints::Udp2TcpObfuscationSettings,
     relay_list::{BridgeEndpointData, Relay, RelayEndpointData},
 };
-use rand::{
-    seq::{IteratorRandom, SliceRandom},
-    thread_rng, Rng,
-};
+use rand::{seq::SliceRandom, thread_rng, Rng};
 use talpid_types::net::{obfuscation::ObfuscatorConfig, proxy::CustomProxy};
 
 use crate::SelectedObfuscator;
 
-/// Pick a random element out of `from`, excluding the element `exclude` from the selection.
-pub fn random<'a, A: PartialEq>(
-    from: impl IntoIterator<Item = &'a A>,
-    exclude: &A,
-) -> Option<&'a A> {
-    from.into_iter()
-        .filter(|&a| a != exclude)
-        .choose(&mut thread_rng())
-}
-
 /// Picks a relay using [pick_random_relay_fn], using the `weight` member of each relay
 /// as the weight function.
 pub fn pick_random_relay(relays: &[Relay]) -> Option<&Relay> {
@@ -49,16 +36,16 @@ pub fn pick_random_relay_weighted<RelayType>(
         // Pick a random number in the range 1..=total_weight. This choses the relay with a
         // non-zero weight.
         //
-        //                                rng(1..=total_weight)
-        //                                |
-        //                                v
-        //   _____________________________i___________________________________________________
-        // 0|_____________|__________________________|___________|_____|___________|__________| total_weight
-        //  ^             ^                          ^                             ^          ^
-        //  |             |                          |                             |          |
-        //  ------------------------------------------                             ------------
-        //         |                    |                                                |
-        //   weight(relay 0)       weight(relay 1)        ..       ..       ..     weight(relay n)
+        //                           rng(1..=total_weight)
+        //                           |
+        //                           v
+        //   ________________________i_______________________________________________
+        // 0|_____________|____________________|___________|_____|________|__________| total_weight
+        //  ^             ^                    ^                          ^          ^
+        //  |             |                    |                          |          |
+        //  ------------------------------------                          ------------
+        //         |                  |                                         |
+        //   weight(relay 0)     weight(relay 1)    ..       ..     ..    weight(relay n)
         let mut i: u64 = rng.gen_range(1..=total_weight);
         Some(
             relays
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 1cb9dca523e8..e0949beae970 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -9,6 +9,7 @@ pub mod query;
 use chrono::{DateTime, Local};
 use itertools::Itertools;
 use once_cell::sync::Lazy;
+use rand::{seq::IteratorRandom, thread_rng};
 use std::{
     path::Path,
     sync::{Arc, Mutex},
@@ -729,25 +730,25 @@ impl RelaySelector {
             config.custom_lists,
         );
 
-        // This algorithm gracefully handles a particular edge case that arise when a constraint on
-        // the exit relay is more specific than on the entry relay which forces the relay selector
-        // to choose one specific relay. The relay selector could end up selecting that specific
-        // relay as the entry relay, thus leaving no remaining exit relay candidates or vice versa.
+        fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
+            list.iter()
+                .filter(|&a| a != exclude)
+                .choose(&mut thread_rng())
+        }
+        // We avoid picking the same relay for entry and exit by choosing one and excluding it when
+        // choosing the other.
         let (exit, entry) = match (exit_candidates.as_slice(), entry_candidates.as_slice()) {
-            ([exit], [entry]) if exit == entry => None,
+            // In the case where there is only one entry to choose from, we have to pick it before
+            // the exit
             (exits, [entry]) if exits.contains(entry) => {
-                let exit = helpers::random(exits, entry).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
-            }
-            ([exit], entrys) if entrys.contains(exit) => {
-                let entry = helpers::random(entrys, exit).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
+                pick_random_excluding(exits, entry).map(|exit| (exit, entry))
             }
-            (exits, entrys) => {
-                let exit = helpers::pick_random_relay(exits).ok_or(Error::NoRelay)?;
-                let entry = helpers::random(entrys, exit).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
+            // Vice versa for the case of only one exit
+            ([exit], entries) if entries.contains(exit) => {
+                pick_random_excluding(entries, exit).map(|entry| (exit, entry))
             }
+            (exits, entries) => helpers::pick_random_relay(exits)
+                .and_then(|exit| pick_random_excluding(entries, exit).map(|entry| (exit, entry))),
         }
         .ok_or(Error::NoRelay)?;