Merge branch 'small-smart-routing-code-cleanup'

2 files changed, 18 insertions, 12 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index c49a7df04d..c8b529e0cd 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -29,6 +29,17 @@ pub enum Error {
     NoMatchingPort,
 }
 
+/// Picks a relay at random from `relays`, but don't pick `exclude`.
+pub fn pick_random_relay_excluding<'a>(
+    relays: &'a [Relay],
+    exclude: &'_ Relay,
+) -> Option<&'a Relay> {
+    relays
+        .iter()
+        .filter(|&a| a != exclude)
+        .choose(&mut thread_rng())
+}
+
 /// Picks a relay using [pick_random_relay_weighted], using the `weight` member of each relay
 /// as the weight function.
 pub fn pick_random_relay(relays: &[Relay]) -> Option<&Relay> {
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index fb3c37064c..6233e2da05 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -9,7 +9,6 @@ pub mod query;
 use chrono::{DateTime, Local};
 use itertools::Itertools;
 use query::ObfuscationQuery;
-use rand::{seq::IteratorRandom, thread_rng};
 use std::{
     path::Path,
     sync::{Arc, LazyLock, Mutex},
@@ -803,7 +802,8 @@ impl RelaySelector {
             .take_while(|relay| relay.distance <= smallest_distance)
             .map(|relay_with_distance| relay_with_distance.relay)
             .collect_vec();
-        let entry = pick_random_excluding(&entry_candidates, exit).ok_or(Error::NoRelay)?;
+        let entry =
+            helpers::pick_random_relay_excluding(&entry_candidates, exit).ok_or(Error::NoRelay)?;
 
         Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
     }
@@ -847,14 +847,15 @@ impl RelaySelector {
             // 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) => {
-                pick_random_excluding(exits, entry).map(|exit| (exit, entry))
+                helpers::pick_random_relay_excluding(exits, entry).map(|exit| (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))
+                helpers::pick_random_relay_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))),
+            (exits, entries) => helpers::pick_random_relay(exits).and_then(|exit| {
+                helpers::pick_random_relay_excluding(entries, exit).map(|entry| (exit, entry))
+            }),
         }
         .ok_or(Error::NoRelay)?;
 
@@ -1154,12 +1155,6 @@ impl RelaySelector {
     }
 }
 
-fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
-    list.iter()
-        .filter(|&a| a != exclude)
-        .choose(&mut thread_rng())
-}
-
 #[derive(Clone)]
 struct RelayWithDistance {
     distance: f64,