Merge branch 'add-provider-constraint'

3 files changed, 119 insertions, 17 deletions

diff --git a/mullvad-cli/src/cmds/bridge.rs b/mullvad-cli/src/cmds/bridge.rs
index 3118c09a39..6a68fd3637 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(())
@@ -388,7 +444,10 @@ impl Bridge {
                     city.name, city.code, city.latitude, city.longitude
                 );
                 for relay in &city.relays {
-                    println!("\t\t{} ({})", relay.hostname, relay.ipv4_addr_in);
+                    println!(
+                        "\t\t{} ({}) - hosted by {}",
+                        relay.hostname, relay.ipv4_addr_in, relay.provider
+                    );
                 }
             }
             println!();
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 7071db4d45..b9948e5fb4 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -9,7 +9,7 @@ use std::{
 use mullvad_management_interface::types::{
     connection_config::{self, OpenvpnConfig, WireguardConfig},
     relay_settings, relay_settings_update, ConnectionConfig, CustomRelaySettings,
-    NormalRelaySettingsUpdate, OpenvpnConstraints, RelayListCountry, RelayLocation,
+    NormalRelaySettingsUpdate, OpenvpnConstraints, ProviderUpdate, RelayListCountry, RelayLocation,
     RelaySettingsUpdate, TransportProtocol, TransportProtocolConstraint, TunnelType,
     TunnelTypeConstraint, TunnelTypeUpdate, WireguardConstraints,
 };
@@ -127,6 +127,16 @@ impl Command for Relay {
                             ),
                     )
                     .subcommand(
+                        clap::SubCommand::with_name("provider")
+                            .about("Set a hosting provider to select 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(
                         clap::SubCommand::with_name("tunnel")
                             .about("Set individual tunnel constraints")
                             .arg(
@@ -195,6 +205,8 @@ impl Relay {
             self.set_location(location_matches).await
         } else if let Some(relay_matches) = matches.subcommand_matches("relay") {
             self.set_relay(relay_matches).await
+        } else if let Some(provider_matches) = matches.subcommand_matches("provider") {
+            self.set_provider(provider_matches).await
         } else if let Some(tunnel_matches) = matches.subcommand_matches("tunnel") {
             self.set_tunnel(tunnel_matches).await
         } else if let Some(tunnel_matches) = matches.subcommand_matches("tunnel-protocol") {
@@ -426,6 +438,26 @@ impl Relay {
         .await
     }
 
+    async fn set_provider(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
+        let provider = value_t!(matches.value_of("provider"), String).unwrap_or_else(|e| e.exit());
+
+        self.update_constraints(RelaySettingsUpdate {
+            r#type: Some(relay_settings_update::Type::Normal(
+                NormalRelaySettingsUpdate {
+                    provider: Some(ProviderUpdate {
+                        provider: if provider == "any" {
+                            "".to_string()
+                        } else {
+                            provider
+                        },
+                    }),
+                    ..Default::default()
+                },
+            )),
+        })
+        .await
+    }
+
     async fn set_tunnel(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let vpn_protocol = matches.value_of("vpn protocol").unwrap();
         let port = parse_port_constraint(matches.value_of("port").unwrap())?;
@@ -507,27 +539,30 @@ impl Relay {
             relay_settings::Endpoint::Normal(settings) => match settings.tunnel_type {
                 None => {
                     println!(
-                        "Any tunnel protocol with OpenVPN over {} and WireGuard over {} in {}",
+                        "Any tunnel protocol with OpenVPN over {} and WireGuard over {} in {} using {}",
                         Self::format_openvpn_constraints(settings.openvpn_constraints.as_ref()),
                         Self::format_wireguard_constraints(settings.wireguard_constraints.as_ref()),
-                        location::format_location(settings.location.as_ref())
+                        location::format_location(settings.location.as_ref()),
+                        location::format_provider(settings.provider.as_ref())
                     );
                 }
                 Some(constraint) => match TunnelType::from_i32(constraint.tunnel_type).unwrap() {
                     TunnelType::Wireguard => {
                         println!(
-                            "WireGuard over {} in {}",
+                            "WireGuard over {} in {} using {}",
                             Self::format_wireguard_constraints(
                                 settings.wireguard_constraints.as_ref()
                             ),
-                            location::format_location(settings.location.as_ref())
+                            location::format_location(settings.location.as_ref()),
+                            location::format_provider(settings.provider.as_ref())
                         );
                     }
                     TunnelType::Openvpn => {
                         println!(
-                            "OpenVPN over {} in {}",
+                            "OpenVPN over {} in {} using {}",
                             Self::format_openvpn_constraints(settings.openvpn_constraints.as_ref()),
-                            location::format_location(settings.location.as_ref())
+                            location::format_location(settings.location.as_ref()),
+                            location::format_provider(settings.provider.as_ref())
                         );
                     }
                 },
@@ -586,8 +621,8 @@ impl Relay {
                         _ => unreachable!("Bug in relay filtering earlier on"),
                     };
                     println!(
-                        "\t\t{} ({}) - {}",
-                        relay.hostname, relay.ipv4_addr_in, support_msg
+                        "\t\t{} ({}) - {}, hosted by {}",
+                        relay.hostname, relay.ipv4_addr_in, support_msg, relay.provider
                     );
                 }
             }
diff --git a/mullvad-cli/src/location.rs b/mullvad-cli/src/location.rs
index 21e6a1c35c..bc4c9bc630 100644
--- a/mullvad-cli/src/location.rs
+++ b/mullvad-cli/src/location.rs
@@ -73,6 +73,14 @@ pub fn format_location(location: Option<&RelayLocation>) -> String {
     "any location".to_string()
 }
 
+pub fn format_provider(provider: &str) -> String {
+    if !provider.is_empty() {
+        format!("provider {}", provider)
+    } else {
+        "any provider".to_string()
+    }
+}
+
 fn country_code_validator(code: String) -> std::result::Result<(), String> {
     if code.len() == 2 || code == "any" {
         Ok(())