Pick bridges from the relay list

2 files changed, 176 insertions, 46 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 327252ada8..50b8dae23f 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -36,8 +36,8 @@ use mullvad_types::{
     endpoint::MullvadEndpoint,
     location::GeoIpLocation,
     relay_constraints::{
-        Constraint, OpenVpnConstraints, RelayConstraintsUpdate, RelaySettings, RelaySettingsUpdate,
-        TunnelConstraints,
+        BridgeSettings, BridgeState, Constraint, InternalBridgeConstraints, OpenVpnConstraints,
+        RelayConstraintsUpdate, RelaySettings, RelaySettingsUpdate, TunnelConstraints,
     },
     relay_list::{Relay, RelayList},
     settings::{self, Settings},
@@ -90,6 +90,9 @@ pub enum Error {
     #[error(display = "No wireguard private key available")]
     NoKeyAvailable,
 
+    #[error(display = "No bridge available")]
+    NoBridgeAvailable,
+
     #[error(display = "Account history problems")]
     AccountHistory(#[error(cause)] account_history::Error),
 
@@ -457,9 +460,16 @@ where
                         )
                     })
                     .and_then(|(relay, endpoint)| {
+                        let result = self
+                            .create_tunnel_parameters(
+                                &relay,
+                                endpoint,
+                                account_token,
+                                retry_attempt,
+                            )
+                            .map_err(|e| e.display_chain());
                         self.last_generated_relay = Some(relay);
-                        self.create_tunnel_parameters(endpoint, account_token)
-                            .map_err(|e| e.display_chain())
+                        result
                     }),
             }
             .and_then(|tunnel_params| {
@@ -476,17 +486,62 @@ where
 
     fn create_tunnel_parameters(
         &mut self,
+        relay: &Relay,
         endpoint: MullvadEndpoint,
         account_token: String,
+        retry_attempt: u32,
     ) -> Result<TunnelParameters> {
-        let tunnel_options = self.settings.get_tunnel_options().clone();
+        let mut tunnel_options = self.settings.get_tunnel_options().clone();
+        let location = relay.location.as_ref().expect("Relay has no location set");
         match endpoint {
-            MullvadEndpoint::OpenVpn(endpoint) => Ok(openvpn::TunnelParameters {
-                config: openvpn::ConnectionConfig::new(endpoint, account_token, "-".to_string()),
-                options: tunnel_options.openvpn,
-                generic_options: tunnel_options.generic,
+            MullvadEndpoint::OpenVpn(endpoint) => {
+                let proxy_settings = match self.settings.get_bridge_settings() {
+                    BridgeSettings::Normal(settings) => {
+                        let bridge_constraints = InternalBridgeConstraints {
+                            location: settings.location.clone(),
+                            transport_protocol: Constraint::Only(endpoint.protocol),
+                        };
+                        match self.settings.get_bridge_state() {
+                            BridgeState::On => Some(
+                                self.relay_selector
+                                    .get_proxy_settings(&bridge_constraints, location)
+                                    .ok_or(Error::NoBridgeAvailable)?,
+                            ),
+                            BridgeState::Auto => self.relay_selector.get_auto_proxy_settings(
+                                &bridge_constraints,
+                                location,
+                                retry_attempt,
+                            ),
+                            BridgeState::Off => None,
+                        }
+                    }
+                    BridgeSettings::Custom(proxy_settings) => {
+                        match self.settings.get_bridge_state() {
+                            BridgeState::On => Some(proxy_settings.clone()),
+                            BridgeState::Auto => {
+                                if self.relay_selector.should_use_bridge(retry_attempt) {
+                                    Some(proxy_settings.clone())
+                                } else {
+                                    None
+                                }
+                            }
+                            BridgeState::Off => None,
+                        }
+                    }
+                };
+                tunnel_options.openvpn.proxy = proxy_settings;
+
+                Ok(openvpn::TunnelParameters {
+                    config: openvpn::ConnectionConfig::new(
+                        endpoint,
+                        account_token,
+                        "-".to_string(),
+                    ),
+                    options: tunnel_options.openvpn,
+                    generic_options: tunnel_options.generic,
+                }
+                .into())
             }
-            .into()),
             MullvadEndpoint::Wireguard {
                 peer,
                 ipv4_gateway,
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index eac68588de..f77dca1bec 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -5,8 +5,8 @@ use mullvad_types::{
     endpoint::MullvadEndpoint,
     location::Location,
     relay_constraints::{
-        Constraint, LocationConstraint, Match, OpenVpnConstraints, RelayConstraints,
-        TunnelConstraints, WireguardConstraints,
+        Constraint, InternalBridgeConstraints, LocationConstraint, Match, OpenVpnConstraints,
+        RelayConstraints, TunnelConstraints, WireguardConstraints,
     },
     relay_list::{Relay, RelayList, RelayTunnels, WireguardEndpointData},
 };
@@ -20,7 +20,7 @@ use std::{
     time::{self, Duration, SystemTime},
 };
 use talpid_types::{
-    net::{all_of_the_internet, wireguard, TransportProtocol},
+    net::{all_of_the_internet, openvpn::ProxySettings, wireguard, TransportProtocol},
     ErrorExt,
 };
 
@@ -90,12 +90,6 @@ impl ParsedRelays {
                 let city_code = city.code.clone();
                 let latitude = city.latitude;
                 let longitude = city.longitude;
-                city.relays = city
-                    .relays
-                    .iter()
-                    .filter(|relay| !relay.tunnels.is_empty())
-                    .cloned()
-                    .collect();
                 for relay in &mut city.relays {
                     let mut relay_with_location = relay.clone();
                     relay_with_location.location = Some(Location {
@@ -214,24 +208,25 @@ impl RelaySelector {
     /// preferences applied.
     pub fn get_tunnel_endpoint(
         &mut self,
-        constraints: &RelayConstraints,
+        relay_constraints: &RelayConstraints,
         retry_attempt: u32,
     ) -> Result<(Relay, MullvadEndpoint), Error> {
-        let preferred_constraints = Self::preferred_constraints(constraints, retry_attempt);
+        let preferred_constraints = Self::preferred_constraints(relay_constraints, retry_attempt);
         if let Some((relay, endpoint)) = self.get_tunnel_endpoint_internal(&preferred_constraints) {
             debug!(
                 "Relay matched on highest preference for retry attempt {}",
                 retry_attempt
             );
             Ok((relay, endpoint))
-        } else if let Some((relay, endpoint)) = self.get_tunnel_endpoint_internal(constraints) {
+        } else if let Some((relay, endpoint)) = self.get_tunnel_endpoint_internal(relay_constraints)
+        {
             debug!(
                 "Relay matched on second preference for retry attempt {}",
                 retry_attempt
             );
             Ok((relay, endpoint))
         } else {
-            warn!("No relays matching {}", constraints);
+            warn!("No relays matching {}", relay_constraints);
             Err(Error::NoRelay)
         }
     }
@@ -280,6 +275,53 @@ impl RelaySelector {
         }
     }
 
+    pub fn get_auto_proxy_settings(
+        &mut self,
+        bridge_constraints: &InternalBridgeConstraints,
+        location: &Location,
+        retry_attempt: u32,
+    ) -> Option<ProxySettings> {
+        if !self.should_use_bridge(retry_attempt) {
+            return None;
+        }
+        self.get_proxy_settings(bridge_constraints, location)
+    }
+
+    pub 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
+            // The test is to see whether the current _couple of connections_ is even or not.
+            // | retry_attempt                | 4 | 5 | 6 | 7 | 8 | 9 |
+            // | (retry_attempt % 4) < 2      | t | t | f | f | t | t |
+            (retry_attempt % 4) < 2
+    }
+
+    pub fn get_proxy_settings(
+        &mut self,
+        constraints: &InternalBridgeConstraints,
+        location: &Location,
+    ) -> Option<ProxySettings> {
+        let mut matching_relays: Vec<Relay> = self
+            .lock_parsed_relays()
+            .relays()
+            .iter()
+            .filter_map(|relay| Self::matching_bridge_relay(relay, constraints))
+            .collect();
+
+        if matching_relays.len() == 0 {
+            return None;
+        }
+
+        matching_relays.sort_by_cached_key(|relay| {
+            (relay.location.as_ref().unwrap().distance_from(&location) * 1000.0) as i64
+        });
+        return matching_relays
+            .get(0)
+            .and_then(|relay| self.pick_random_bridge(&relay));
+    }
+
+
     /// Returns a random relay endpoint if any is matching the given constraints.
     fn get_tunnel_endpoint_internal(
         &mut self,
@@ -306,31 +348,10 @@ impl RelaySelector {
     /// Takes a `Relay` and a corresponding `RelayConstraints` and returns a new `Relay` if the
     /// given relay matches the constraints.
     fn matching_relay(relay: &Relay, constraints: &RelayConstraints) -> Option<Relay> {
-        let matches_location = match constraints.location {
-            Constraint::Any => true,
-            Constraint::Only(LocationConstraint::Country(ref country)) => {
-                relay
-                    .location
-                    .as_ref()
-                    .map_or(false, |loc| loc.country_code == *country)
-                    && relay.include_in_country
-            }
-            Constraint::Only(LocationConstraint::City(ref country, ref city)) => {
-                relay.location.as_ref().map_or(false, |loc| {
-                    loc.country_code == *country && loc.city_code == *city
-                })
-            }
-            Constraint::Only(LocationConstraint::Hostname(ref country, ref city, ref hostname)) => {
-                relay.location.as_ref().map_or(false, |loc| {
-                    loc.country_code == *country
-                        && loc.city_code == *city
-                        && relay.hostname == *hostname
-                })
-            }
-        };
-        if !matches_location {
+        if !Self::relay_matches_location(relay, &constraints.location) {
             return None;
         }
+
         let relay = match constraints.tunnel {
             Constraint::Any => relay.clone(),
             Constraint::Only(ref tunnel_constraints) => {
@@ -356,6 +377,51 @@ impl RelaySelector {
         }
     }
 
+    fn relay_matches_location(relay: &Relay, location: &Constraint<LocationConstraint>) -> bool {
+        match location {
+            Constraint::Any => true,
+            Constraint::Only(LocationConstraint::Country(ref country)) => {
+                relay
+                    .location
+                    .as_ref()
+                    .map_or(false, |loc| loc.country_code == *country)
+                    && relay.include_in_country
+            }
+            Constraint::Only(LocationConstraint::City(ref country, ref city)) => {
+                relay.location.as_ref().map_or(false, |loc| {
+                    loc.country_code == *country && loc.city_code == *city
+                })
+            }
+            Constraint::Only(LocationConstraint::Hostname(ref country, ref city, ref hostname)) => {
+                relay.location.as_ref().map_or(false, |loc| {
+                    loc.country_code == *country
+                        && loc.city_code == *city
+                        && relay.hostname == *hostname
+                })
+            }
+        }
+    }
+
+    fn matching_bridge_relay(
+        relay: &Relay,
+        constraints: &InternalBridgeConstraints,
+    ) -> Option<Relay> {
+        if !Self::relay_matches_location(relay, &constraints.location) {
+            return None;
+        }
+
+        let mut filtered_relay = relay.clone();
+        filtered_relay
+            .bridges
+            .shadowsocks
+            .retain(|bridge| constraints.transport_protocol.matches(&bridge.protocol));
+        if filtered_relay.bridges.shadowsocks.len() == 0 {
+            return None;
+        }
+
+        Some(filtered_relay)
+    }
+
     /// Takes a `RelayTunnels` object which in turn is a collection of tunnel configurations for
     /// a given relay. Then returns a new `RelayTunnels` instance with only the entries that
     /// matches the given `TunnelConstraints`.
@@ -405,6 +471,15 @@ impl RelaySelector {
         }
     }
 
+    /// Picks a random bridge from a relay.
+    fn pick_random_bridge(&mut self, relay: &Relay) -> Option<ProxySettings> {
+        relay
+            .bridges
+            .shadowsocks
+            .choose(&mut self.rng)
+            .map(|data| data.clone().to_proxy_settings(relay.ipv4_addr_in.into()))
+    }
+
     fn get_random_tunnel(
         &mut self,
         relay: &Relay,