Manage relay, bridge, and obfuscation settings in the relay selector

2 files changed, 202 insertions, 163 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 97c87914c2..9db0afd719 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -36,7 +36,7 @@ use mullvad_api::{
 };
 use mullvad_relay_selector::{
     updater::{RelayListUpdater, RelayListUpdaterHandle},
-    RelaySelector, RelaySelectorResult,
+    RelaySelector, RelaySelectorResult, SelectedBridge, SelectedRelay, SelectorConfig,
 };
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
@@ -44,8 +44,9 @@ use mullvad_types::{
     endpoint::MullvadEndpoint,
     location::{Coordinates, GeoIpLocation},
     relay_constraints::{
-        BridgeSettings, BridgeState, Constraint, InternalBridgeConstraints, ObfuscationSettings,
-        RelaySettings, RelaySettingsUpdate, SelectedObfuscation,
+        BridgeConstraints, BridgeSettings, BridgeState, Constraint, InternalBridgeConstraints,
+        ObfuscationSettings, RelayConstraints, RelaySettings, RelaySettingsUpdate,
+        SelectedObfuscation,
     },
     relay_list::{Relay, RelayList},
     settings::{DnsOptions, DnsState, Settings},
@@ -758,7 +759,9 @@ where
             relay_list_listener.notify_relay_list(relay_list.clone());
         };
 
-        let relay_selector = RelaySelector::new(&resource_dir, &cache_dir);
+        let initial_selector_config = new_selector_config(&settings);
+        let relay_selector = RelaySelector::new(initial_selector_config, &resource_dir, &cache_dir);
+
         let mut relay_list_updater = RelayListUpdater::new(
             relay_selector.clone(),
             api_handle.clone(),
@@ -1039,8 +1042,8 @@ where
         retry_attempt: u32,
     ) {
         if let Ok(Some(device)) = self.account_manager.data().await {
-            let result = match self.settings.get_relay_settings() {
-                RelaySettings::CustomTunnelEndpoint(custom_relay) => {
+            let result = match self.relay_selector.get_relay(retry_attempt) {
+                Ok(SelectedRelay::Custom(custom_relay)) => {
                     custom_relay
                         // TODO(emilsp): generate proxy settings for custom tunnels
                         .to_tunnel_parameters(self.settings.tunnel_options.clone(), None)
@@ -1049,52 +1052,30 @@ where
                             ParameterGenerationError::CustomTunnelHostResultionError
                         })
                 }
-                RelaySettings::Normal(constraints) => {
-                    let endpoint = self
-                        .relay_selector
-                        .get_tunnel_endpoint(
-                            &constraints,
-                            self.settings.get_bridge_state(),
+                Ok(SelectedRelay::Normal(constraints)) => {
+                    let result = self
+                        .create_tunnel_parameters(
+                            &constraints.exit_relay,
+                            &constraints.entry_relay,
+                            constraints.endpoint,
+                            device.token,
                             retry_attempt,
                         )
-                        .ok();
-                    if let Some(RelaySelectorResult {
-                        exit_relay,
-                        entry_relay,
-                        endpoint,
-                    }) = endpoint
-                    {
-                        let result = self
-                            .create_tunnel_parameters(
-                                &exit_relay,
-                                &entry_relay,
-                                endpoint,
-                                device.token,
-                                retry_attempt,
-                            )
-                            .await;
-                        match result {
-                            Ok(result) => Ok(result),
-                            Err(Error::NoKeyAvailable) => {
-                                Err(ParameterGenerationError::NoWireguardKey)
-                            }
-                            Err(Error::NoBridgeAvailable) => {
-                                Err(ParameterGenerationError::NoMatchingBridgeRelay)
-                            }
-                            Err(err) => {
-                                log::error!(
-                                    "{}",
-                                    err.display_chain_with_msg(
-                                        "Failed to generate tunnel parameters"
-                                    )
-                                );
-                                Err(ParameterGenerationError::NoMatchingRelay)
-                            }
+                        .await;
+                    result.map_err(|error| match error {
+                        Error::NoKeyAvailable => ParameterGenerationError::NoWireguardKey,
+                        Error::NoBridgeAvailable => ParameterGenerationError::NoMatchingBridgeRelay,
+                        error => {
+                            log::error!(
+                                "{}",
+                                error
+                                    .display_chain_with_msg("Failed to generate tunnel parameters")
+                            );
+                            ParameterGenerationError::NoMatchingRelay
                         }
-                    } else {
-                        Err(ParameterGenerationError::NoMatchingRelay)
-                    }
+                    })
                 }
+                Err(_error) => Err(ParameterGenerationError::NoMatchingRelay),
             };
             if tunnel_parameters_tx.send(result).is_err() {
                 log::error!("Failed to send tunnel parameters");
@@ -1118,8 +1099,17 @@ where
         match endpoint {
             #[cfg(not(target_os = "android"))]
             MullvadEndpoint::OpenVpn(endpoint) => {
-                let (bridge_settings, bridge_relay) =
-                    self.generate_bridge_parameters(&location, retry_attempt)?;
+                let (bridge_settings, bridge_relay) = match self
+                    .relay_selector
+                    .get_bridge(location, retry_attempt)
+                    .map_err(|_error| Error::NoBridgeAvailable)?
+                {
+                    Some(SelectedBridge::Normal(bridge)) => {
+                        (Some(bridge.settings), Some(bridge.relay))
+                    }
+                    Some(SelectedBridge::Custom(settings)) => (Some(settings), None),
+                    None => (None, None),
+                };
 
                 self.last_generated_relays = Some(LastSelectedRelays::OpenVpn {
                     relay: relay.clone(),
@@ -1158,38 +1148,18 @@ where
                     ],
                 };
 
-                let selected_obfuscator =
-                    match self.settings.obfuscation_settings.selected_obfuscation {
-                        SelectedObfuscation::Off => None,
-                        SelectedObfuscation::Auto
-                            if !self
-                                .relay_selector
-                                .should_use_auto_obfuscator(retry_attempt) =>
-                        {
-                            None
-                        }
-                        _ => {
-                            let (obfuscator_config, obfuscator_relay) = self
-                                .relay_selector
-                                .get_obfuscator(
-                                    &self.settings.obfuscation_settings,
-                                    entry_relay.as_ref().unwrap_or(relay),
-                                    &endpoint,
-                                    retry_attempt,
-                                )
-                                .ok_or(Error::NoObfuscator)?;
-                            Some((obfuscator_config, obfuscator_relay))
-                        }
-                    };
-                let (obfuscation, obfuscator_relay) = match selected_obfuscator {
-                    Some((obfuscation, relay)) => (Some(obfuscation), Some(relay)),
-                    None => (None, None),
-                };
+                let obfuscator_relay = entry_relay.as_ref().unwrap_or(relay);
+                let selected_obfuscator = self
+                    .relay_selector
+                    .get_obfuscator(obfuscator_relay, &endpoint, retry_attempt)
+                    .map_err(|_error| Error::NoObfuscator)?;
 
                 self.last_generated_relays = Some(LastSelectedRelays::WireGuard {
                     wg_entry: entry_relay.clone(),
                     wg_exit: relay.clone(),
-                    obfuscator: obfuscator_relay,
+                    obfuscator: selected_obfuscator
+                        .as_ref()
+                        .map(|_| obfuscator_relay.clone()),
                 });
 
                 Ok(wireguard::TunnelParameters {
@@ -1202,67 +1172,13 @@ where
                     },
                     options: tunnel_options.wireguard.options,
                     generic_options: tunnel_options.generic,
-                    obfuscation,
+                    obfuscation: selected_obfuscator,
                 }
                 .into())
             }
         }
     }
 
-    /// Generates bridge parameters for the chosen OpenVpn tunnel relay
-    #[cfg(not(target_os = "android"))]
-    fn generate_bridge_parameters(
-        &self,
-        location: &mullvad_types::location::Location,
-        retry_attempt: u32,
-    ) -> Result<(Option<ProxySettings>, Option<Relay>), Error> {
-        let result = match &self.settings.bridge_settings {
-            BridgeSettings::Normal(settings) => {
-                let bridge_constraints = InternalBridgeConstraints {
-                    location: settings.location.clone(),
-                    providers: settings.providers.clone(),
-                    // FIXME: This is temporary while talpid-core only supports TCP proxies
-                    transport_protocol: Constraint::Only(TransportProtocol::Tcp),
-                };
-                match self.settings.get_bridge_state() {
-                    BridgeState::On => {
-                        let (bridge_settings, bridge_relay) = self
-                            .relay_selector
-                            .get_proxy_settings(&bridge_constraints, Some(location))
-                            .ok_or(Error::NoBridgeAvailable)?;
-                        (Some(bridge_settings), Some(bridge_relay))
-                    }
-                    BridgeState::Auto => {
-                        if let Some((bridge_settings, bridge_relay)) =
-                            self.relay_selector.get_auto_proxy_settings(
-                                &bridge_constraints,
-                                Some(location),
-                                retry_attempt,
-                            )
-                        {
-                            (Some(bridge_settings), Some(bridge_relay))
-                        } else {
-                            (None, None)
-                        }
-                    }
-                    BridgeState::Off => (None, None),
-                }
-            }
-            BridgeSettings::Custom(bridge_settings) => match self.settings.get_bridge_state() {
-                BridgeState::On => (Some(bridge_settings.clone()), None),
-                BridgeState::Auto => {
-                    if self.relay_selector.should_use_bridge(retry_attempt) {
-                        (Some(bridge_settings.clone()), None)
-                    } else {
-                        (None, None)
-                    }
-                }
-                BridgeState::Off => (None, None),
-            },
-        };
-        Ok(result)
-    }
-
     fn schedule_reconnect(&mut self, delay: Duration) {
         self.unschedule_reconnect();
 
@@ -2182,6 +2098,8 @@ where
                 if settings_changed {
                     self.event_listener
                         .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
                     log::info!("Initiating tunnel restart because the relay settings changed");
                     self.reconnect_tunnel();
                 }
@@ -2319,6 +2237,8 @@ where
                 if settings_changes {
                     self.event_listener
                         .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
                     self.reconnect_tunnel();
                 };
                 Self::oneshot_send(tx, Ok(()), "set_bridge_settings");
@@ -2344,6 +2264,8 @@ where
                 if settings_changed {
                     self.event_listener
                         .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
                     self.reconnect_tunnel();
                 }
                 Self::oneshot_send(tx, Ok(()), "set_obfuscation_settings");
@@ -2368,6 +2290,8 @@ where
                 if settings_changed {
                     self.event_listener
                         .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
                     log::info!("Initiating tunnel restart because bridge state changed");
                     self.reconnect_tunnel();
                 }
@@ -2784,6 +2708,15 @@ impl LastSelectedRelays {
     }
 }
 
+fn new_selector_config(settings: &Settings) -> SelectorConfig {
+    SelectorConfig {
+        relay_settings: settings.get_relay_settings(),
+        bridge_state: settings.get_bridge_state(),
+        bridge_settings: settings.bridge_settings.clone(),
+        obfuscation_settings: settings.obfuscation_settings.clone(),
+    }
+}
+
 /// Bump filehandle limit
 #[cfg(target_os = "macos")]
 pub fn bump_filehandle_limit() {
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 7375de9c26..988e5a2368 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -7,11 +7,12 @@ use mullvad_types::{
     endpoint::{MullvadEndpoint, MullvadWireguardEndpoint},
     location::{Coordinates, Location},
     relay_constraints::{
-        BridgeState, Constraint, InternalBridgeConstraints, LocationConstraint, Match,
-        ObfuscationSettings, OpenVpnConstraints, Providers, RelayConstraints, SelectedObfuscation,
-        Set, TransportPort, Udp2TcpObfuscationSettings, WireguardConstraints,
+        BridgeSettings, BridgeState, Constraint, InternalBridgeConstraints, LocationConstraint,
+        Match, ObfuscationSettings, OpenVpnConstraints, Providers, RelayConstraints, RelaySettings,
+        SelectedObfuscation, Set, TransportPort, Udp2TcpObfuscationSettings, WireguardConstraints,
     },
     relay_list::{Relay, RelayList, Udp2TcpEndpointData},
+    CustomTunnelEndpoint,
 };
 use parking_lot::Mutex;
 use rand::{self, seq::SliceRandom, Rng};
@@ -64,6 +65,12 @@ pub enum Error {
     #[error(display = "No relays matching current constraints")]
     NoRelay,
 
+    #[error(display = "No bridges matching current constraints")]
+    NoBridge,
+
+    #[error(display = "No obfuscators matching current constraints")]
+    NoObfuscator,
+
     #[error(display = "Failure in serialization of the relay list")]
     Serialize(#[error(source)] serde_json::Error),
 
@@ -202,13 +209,22 @@ impl ParsedRelays {
 }
 
 #[derive(Clone)]
+pub struct SelectorConfig {
+    pub relay_settings: RelaySettings,
+    pub bridge_state: BridgeState,
+    pub bridge_settings: BridgeSettings,
+    pub obfuscation_settings: ObfuscationSettings,
+}
+
+#[derive(Clone)]
 pub struct RelaySelector {
+    config: Arc<Mutex<SelectorConfig>>,
     parsed_relays: Arc<Mutex<ParsedRelays>>,
 }
 
 impl RelaySelector {
     /// Returns a new `RelaySelector` backed by relays cached on disk.
-    pub fn new(resource_dir: &Path, cache_dir: &Path) -> Self {
+    pub fn new(config: SelectorConfig, resource_dir: &Path, cache_dir: &Path) -> Self {
         let cache_path = cache_dir.join(RELAYS_FILENAME);
         let resource_path = resource_dir.join(RELAYS_FILENAME);
         let unsynchronized_parsed_relays = Self::read_relays_from_disk(&cache_path, &resource_path)
@@ -225,9 +241,15 @@ impl RelaySelector {
             DateTime::<Local>::from(unsynchronized_parsed_relays.last_updated())
                 .format(DATE_TIME_FORMAT_STR)
         );
-        let parsed_relays = Arc::new(Mutex::new(unsynchronized_parsed_relays));
 
-        RelaySelector { parsed_relays }
+        RelaySelector {
+            config: Arc::new(Mutex::new(config)),
+            parsed_relays: Arc::new(Mutex::new(unsynchronized_parsed_relays)),
+        }
+    }
+
+    pub fn set_config(&mut self, config: SelectorConfig) {
+        *self.config.lock() = config;
     }
 
     /// Returns all countries and cities. The cities in the object returned does not have any
@@ -236,9 +258,22 @@ impl RelaySelector {
         self.parsed_relays.lock().locations().clone()
     }
 
+    /// Returns a random relay and relay endpoint matching the current constraints.
+    pub fn get_relay(&self, retry_attempt: u32) -> Result<SelectedRelay, Error> {
+        let config = self.config.lock();
+        match &config.relay_settings {
+            RelaySettings::CustomTunnelEndpoint(custom_relay) => {
+                Ok(SelectedRelay::Custom(custom_relay.clone()))
+            }
+            RelaySettings::Normal(constraints) => self
+                .get_tunnel_endpoint(&constraints, config.bridge_state, retry_attempt)
+                .map(SelectedRelay::Normal),
+        }
+    }
+
     /// Returns a random relay and relay endpoint matching the given constraints and with
     /// preferences applied.
-    pub fn get_tunnel_endpoint(
+    fn get_tunnel_endpoint(
         &self,
         relay_constraints: &RelayConstraints,
         bridge_state: BridgeState,
@@ -612,8 +647,59 @@ impl RelaySelector {
         entry_endpoint.exit_peer = Some(exit_peer.clone());
     }
 
+    pub fn get_bridge(
+        &self,
+        location: &mullvad_types::location::Location,
+        retry_attempt: u32,
+    ) -> Result<Option<SelectedBridge>, Error> {
+        let config = self.config.lock();
+        match &config.bridge_settings {
+            BridgeSettings::Normal(settings) => {
+                let bridge_constraints = InternalBridgeConstraints {
+                    location: settings.location.clone(),
+                    providers: settings.providers.clone(),
+                    // FIXME: This is temporary while talpid-core only supports TCP proxies
+                    transport_protocol: Constraint::Only(TransportProtocol::Tcp),
+                };
+                match config.bridge_state {
+                    BridgeState::On => {
+                        let (settings, relay) = self
+                            .get_proxy_settings(&bridge_constraints, Some(location))
+                            .ok_or(Error::NoBridge)?;
+                        Ok(Some(SelectedBridge::Normal(NormalSelectedBridge {
+                            settings,
+                            relay,
+                        })))
+                    }
+                    BridgeState::Auto => {
+                        if let Some((settings, relay)) = self.get_auto_proxy_settings(
+                            &bridge_constraints,
+                            Some(location),
+                            retry_attempt,
+                        ) {
+                            Ok(Some(SelectedBridge::Normal(NormalSelectedBridge {
+                                settings,
+                                relay,
+                            })))
+                        } else {
+                            Ok(None)
+                        }
+                    }
+                    BridgeState::Off => Ok(None),
+                }
+            }
+            BridgeSettings::Custom(bridge_settings) => match config.bridge_state {
+                BridgeState::On => Ok(Some(SelectedBridge::Custom(bridge_settings.clone()))),
+                BridgeState::Auto if self.should_use_bridge(retry_attempt) => {
+                    Ok(Some(SelectedBridge::Custom(bridge_settings.clone())))
+                }
+                BridgeState::Auto | BridgeState::Off => Ok(None),
+            },
+        }
+    }
+
     #[cfg(not(target_os = "android"))]
-    pub fn get_auto_proxy_settings<T: Into<Coordinates>>(
+    fn get_auto_proxy_settings<T: Into<Coordinates>>(
         &self,
         bridge_constraints: &InternalBridgeConstraints,
         location: Option<T>,
@@ -632,7 +718,7 @@ impl RelaySelector {
     }
 
     #[cfg(not(target_os = "android"))]
-    pub fn should_use_bridge(&self, retry_attempt: u32) -> bool {
+    fn should_use_bridge(&self, retry_attempt: u32) -> bool {
         // shouldn't use a bridge for the first 3 times
         retry_attempt > 3 &&
             // i.e. 4th and 5th with bridge, 6th & 7th without
@@ -642,7 +728,7 @@ impl RelaySelector {
             (retry_attempt % 4) < 2
     }
 
-    pub fn get_proxy_settings<T: Into<Coordinates>>(
+    fn get_proxy_settings<T: Into<Coordinates>>(
         &self,
         constraints: &InternalBridgeConstraints,
         location: Option<T>,
@@ -682,22 +768,29 @@ impl RelaySelector {
 
     pub fn get_obfuscator(
         &self,
-        obfuscation_settings: &ObfuscationSettings,
         relay: &Relay,
         endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Option<(ObfuscatorConfig, Relay)> {
-        match obfuscation_settings.selected_obfuscation {
-            SelectedObfuscation::Auto => {
-                self.get_auto_obfuscator(obfuscation_settings, relay, endpoint, retry_attempt)
-            }
-            SelectedObfuscation::Off => None,
-            SelectedObfuscation::Udp2Tcp => self.get_udp2tcp_obfuscator(
-                &obfuscation_settings.udp2tcp,
+    ) -> Result<Option<ObfuscatorConfig>, Error> {
+        let config = self.config.lock();
+
+        match &config.obfuscation_settings.selected_obfuscation {
+            SelectedObfuscation::Auto => Ok(self.get_auto_obfuscator(
+                &config.obfuscation_settings,
                 relay,
                 endpoint,
                 retry_attempt,
-            ),
+            )),
+            SelectedObfuscation::Off => Ok(None),
+            SelectedObfuscation::Udp2Tcp => Ok(Some(
+                self.get_udp2tcp_obfuscator(
+                    &config.obfuscation_settings.udp2tcp,
+                    relay,
+                    endpoint,
+                    retry_attempt,
+                )
+                .ok_or(Error::NoObfuscator)?,
+            )),
         }
     }
 
@@ -707,7 +800,7 @@ impl RelaySelector {
         relay: &Relay,
         endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Option<(ObfuscatorConfig, Relay)> {
+    ) -> Option<ObfuscatorConfig> {
         if !self.should_use_auto_obfuscator(retry_attempt) {
             return None;
         }
@@ -723,7 +816,7 @@ impl RelaySelector {
         )
     }
 
-    pub fn should_use_auto_obfuscator(&self, retry_attempt: u32) -> bool {
+    fn should_use_auto_obfuscator(&self, retry_attempt: u32) -> bool {
         self.get_auto_obfuscator_retry_attempt(retry_attempt)
             .is_some()
     }
@@ -741,7 +834,7 @@ impl RelaySelector {
         relay: &Relay,
         _endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Option<(ObfuscatorConfig, Relay)> {
+    ) -> Option<ObfuscatorConfig> {
         let udp2tcp_endpoint = if obfuscation_settings.port.is_only() {
             relay
                 .obfuscators
@@ -754,13 +847,8 @@ impl RelaySelector {
                 .udp2tcp
                 .get(retry_attempt as usize % relay.obfuscators.udp2tcp.len())
         };
-        udp2tcp_endpoint.map(|udp2tcp_endpoint| {
-            (
-                ObfuscatorConfig::Udp2Tcp {
-                    endpoint: SocketAddr::new(relay.ipv4_addr_in.into(), udp2tcp_endpoint.port),
-                },
-                relay.clone(),
-            )
+        udp2tcp_endpoint.map(|udp2tcp_endpoint| ObfuscatorConfig::Udp2Tcp {
+            endpoint: SocketAddr::new(relay.ipv4_addr_in.into(), udp2tcp_endpoint.port),
         })
     }
 
@@ -991,6 +1079,24 @@ impl RelaySelector {
 }
 
 #[derive(Debug)]
+pub enum SelectedBridge {
+    Normal(NormalSelectedBridge),
+    Custom(ProxySettings),
+}
+
+#[derive(Debug)]
+pub struct NormalSelectedBridge {
+    pub settings: ProxySettings,
+    pub relay: Relay,
+}
+
+#[derive(Debug)]
+pub enum SelectedRelay {
+    Normal(RelaySelectorResult),
+    Custom(CustomTunnelEndpoint),
+}
+
+#[derive(Debug)]
 pub struct RelaySelectorResult {
     pub exit_relay: Relay,
     pub endpoint: MullvadEndpoint,