Return obfuscation config along with relay

2 files changed, 80 insertions, 32 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 106f3c1784..19f17e3e45 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -33,7 +33,7 @@ use futures::{
 use mullvad_api::availability::ApiAvailabilityHandle;
 use mullvad_relay_selector::{
     updater::{RelayListUpdater, RelayListUpdaterHandle},
-    RelaySelector, SelectedBridge, SelectedRelay, SelectorConfig,
+    RelaySelector, SelectedBridge, SelectedObfuscator, SelectedRelay, SelectorConfig,
 };
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
@@ -156,9 +156,6 @@ pub enum Error {
     #[error(display = "No bridge available")]
     NoBridgeAvailable,
 
-    #[error(display = "Failed to select a compatible obfuscator")]
-    NoObfuscator,
-
     #[error(display = "No matching entry relay was found")]
     NoEntryRelayAvailable,
 
@@ -1028,7 +1025,7 @@ where
         };
 
         let result = match self.relay_selector.get_relay(retry_attempt) {
-            Ok((SelectedRelay::Custom(custom_relay), _bridge)) => {
+            Ok((SelectedRelay::Custom(custom_relay), _bridge, _obfsucator)) => {
                 custom_relay
                     // TODO(emilsp): generate proxy settings for custom tunnels
                     .to_tunnel_parameters(self.settings.tunnel_options.clone(), None)
@@ -1037,15 +1034,15 @@ where
                         ParameterGenerationError::CustomTunnelHostResultionError
                     })
             }
-            Ok((SelectedRelay::Normal(constraints), bridge)) => {
+            Ok((SelectedRelay::Normal(constraints), bridge, obfuscator)) => {
                 let result = self
                     .create_tunnel_parameters(
                         &constraints.exit_relay,
                         &constraints.entry_relay,
                         constraints.endpoint,
                         bridge,
+                        obfuscator,
                         data,
-                        retry_attempt,
                     )
                     .await;
                 result.map_err(|error| match error {
@@ -1077,8 +1074,8 @@ where
         entry_relay: &Option<Relay>,
         endpoint: MullvadEndpoint,
         bridge: Option<SelectedBridge>,
+        obfuscator: Option<SelectedObfuscator>,
         device: DeviceData,
-        retry_attempt: u32,
     ) -> Result<TunnelParameters, Error> {
         let tunnel_options = self.settings.tunnel_options.clone();
         match endpoint {
@@ -1118,18 +1115,15 @@ where
                     ],
                 };
 
-                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)?;
+                let (obfuscator_relay, obfuscator_config) = match obfuscator {
+                    Some(obfuscator) => (Some(obfuscator.relay), Some(obfuscator.config)),
+                    None => (None, None),
+                };
 
                 self.last_generated_relays = Some(LastSelectedRelays::WireGuard {
                     wg_entry: entry_relay.clone(),
                     wg_exit: relay.clone(),
-                    obfuscator: selected_obfuscator
-                        .as_ref()
-                        .map(|_| obfuscator_relay.clone()),
+                    obfuscator: obfuscator_relay,
                 });
 
                 Ok(wireguard::TunnelParameters {
@@ -1142,7 +1136,7 @@ where
                     },
                     options: tunnel_options.wireguard.options,
                     generic_options: tunnel_options.generic,
-                    obfuscation: selected_obfuscator,
+                    obfuscation: obfuscator_config,
                 }
                 .into())
             }
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 9b257e7c74..48464f5616 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -262,11 +262,18 @@ impl RelaySelector {
     pub fn get_relay(
         &self,
         retry_attempt: u32,
-    ) -> Result<(SelectedRelay, Option<SelectedBridge>), Error> {
+    ) -> Result<
+        (
+            SelectedRelay,
+            Option<SelectedBridge>,
+            Option<SelectedObfuscator>,
+        ),
+        Error,
+    > {
         let config = self.config.lock();
         match &config.relay_settings {
             RelaySettings::CustomTunnelEndpoint(custom_relay) => {
-                Ok((SelectedRelay::Custom(custom_relay.clone()), None))
+                Ok((SelectedRelay::Custom(custom_relay.clone()), None, None))
             }
             RelaySettings::Normal(constraints) => {
                 let relay =
@@ -280,11 +287,24 @@ impl RelaySelector {
                             .location
                             .as_ref()
                             .expect("Relay has no location set");
-                        self.get_bridge_for(config, location, retry_attempt)?
+                        self.get_bridge_for(&config, location, retry_attempt)?
                     }
                     _ => None,
                 };
-                Ok((SelectedRelay::Normal(relay), bridge))
+                let obfuscator = match relay.endpoint {
+                    MullvadEndpoint::Wireguard(ref endpoint) => {
+                        let obfuscator_relay =
+                            relay.entry_relay.as_ref().unwrap_or(&relay.exit_relay);
+                        self.get_obfuscator_inner(
+                            &config,
+                            obfuscator_relay,
+                            &endpoint,
+                            retry_attempt,
+                        )?
+                    }
+                    _ => None,
+                };
+                Ok((SelectedRelay::Normal(relay), bridge, obfuscator))
             }
         }
     }
@@ -667,7 +687,7 @@ impl RelaySelector {
 
     fn get_bridge_for(
         &self,
-        config: MutexGuard<'_, SelectorConfig>,
+        config: &MutexGuard<'_, SelectorConfig>,
         location: &mullvad_types::location::Location,
         retry_attempt: u32,
     ) -> Result<Option<SelectedBridge>, Error> {
@@ -786,9 +806,17 @@ impl RelaySelector {
         relay: &Relay,
         endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Result<Option<ObfuscatorConfig>, Error> {
-        let config = self.config.lock();
+    ) -> Result<Option<SelectedObfuscator>, Error> {
+        self.get_obfuscator_inner(&self.config.lock(), relay, endpoint, retry_attempt)
+    }
 
+    fn get_obfuscator_inner(
+        &self,
+        config: &MutexGuard<'_, SelectorConfig>,
+        relay: &Relay,
+        endpoint: &MullvadWireguardEndpoint,
+        retry_attempt: u32,
+    ) -> Result<Option<SelectedObfuscator>, Error> {
         match &config.obfuscation_settings.selected_obfuscation {
             SelectedObfuscation::Auto => Ok(self.get_auto_obfuscator(
                 &config.obfuscation_settings,
@@ -815,7 +843,7 @@ impl RelaySelector {
         relay: &Relay,
         endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Option<ObfuscatorConfig> {
+    ) -> Option<SelectedObfuscator> {
         if !self.should_use_auto_obfuscator(retry_attempt) {
             return None;
         }
@@ -849,7 +877,7 @@ impl RelaySelector {
         relay: &Relay,
         _endpoint: &MullvadWireguardEndpoint,
         retry_attempt: u32,
-    ) -> Option<ObfuscatorConfig> {
+    ) -> Option<SelectedObfuscator> {
         let udp2tcp_endpoint = if obfuscation_settings.port.is_only() {
             relay
                 .obfuscators
@@ -862,9 +890,14 @@ 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),
-        })
+        udp2tcp_endpoint
+            .map(|udp2tcp_endpoint| ObfuscatorConfig::Udp2Tcp {
+                endpoint: SocketAddr::new(relay.ipv4_addr_in.into(), udp2tcp_endpoint.port),
+            })
+            .map(|config| SelectedObfuscator {
+                config,
+                relay: relay.clone(),
+            })
     }
 
     /// Returns preferred constraints
@@ -1119,6 +1152,12 @@ pub struct NormalSelectedRelay {
     pub entry_relay: Option<Relay>,
 }
 
+#[derive(Debug)]
+pub struct SelectedObfuscator {
+    pub config: ObfuscatorConfig,
+    pub relay: Relay,
+}
+
 impl NormalSelectedRelay {
     fn new(endpoint: MullvadEndpoint, exit_relay: Relay) -> Self {
         Self {
@@ -1690,7 +1729,13 @@ mod test {
             .unwrap()
             .unwrap();
 
-        assert!(matches!(obfs_config, ObfuscatorConfig::Udp2Tcp { .. }));
+        assert!(matches!(
+            obfs_config,
+            SelectedObfuscator {
+                config: ObfuscatorConfig::Udp2Tcp { .. },
+                ..
+            }
+        ));
     }
 
     #[test]
@@ -1750,9 +1795,18 @@ mod test {
                 .unwrap()
                 .expect("Failed to get Tcp2Udp endpoint");
 
-            assert!(matches!(obfs_config, ObfuscatorConfig::Udp2Tcp { .. }));
+            assert!(matches!(
+                obfs_config,
+                SelectedObfuscator {
+                    config: ObfuscatorConfig::Udp2Tcp { .. },
+                    ..
+                }
+            ));
 
-            let ObfuscatorConfig::Udp2Tcp { endpoint } = obfs_config;
+            let SelectedObfuscator {
+                config: ObfuscatorConfig::Udp2Tcp { endpoint },
+                ..
+            } = obfs_config;
             assert!(TCP2UDP_PORTS.contains(&endpoint.port()));
         }
     }