Return bridge from get_relay()

3 files changed, 79 insertions, 71 deletions

diff --git a/mullvad-daemon/src/api.rs b/mullvad-daemon/src/api.rs
index 48870049c7..41ac64ea32 100644
--- a/mullvad-daemon/src/api.rs
+++ b/mullvad-daemon/src/api.rs
@@ -78,7 +78,7 @@ impl Stream for ApiConnectionModeProvider {
 
         let config = if Self::should_use_bridge(self.retry_attempt) {
             selector
-                .get_api_bridge()
+                .get_bridge_forced()
                 .map(|settings| match settings {
                     ProxySettings::Shadowsocks(ss_settings) => {
                         ApiConnectionMode::Proxied(ProxyConfig::Shadowsocks(ss_settings))
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 315ea37ffa..77c1be91fc 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1018,47 +1018,54 @@ where
         >,
         retry_attempt: u32,
     ) {
-        if let Ok(Some(device)) = self.account_manager.data().await {
-            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)
-                        .map_err(|e| {
-                            log::error!("Failed to resolve hostname for custom tunnel config: {}", e);
-                            ParameterGenerationError::CustomTunnelHostResultionError
-                        })
-                }
-                Ok(SelectedRelay::Normal(constraints)) => {
-                    let result = self
-                        .create_tunnel_parameters(
-                            &constraints.exit_relay,
-                            &constraints.entry_relay,
-                            constraints.endpoint,
-                            device.token,
-                            retry_attempt,
-                        )
-                        .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
-                        }
+        let data = match self.account_manager.data().await {
+            Ok(Some(data)) => data,
+            _ => {
+                log::error!("No account token configured");
+                return;
+            }
+        };
+
+        let result = match self.relay_selector.get_relay(retry_attempt) {
+            Ok((SelectedRelay::Custom(custom_relay), _bridge)) => {
+                custom_relay
+                    // TODO(emilsp): generate proxy settings for custom tunnels
+                    .to_tunnel_parameters(self.settings.tunnel_options.clone(), None)
+                    .map_err(|e| {
+                        log::error!("Failed to resolve hostname for custom tunnel config: {}", e);
+                        ParameterGenerationError::CustomTunnelHostResultionError
                     })
-                }
-                Err(_error) => Err(ParameterGenerationError::NoMatchingRelay),
-            };
-            if tunnel_parameters_tx.send(result).is_err() {
-                log::error!("Failed to send tunnel parameters");
             }
-        } else {
-            log::error!("No account token configured");
+            Ok((SelectedRelay::Normal(constraints), bridge)) => {
+                let result = self
+                    .create_tunnel_parameters(
+                        &constraints.exit_relay,
+                        &constraints.entry_relay,
+                        constraints.endpoint,
+                        bridge,
+                        data,
+                        retry_attempt,
+                    )
+                    .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
+                    }
+                })
+            }
+            Err(mullvad_relay_selector::Error::NoBridge) => {
+                Err(ParameterGenerationError::NoMatchingBridgeRelay)
+            }
+            Err(_error) => Err(ParameterGenerationError::NoMatchingRelay),
+        };
+        if tunnel_parameters_tx.send(result).is_err() {
+            log::error!("Failed to send tunnel parameters");
         }
     }
 
@@ -1068,19 +1075,15 @@ where
         relay: &Relay,
         entry_relay: &Option<Relay>,
         endpoint: MullvadEndpoint,
-        account_token: String,
+        bridge: Option<SelectedBridge>,
+        device: DeviceData,
         retry_attempt: u32,
     ) -> Result<TunnelParameters, Error> {
         let tunnel_options = self.settings.tunnel_options.clone();
-        let location = relay.location.as_ref().expect("Relay has no location set");
         match endpoint {
             #[cfg(not(target_os = "android"))]
             MullvadEndpoint::OpenVpn(endpoint) => {
-                let (bridge_settings, bridge_relay) = match self
-                    .relay_selector
-                    .get_bridge(location, retry_attempt)
-                    .map_err(|_error| Error::NoBridgeAvailable)?
-                {
+                let (bridge_settings, bridge_relay) = match bridge {
                     Some(SelectedBridge::Normal(bridge)) => {
                         (Some(bridge.settings), Some(bridge.relay))
                     }
@@ -1094,11 +1097,7 @@ where
                 });
 
                 Ok(openvpn::TunnelParameters {
-                    config: openvpn::ConnectionConfig::new(
-                        endpoint,
-                        account_token,
-                        "-".to_string(),
-                    ),
+                    config: openvpn::ConnectionConfig::new(endpoint, device.token, "-".to_string()),
                     options: tunnel_options.openvpn,
                     generic_options: tunnel_options.generic,
                     proxy: bridge_settings,
@@ -1110,18 +1109,11 @@ where
                 unreachable!("OpenVPN is not supported on Android");
             }
             MullvadEndpoint::Wireguard(endpoint) => {
-                let wg_data = self
-                    .account_manager
-                    .data()
-                    .await
-                    .map_err(|_| Error::NoKeyAvailable)?
-                    .map(|device| device.wg_data)
-                    .ok_or(Error::NoKeyAvailable)?;
                 let tunnel = wireguard::TunnelConfig {
-                    private_key: wg_data.private_key,
+                    private_key: device.wg_data.private_key,
                     addresses: vec![
-                        wg_data.addresses.ipv4_address.ip().into(),
-                        wg_data.addresses.ipv6_address.ip().into(),
+                        device.wg_data.addresses.ipv4_address.ip().into(),
+                        device.wg_data.addresses.ipv6_address.ip().into(),
                     ],
                 };
 
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index de85bcbed7..e06b88f41e 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -14,7 +14,7 @@ use mullvad_types::{
     relay_list::{Relay, RelayList, Udp2TcpEndpointData},
     CustomTunnelEndpoint,
 };
-use parking_lot::Mutex;
+use parking_lot::{Mutex, MutexGuard};
 use rand::{self, seq::SliceRandom, Rng};
 use std::{
     io,
@@ -259,15 +259,31 @@ impl RelaySelector {
     }
 
     /// Returns a random relay and relay endpoint matching the current constraints.
-    pub fn get_relay(&self, retry_attempt: u32) -> Result<SelectedRelay, Error> {
+    pub fn get_relay(
+        &self,
+        retry_attempt: u32,
+    ) -> Result<(SelectedRelay, Option<SelectedBridge>), Error> {
         let config = self.config.lock();
         match &config.relay_settings {
             RelaySettings::CustomTunnelEndpoint(custom_relay) => {
-                Ok(SelectedRelay::Custom(custom_relay.clone()))
+                Ok((SelectedRelay::Custom(custom_relay.clone()), None))
+            }
+            RelaySettings::Normal(constraints) => {
+                let relay =
+                    self.get_tunnel_endpoint(&constraints, config.bridge_state, retry_attempt)?;
+                let bridge = match relay.endpoint {
+                    MullvadEndpoint::OpenVpn(_endpoint) => {
+                        let location = relay
+                            .exit_relay
+                            .location
+                            .as_ref()
+                            .expect("Relay has no location set");
+                        self.get_bridge_for(config, location, retry_attempt)?
+                    }
+                    MullvadEndpoint::Wireguard(ref _endpoint) => None,
+                };
+                Ok((SelectedRelay::Normal(relay), bridge))
             }
-            RelaySettings::Normal(constraints) => self
-                .get_tunnel_endpoint(&constraints, config.bridge_state, retry_attempt)
-                .map(SelectedRelay::Normal),
         }
     }
 
@@ -647,12 +663,12 @@ impl RelaySelector {
         entry_endpoint.exit_peer = Some(exit_peer.clone());
     }
 
-    pub fn get_bridge(
+    fn get_bridge_for(
         &self,
+        config: MutexGuard<'_, SelectorConfig>,
         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 {
@@ -690,7 +706,7 @@ impl RelaySelector {
     }
 
     /// Returns a bridge based on the relay and bridge constraints, ignoring the bridge state.
-    pub fn get_api_bridge(&self) -> Option<ProxySettings> {
+    pub fn get_bridge_forced(&self) -> Option<ProxySettings> {
         let config = self.config.lock();
 
         let near_location = match &config.relay_settings {