Small refactor of entry exit picking algorithm

Author: Serock3

gui/build-e2e.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+set -ex
+
+COMMIT=$(git rev-parse --short=6 HEAD)
+npm run build-test-executable
+cp ../dist/mullvadvpn-app-e2e-tests "$HOME/.cache/mullvad-test/packages/app-e2e-tests-2024.1-beta2-dev-${COMMIT}_amd64-unknown-linux-gnu"
+cp ../dist/mullvadvpn-app-e2e-tests "$HOME/.cache/mullvad-test/packages/app-e2e-tests-2024.1-beta2-dev-${COMMIT}_x86_64-unknown-linux-gnu"

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> {
@@ -53,11 +40,11 @@ pub fn pick_random_relay_weighted<RelayType>(
         //                                |
         //                                v
         //   _____________________________i___________________________________________________
-        // 0|_____________|__________________________|___________|_____|___________|__________| total_weight
-        //  ^             ^                          ^                             ^          ^
-        //  |             |                          |                             |          |
-        //  ------------------------------------------                             ------------
-        //         |                    |                                                |
+        // 0|_____________|__________________________|___________|_____|___________|__________|
+        // total_weight  ^             ^                          ^
+        // ^          ^  |             |                          |                             |
+        // |  ------------------------------------------
+        // ------------         |                    |                                                |
         //   weight(relay 0)       weight(relay 1)        ..       ..       ..     weight(relay n)
         let mut i: u64 = rng.gen_range(1..=total_weight);
         Some(

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},
@@ -694,31 +695,32 @@ 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 exit to choose from, we have to pick it before
+            // the entry
             (exits, [entry]) if exits.contains(entry) => {
-                let exit = helpers::random(exits, entry).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
+                pick_random_excluding(exits, entry).map(|exit| (exit, entry))
             }
+            // Vice versa for the case of only one entry
             ([exit], entries) if entries.contains(exit) => {
-                let entry = helpers::random(entries, exit).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
-            }
-            (exits, entries) => {
-                let exit = helpers::pick_random_relay(exits).ok_or(Error::NoRelay)?;
-                let entry = helpers::random(entries, exit).ok_or(Error::NoRelay)?;
-                Some((exit, entry))
+                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)?;
 
         Ok(WireguardConfig::Multihop {
-            exit: exit.clone(),
-            entry: entry.clone(),
+            exit: exit.to_owned(),
+            entry: entry.to_owned(),
         })
     }