Skip to content

Commit

Permalink
Merge branch 'relay-selector-bug'
Browse files Browse the repository at this point in the history
  • Loading branch information
MarkusPettersson98 committed Apr 17, 2024
2 parents f1d97bc + ce59eb2 commit 52b70e8
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 77 deletions.
138 changes: 61 additions & 77 deletions mullvad-relay-selector/src/relay_selector/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -472,7 +472,7 @@ impl RelaySelector {
let near_location = match specialized_config {
SpecializedSelectorConfig::Normal(config) => {
let user_preferences = RelayQuery::from(config.clone());
Self::get_relay_midpoint(&user_preferences, parsed_relays, &config)
Self::get_relay_midpoint(&user_preferences, parsed_relays, config.custom_lists)
}
SpecializedSelectorConfig::Custom(_) => None,
};
Expand Down Expand Up @@ -508,9 +508,9 @@ impl RelaySelector {
SpecializedSelectorConfig::Custom(custom_config) => {
Ok(GetRelay::Custom(custom_config.clone()))
}
SpecializedSelectorConfig::Normal(pure_config) => {
SpecializedSelectorConfig::Normal(normal_config) => {
let parsed_relays = &self.parsed_relays.lock().unwrap();
Self::get_relay_inner(query, parsed_relays, &pure_config)
Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
}
}
}
Expand All @@ -527,30 +527,6 @@ impl RelaySelector {
self.get_relay_with_custom_params(retry_attempt, &RETRY_ORDER, runtime_params)
}

/// Peek at which [`TunnelType`] that would be returned for a certain connection attempt for a
/// given [`SelectorConfig`]. Returns [`Option::None`] if the given config would return a
/// custom tunnel endpoint.
///
/// # Note
/// This function is only really useful for testing-purposes. It is exposed to ease testing of
/// other mullvad crates which depend on the retry behaviour of [`RelaySelector`].
pub fn would_return(connection_attempt: usize, config: &SelectorConfig) -> Option<TunnelType> {
match SpecializedSelectorConfig::from(config) {
// This case is not really interesting
SpecializedSelectorConfig::Custom(_) => None,
SpecializedSelectorConfig::Normal(config) => Some(
Self::pick_and_merge_query(
connection_attempt,
&RETRY_ORDER,
RuntimeParameters::default(),
RelayQuery::from(config),
)
.tunnel_protocol
.unwrap_or(TunnelType::Wireguard),
),
}
}

/// Returns a random relay and relay endpoint matching the current constraints defined by
/// `retry_order` corresponding to `retry_attempt`.
pub fn get_relay_with_custom_params(
Expand All @@ -571,14 +547,14 @@ impl RelaySelector {
SpecializedSelectorConfig::Normal(normal_config) => {
let parsed_relays = &self.parsed_relays.lock().unwrap();
// Merge user preferences with the relay selector's default preferences.
let user_preferences = RelayQuery::from(normal_config.clone());
let query = Self::pick_and_merge_query(
retry_attempt,
retry_order,
runtime_params,
user_preferences,
);
Self::get_relay_inner(query, parsed_relays, &normal_config)
&normal_config,
parsed_relays,
)?;
Self::get_relay_inner(&query, parsed_relays, normal_config.custom_lists)
}
}
}
Expand All @@ -593,22 +569,29 @@ impl RelaySelector {
/// Runtime parameters may affect which of the default queries that are considered. For example,
/// queries which rely on IPv6 will not be considered if working IPv6 is not available at
/// runtime.
///
/// Returns an error iff the intersection between the user's preferences and every default retry
/// attempt-query yields queries with no matching relays. I.e., no retry attempt could ever
/// resolve to a relay.
fn pick_and_merge_query(
retry_attempt: usize,
retry_order: &[RelayQuery],
runtime_params: RuntimeParameters,
user_preferences: RelayQuery,
) -> RelayQuery {
log::trace!("Merging user preferences {user_preferences:?} with default retry strategy");
user_config: &NormalSelectorConfig<'_>,
parsed_relays: &ParsedRelays,
) -> Result<RelayQuery, Error> {
let user_query = RelayQuery::from(user_config.clone());
log::trace!("Merging user preferences {user_query:?} with default retry strategy");
retry_order
.iter()
// Remove candidate queries based on runtime parameters before trying to merge user
// settings
.filter(|query| runtime_params.compatible(query))
.cycle()
.filter_map(|query| query.clone().intersection(user_preferences.clone()))
.filter_map(|query| query.clone().intersection(user_query.clone()))
.filter(|query| Self::get_relay_inner(query, parsed_relays, user_config.custom_lists).is_ok())
.cycle() // If the above filters remove all relays, cycle will also return an empty iterator
.nth(retry_attempt)
.unwrap_or(user_preferences)
.ok_or(Error::NoRelay)
}

/// "Execute" the given query, yielding a final set of relays and/or bridges which the VPN
Expand All @@ -629,24 +612,26 @@ impl RelaySelector {
/// * An `Err` if no suitable bridge is found
#[cfg(not(target_os = "android"))]
fn get_relay_inner(
query: RelayQuery,
query: &RelayQuery,
parsed_relays: &ParsedRelays,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
) -> Result<GetRelay, Error> {
match query.tunnel_protocol {
Constraint::Only(TunnelType::Wireguard) => {
Self::get_wireguard_relay(query, config, parsed_relays)
Self::get_wireguard_relay(query, custom_lists, parsed_relays)
}
Constraint::Only(TunnelType::OpenVpn) => {
Self::get_openvpn_relay(query, config, parsed_relays)
Self::get_openvpn_relay(query, custom_lists, parsed_relays)
}
Constraint::Any => {
// Try Wireguard, then OpenVPN, then fail
for tunnel_type in [TunnelType::Wireguard, TunnelType::OpenVpn] {
let mut new_query = query.clone();
new_query.tunnel_protocol = Constraint::Only(tunnel_type);
// If a suitable relay is found, short-circuit and return it
if let Ok(relay) = Self::get_relay_inner(new_query, parsed_relays, config) {
if let Ok(relay) =
Self::get_relay_inner(&new_query, parsed_relays, custom_lists)
{
return Ok(relay);
}
}
Expand All @@ -657,16 +642,17 @@ impl RelaySelector {

#[cfg(target_os = "android")]
fn get_relay_inner(
mut query: RelayQuery,
query: &RelayQuery,
parsed_relays: &ParsedRelays,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
) -> Result<GetRelay, Error> {
// FIXME: A bit of defensive programming - calling `get_wiregurad_relay` with a query that
// doesn't specify Wireguard as the desired tunnel type is not valid and will lead
// to unwanted behavior. This should be seen as a workaround, and it would be nicer
// to lift this invariant to be checked by the type system instead.
let mut query = query.clone();
query.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
Self::get_wireguard_relay(query, config, parsed_relays)
Self::get_wireguard_relay(&query, custom_lists, parsed_relays)
}

/// Derive a valid Wireguard relay configuration from `query`.
Expand All @@ -683,22 +669,22 @@ impl RelaySelector {
///
/// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
fn get_wireguard_relay(
query: RelayQuery,
config: &NormalSelectorConfig<'_>,
query: &RelayQuery,
custom_lists: &CustomListsSettings,
parsed_relays: &ParsedRelays,
) -> Result<GetRelay, Error> {
assert_eq!(
query.tunnel_protocol,
Constraint::Only(TunnelType::Wireguard)
);
let inner = if !query.wireguard_constraints.multihop() {
Self::get_wireguard_singlehop_config(&query, config, parsed_relays)?
Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays)?
} else {
Self::get_wireguard_multihop_config(&query, config, parsed_relays)?
Self::get_wireguard_multihop_config(query, custom_lists, parsed_relays)?
};
let endpoint = Self::get_wireguard_endpoint(&query, parsed_relays, &inner)?;
let endpoint = Self::get_wireguard_endpoint(query, parsed_relays, &inner)?;
let obfuscator =
Self::get_wireguard_obfuscator(&query, inner.clone(), &endpoint, parsed_relays)?;
Self::get_wireguard_obfuscator(query, inner.clone(), &endpoint, parsed_relays)?;

Ok(GetRelay::Wireguard {
endpoint,
Expand All @@ -714,11 +700,10 @@ impl RelaySelector {
/// * `Ok(WireguardInner::Singlehop)` otherwise
fn get_wireguard_singlehop_config(
query: &RelayQuery,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
parsed_relays: &ParsedRelays,
) -> Result<WireguardConfig, Error> {
let candidates =
filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists);
let candidates = filter_matching_relay_list(query, parsed_relays.relays(), custom_lists);
helpers::pick_random_relay(&candidates)
.cloned()
.map(WireguardConfig::singlehop)
Expand All @@ -734,7 +719,7 @@ impl RelaySelector {
/// * `Ok(WireguardInner::Multihop)` otherwise
fn get_wireguard_multihop_config(
query: &RelayQuery,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
parsed_relays: &ParsedRelays,
) -> Result<WireguardConfig, Error> {
// Here, we modify the original query just a bit.
Expand All @@ -749,16 +734,10 @@ impl RelaySelector {
let mut exit_relay_query = query.clone();
// DAITA should only be enabled for the entry relay
exit_relay_query.wireguard_constraints.daita = Constraint::Only(false);
let exit_candidates = filter_matching_relay_list(
&exit_relay_query,
parsed_relays.relays(),
config.custom_lists,
);
let entry_candidates = filter_matching_relay_list(
&entry_relay_query,
parsed_relays.relays(),
config.custom_lists,
);
let exit_candidates =
filter_matching_relay_list(&exit_relay_query, parsed_relays.relays(), custom_lists);
let entry_candidates =
filter_matching_relay_list(&entry_relay_query, parsed_relays.relays(), custom_lists);

fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
list.iter()
Expand Down Expand Up @@ -846,16 +825,21 @@ impl RelaySelector {
/// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
#[cfg(not(target_os = "android"))]
fn get_openvpn_relay(
query: RelayQuery,
config: &NormalSelectorConfig<'_>,
query: &RelayQuery,
custom_lists: &CustomListsSettings,
parsed_relays: &ParsedRelays,
) -> Result<GetRelay, Error> {
assert_eq!(query.tunnel_protocol, Constraint::Only(TunnelType::OpenVpn));
let exit =
Self::choose_openvpn_relay(&query, config, parsed_relays).ok_or(Error::NoRelay)?;
let endpoint = Self::get_openvpn_endpoint(&query, &exit, parsed_relays)?;
let bridge =
Self::get_openvpn_bridge(&query, &exit, &endpoint.protocol, parsed_relays, config)?;
Self::choose_openvpn_relay(query, custom_lists, parsed_relays).ok_or(Error::NoRelay)?;
let endpoint = Self::get_openvpn_endpoint(query, &exit, parsed_relays)?;
let bridge = Self::get_openvpn_bridge(
query,
&exit,
&endpoint.protocol,
parsed_relays,
custom_lists,
)?;

// FIXME: This assert would be better to encode at the type level.
assert!(matches!(exit.endpoint_data, RelayEndpointData::Openvpn));
Expand Down Expand Up @@ -908,13 +892,13 @@ impl RelaySelector {
relay: &Relay,
protocol: &TransportProtocol,
parsed_relays: &ParsedRelays,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
) -> Result<Option<SelectedBridge>, Error> {
if !BridgeQuery::should_use_bridge(&query.openvpn_constraints.bridge_settings) {
Ok(None)
} else {
let bridge_query = &query.openvpn_constraints.bridge_settings.clone().unwrap();
let custom_lists = &config.custom_lists;
let custom_lists = &custom_lists;
match protocol {
TransportProtocol::Udp => {
log::error!("Can not use OpenVPN bridges over UDP");
Expand Down Expand Up @@ -1032,15 +1016,15 @@ impl RelaySelector {
fn get_relay_midpoint(
query: &RelayQuery,
parsed_relays: &ParsedRelays,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
) -> Option<Coordinates> {
use std::ops::Not;
if query.location.is_any() {
return None;
}

let matching_locations: Vec<Location> =
filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists)
filter_matching_relay_list(query, parsed_relays.relays(), custom_lists)
.into_iter()
.filter_map(|relay| relay.location)
.unique_by(|location| location.city.clone())
Expand All @@ -1058,12 +1042,12 @@ impl RelaySelector {
#[cfg(not(target_os = "android"))]
fn choose_openvpn_relay(
query: &RelayQuery,
config: &NormalSelectorConfig<'_>,
custom_lists: &CustomListsSettings,
parsed_relays: &ParsedRelays,
) -> Option<Relay> {
// Filter among all valid relays
let relays = parsed_relays.relays();
let candidates = filter_matching_relay_list(query, relays, config.custom_lists);
let candidates = filter_matching_relay_list(query, relays, custom_lists);
// Pick one of the valid relays.
helpers::pick_random_relay(&candidates).cloned()
}
Expand Down
27 changes: 27 additions & 0 deletions mullvad-relay-selector/tests/relay_selector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1296,3 +1296,30 @@ fn test_daita() {
),
}
}

/// Check that if the original user query would yield a relay, the result of running the query
/// which is the intersection between the user query and any of the default queries shall never
/// fail.
#[test]
fn valid_user_setting_should_yield_relay() {
// Make a valid user relay constraint
let location = GeographicLocationConstraint::hostname("se", "got", "se9-wireguard");
let user_query = RelayQueryBuilder::new().location(location.clone()).build();
let user_constraints = RelayQueryBuilder::new()
.location(location.clone())
.into_constraint();

let config = SelectorConfig {
relay_settings: user_constraints.into(),
..SelectorConfig::default()
};
let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
let user_result = relay_selector.get_relay_by_query(user_query.clone());
for retry_attempt in 0..RETRY_ORDER.len() {
let post_unification_result =
relay_selector.get_relay(retry_attempt, RuntimeParameters::default());
if user_result.is_ok() {
assert!(post_unification_result.is_ok(), "Expected Post-unification query to be valid because original query {:#?} yielded a connection configuration", user_query)
}
}
}

0 comments on commit 52b70e8

Please sign in to comment.