Add bridge provider constraint

8 files changed, 90 insertions, 23 deletions

diff --git a/mullvad-cli/src/cmds/bridge.rs b/mullvad-cli/src/cmds/bridge.rs
index 3118c09a39..0ee961a218 100644
--- a/mullvad-cli/src/cmds/bridge.rs
+++ b/mullvad-cli/src/cmds/bridge.rs
@@ -4,7 +4,7 @@ use clap::value_t;
 use mullvad_management_interface::types::{
     bridge_settings::{Type as BridgeSettingsType, *},
     bridge_state::State as BridgeStateType,
-    BridgeSettings, BridgeState,
+    BridgeSettings, BridgeState, RelayLocation,
 };
 use talpid_types::net::openvpn::SHADOWSOCKS_CIPHERS;
 
@@ -45,6 +45,18 @@ fn create_bridge_set_subcommand() -> clap::App<'static, 'static> {
         .setting(clap::AppSettings::SubcommandRequiredElseHelp)
         .subcommand(create_set_state_subcommand())
         .subcommand(create_set_custom_settings_subcommand())
+        .subcommand(
+            clap::SubCommand::with_name("provider")
+                .about(
+                    "Set a hosting provider to select bridge relays from. The 'list' \
+                        command shows the available relays and their providers.",
+                )
+                .arg(
+                    clap::Arg::with_name("provider")
+                        .help("The hosting provider to use, or 'any' for no preference.")
+                        .required(true),
+                ),
+        )
         .subcommand(location::get_subcommand().about(
             "Set country or city to select bridge relays from. Use the 'list' \
              command to show available alternatives.",
@@ -155,6 +167,9 @@ impl Bridge {
             ("location", Some(location_matches)) => {
                 Self::handle_set_bridge_location(location_matches).await
             }
+            ("provider", Some(provider_matches)) => {
+                Self::handle_set_bridge_provider(provider_matches).await
+            }
             ("custom", Some(custom_matches)) => {
                 Self::handle_bridge_set_custom_settings(custom_matches).await
             }
@@ -175,8 +190,9 @@ impl Bridge {
             }
             BridgeSettingsType::Normal(constraints) => {
                 println!(
-                    "Bridge constraints - {}",
-                    location::format_location(constraints.location.as_ref())
+                    "Bridge constraints - {}, {}",
+                    location::format_location(constraints.location.as_ref()),
+                    location::format_provider(constraints.provider.as_ref())
                 );
             }
         };
@@ -184,12 +200,52 @@ impl Bridge {
     }
 
     async fn handle_set_bridge_location(matches: &clap::ArgMatches<'_>) -> Result<()> {
-        let constraints = location::get_constraint(matches);
+        Self::update_bridge_settings(Some(location::get_constraint(matches)), None).await
+    }
+
+    async fn handle_set_bridge_provider(matches: &clap::ArgMatches<'_>) -> Result<()> {
+        let new_provider =
+            value_t!(matches.value_of("provider"), String).unwrap_or_else(|e| e.exit());
+        let new_provider = if new_provider == "any" {
+            "".to_string()
+        } else {
+            new_provider
+        };
+
+        Self::update_bridge_settings(None, Some(new_provider)).await
+    }
+
+    async fn update_bridge_settings(
+        location: Option<RelayLocation>,
+        provider: Option<String>,
+    ) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+        let settings = rpc.get_settings(()).await?.into_inner();
+
+        let bridge_settings = settings.bridge_settings.unwrap();
+        let constraints = match bridge_settings.r#type.unwrap() {
+            BridgeSettingsType::Normal(mut constraints) => {
+                if let Some(new_location) = location {
+                    constraints.location = Some(new_location);
+                }
+                if let Some(new_provider) = provider {
+                    constraints.provider = new_provider;
+                }
+                constraints
+            }
+            _ => {
+                let location = location.unwrap_or_default();
+                let provider = provider.unwrap_or_default();
+
+                BridgeConstraints {
+                    location: Some(location),
+                    provider,
+                }
+            }
+        };
+
         rpc.set_bridge_settings(BridgeSettings {
-            r#type: Some(BridgeSettingsType::Normal(BridgeConstraints {
-                location: Some(constraints),
-            })),
+            r#type: Some(BridgeSettingsType::Normal(constraints)),
         })
         .await?;
         Ok(())
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 4dc7dcfb47..8b3aaeffd0 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -857,6 +857,7 @@ where
                     BridgeSettings::Normal(settings) => {
                         let bridge_constraints = InternalBridgeConstraints {
                             location: settings.location.clone(),
+                            provider: settings.provider.clone(),
                             // FIXME: This is temporary while talpid-core only supports TCP proxies
                             transport_protocol: Constraint::Only(TransportProtocol::Tcp),
                         };
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index a2414bae66..dee8cb3eea 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -12,7 +12,7 @@ use mullvad_types::{
     location::GeoIpLocation,
     relay_constraints::{
         BridgeConstraints, BridgeSettings, BridgeState, Constraint, LocationConstraint,
-        OpenVpnConstraints, RelayConstraintsUpdate, RelaySettings, RelaySettingsUpdate,
+        OpenVpnConstraints, Provider, RelayConstraintsUpdate, RelaySettings, RelaySettingsUpdate,
         WireguardConstraints,
     },
     relay_list::{Relay, RelayList, RelayListCountry},
@@ -271,14 +271,16 @@ impl ManagementService for ManagementServiceImpl {
 
         let settings = match settings {
             BridgeSettingType::Normal(constraints) => {
-                let constraint = match constraints.location {
+                let location = match constraints.location {
                     None => Constraint::Any,
                     Some(location) => convert_proto_location(location),
                 };
+                let provider = match constraints.provider.as_ref() {
+                    "" => Constraint::Any,
+                    provider => Constraint::Only(String::from(provider)),
+                };
 
-                BridgeSettings::Normal(BridgeConstraints {
-                    location: constraint,
-                })
+                BridgeSettings::Normal(BridgeConstraints { location, provider })
             }
             BridgeSettingType::Local(proxy_settings) => {
                 let peer = proxy_settings
@@ -1012,10 +1014,7 @@ fn convert_relay_settings(settings: &RelaySettings) -> types::RelaySettings {
         RelaySettings::Normal(constraints) => {
             relay_settings::Endpoint::Normal(types::NormalRelaySettings {
                 location: convert_location_constraint(&constraints.location),
-                provider: match &constraints.provider {
-                    Constraint::Any => "".to_string(),
-                    Constraint::Only(ref provider) => provider.clone(),
-                },
+                provider: convert_provider_constraint(&constraints.provider),
                 tunnel_type: match constraints.tunnel_protocol {
                     Constraint::Any => None,
                     Constraint::Only(TunnelType::Wireguard) => Some(types::TunnelType::Wireguard),
@@ -1110,6 +1109,7 @@ fn convert_bridge_settings(settings: &BridgeSettings) -> types::BridgeSettings {
         BridgeSettings::Normal(constraints) => {
             BridgeSettingType::Normal(types::bridge_settings::BridgeConstraints {
                 location: convert_location_constraint(&constraints.location),
+                provider: convert_provider_constraint(&constraints.provider),
             })
         }
         BridgeSettings::Custom(proxy_settings) => match proxy_settings {
@@ -1196,6 +1196,13 @@ fn convert_location_constraint(
     })
 }
 
+fn convert_provider_constraint(provider: &Constraint<Provider>) -> String {
+    match provider.as_ref() {
+        Constraint::Any => "".to_string(),
+        Constraint::Only(ref provider) => provider.to_string(),
+    }
+}
+
 fn convert_bridge_state(state: &BridgeState) -> types::BridgeState {
     let state = match state {
         BridgeState::Auto => types::bridge_state::State::Auto,
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 075fa2468e..82a30d3f7c 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -543,6 +543,9 @@ impl RelaySelector {
         if !constraints.location.matches(relay) {
             return None;
         }
+        if !constraints.provider.matches_eq(&relay.provider) {
+            return None;
+        }
 
         let mut filtered_relay = relay.clone();
         filtered_relay
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index fae924238f..857499afd4 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -201,6 +201,7 @@ message AccountHistory {
 message BridgeSettings {
 	message BridgeConstraints {
 		RelayLocation location = 1;
+		string provider = 2;
 	}
 
 	message LocalProxySettings {
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index cbd91d3def..8f57bb58fa 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -425,10 +425,12 @@ pub enum BridgeSettings {
 
 
 /// Limits the set of bridge servers to use in `mullvad-daemon`.
-#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(default)]
 #[serde(rename_all = "snake_case")]
 pub struct BridgeConstraints {
     pub location: Constraint<LocationConstraint>,
+    pub provider: Constraint<Provider>,
 }
 
 impl fmt::Display for BridgeConstraints {
@@ -466,6 +468,7 @@ impl fmt::Display for BridgeState {
 #[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
 pub struct InternalBridgeConstraints {
     pub location: Constraint<LocationConstraint>,
+    pub provider: Constraint<Provider>,
     pub transport_protocol: Constraint<TransportProtocol>,
 }
 
diff --git a/mullvad-types/src/settings/migrations/v1.rs b/mullvad-types/src/settings/migrations/v1.rs
index 22aa4d29a0..3b4f1a2075 100644
--- a/mullvad-types/src/settings/migrations/v1.rs
+++ b/mullvad-types/src/settings/migrations/v1.rs
@@ -40,9 +40,7 @@ impl Default for Settings {
                 location: Constraint::Only(LocationConstraint::Country("se".to_owned())),
                 tunnel: Constraint::Any,
             }),
-            bridge_settings: BridgeSettings::Normal(BridgeConstraints {
-                location: Constraint::Any,
-            }),
+            bridge_settings: BridgeSettings::Normal(BridgeConstraints::default()),
             bridge_state: BridgeState::Auto,
             allow_lan: false,
             block_when_disconnected: false,
diff --git a/mullvad-types/src/settings/mod.rs b/mullvad-types/src/settings/mod.rs
index 1e5cb458ab..d32cf5fead 100644
--- a/mullvad-types/src/settings/mod.rs
+++ b/mullvad-types/src/settings/mod.rs
@@ -63,9 +63,7 @@ impl Default for Settings {
                 location: Constraint::Only(LocationConstraint::Country("se".to_owned())),
                 ..Default::default()
             }),
-            bridge_settings: BridgeSettings::Normal(BridgeConstraints {
-                location: Constraint::Any,
-            }),
+            bridge_settings: BridgeSettings::Normal(BridgeConstraints::default()),
             bridge_state: BridgeState::Auto,
             allow_lan: false,
             block_when_disconnected: false,