Skip to content

Commit

Permalink
Move pick_random_relay_excluding to helpers module for consistency
Browse files Browse the repository at this point in the history
All other helper methods for selecting random relays are there
  • Loading branch information
faern committed Sep 26, 2024
1 parent c61a054 commit cf0278b
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 12 deletions.
11 changes: 11 additions & 0 deletions mullvad-relay-selector/src/relay_selector/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
19 changes: 7 additions & 12 deletions mullvad-relay-selector/src/relay_selector/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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()))
}
Expand Down Expand Up @@ -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)?;

Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit cf0278b

Please sign in to comment.