Add exit location CLI option

5 files changed, 59 insertions, 14 deletions

diff --git a/mullvad-cli/src/cmds/bridge.rs b/mullvad-cli/src/cmds/bridge.rs
index 1e2a223817..0599c7aef6 100644
--- a/mullvad-cli/src/cmds/bridge.rs
+++ b/mullvad-cli/src/cmds/bridge.rs
@@ -218,7 +218,7 @@ impl Bridge {
     }
 
     async fn handle_set_bridge_location(matches: &clap::ArgMatches<'_>) -> Result<()> {
-        Self::update_bridge_settings(Some(location::get_constraint(matches)), None).await
+        Self::update_bridge_settings(Some(location::get_constraint_from_args(matches)), None).await
     }
 
     async fn handle_set_bridge_provider(matches: &clap::ArgMatches<'_>) -> Result<()> {
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 763e1d7a11..8114e656f9 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -171,6 +171,17 @@ impl Command for Relay {
                                             .default_value("any")
                                             .possible_values(&["any", "4", "6"]),
                                     )
+                                    .arg(
+                                        clap::Arg::with_name("exit location")
+                                            .help("Exit endpoint to use. This can be 'any', 'none', or \
+                                                any location that is valid with 'set location', \
+                                                such as 'se got'.")
+                                            .default_value("none")
+                                            .long("exit-location")
+                                            .multiple(true)
+                                            .min_values(1)
+                                            .max_values(3),
+                                    )
                             )
                     )
                     .subcommand(clap::SubCommand::with_name("tunnel-protocol")
@@ -415,7 +426,7 @@ impl Relay {
     }
 
     async fn set_location(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
-        let location_constraint = location::get_constraint(matches);
+        let location_constraint = location::get_constraint_from_args(matches);
         let mut found = false;
 
         if !location_constraint.country.is_empty() {
@@ -517,6 +528,8 @@ impl Relay {
     async fn set_wireguard_constraints(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let port = parse_port_constraint(matches.value_of("port").unwrap())?;
         let ip_version = parse_ip_version_constraint(matches.value_of("ip version").unwrap());
+        let exit_location =
+            parse_exit_location_constraint(matches.values_of("exit location").unwrap());
 
         self.update_constraints(RelaySettingsUpdate {
             r#type: Some(relay_settings_update::Type::Normal(
@@ -526,7 +539,7 @@ impl Relay {
                         ip_version: ip_version.option().map(|protocol| IpVersionConstraint {
                             protocol: protocol as i32,
                         }),
-                        exit_location: None,
+                        exit_location,
                     }),
                     ..Default::default()
                 },
@@ -801,3 +814,19 @@ fn parse_ip_version_constraint(raw_protocol: &str) -> Constraint<IpVersion> {
         _ => unreachable!(),
     }
 }
+
+fn parse_exit_location_constraint<'a, T: Iterator<Item = &'a str>>(
+    mut location: T,
+) -> Option<RelayLocation> {
+    let country = location.next().unwrap();
+
+    if country == "none" {
+        return None;
+    }
+
+    Some(location::get_constraint(
+        country,
+        location.next(),
+        location.next(),
+    ))
+}
diff --git a/mullvad-cli/src/location.rs b/mullvad-cli/src/location.rs
index 09f39720ba..cadeaa24fe 100644
--- a/mullvad-cli/src/location.rs
+++ b/mullvad-cli/src/location.rs
@@ -22,11 +22,22 @@ pub fn get_subcommand() -> clap::App<'static, 'static> {
         )
 }
 
-pub fn get_constraint(matches: &clap::ArgMatches<'_>) -> RelayLocation {
-    let country_original = matches.value_of("country").unwrap();
+pub fn get_constraint_from_args(matches: &clap::ArgMatches<'_>) -> RelayLocation {
+    let country = matches.value_of("country").unwrap();
+    let city = matches.value_of("city");
+    let hostname = matches.value_of("hostname");
+    get_constraint(country, city, hostname)
+}
+
+pub fn get_constraint<T: AsRef<str>>(
+    country: T,
+    city: Option<T>,
+    hostname: Option<T>,
+) -> RelayLocation {
+    let country_original = country.as_ref();
     let country = country_original.to_lowercase();
-    let city = matches.value_of("city").map(str::to_lowercase);
-    let hostname = matches.value_of("hostname").map(str::to_lowercase);
+    let city = city.map(|s| s.as_ref().to_lowercase());
+    let hostname = hostname.map(|s| s.as_ref().to_lowercase());
 
     match (country_original, city, hostname) {
         ("any", None, None) => RelayLocation::default(),
@@ -81,16 +92,16 @@ pub fn format_providers(providers: &Vec<String>) -> String {
     }
 }
 
-fn country_code_validator(code: String) -> std::result::Result<(), String> {
-    if code.len() == 2 || code == "any" {
+pub fn country_code_validator<T: AsRef<str>>(code: T) -> std::result::Result<(), String> {
+    if code.as_ref().len() == 2 || code.as_ref() == "any" {
         Ok(())
     } else {
         Err(String::from("Country codes must be two letters, or 'any'."))
     }
 }
 
-fn city_code_validator(code: String) -> std::result::Result<(), String> {
-    if code.len() == 3 {
+pub fn city_code_validator<T: AsRef<str>>(code: T) -> std::result::Result<(), String> {
+    if code.as_ref().len() == 3 {
         Ok(())
     } else {
         Err(String::from("City codes must be three letters"))
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index 07df610045..1309e8bf70 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -450,7 +450,10 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for RelaySettings {
                             .option()
                             .map(IpVersion::from)
                             .map(IpVersionConstraint::from),
-                        exit_location: None,
+                        exit_location: constraints
+                            .wireguard_constraints
+                            .exit_location
+                            .and_then(|constraint| constraint.option().map(RelayLocation::from)),
                     }),
 
                     openvpn_constraints: Some(OpenvpnConstraints {
@@ -883,7 +886,9 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                                     Constraint::Any
                                 },
                                 ip_version: Constraint::from(ip_version),
-                                exit_location: None,
+                                exit_location: constraints
+                                    .exit_location
+                                    .map(Constraint::<mullvad_types::relay_constraints::LocationConstraint>::from),
                             }
                         }),
                         openvpn_constraints: settings.openvpn_constraints.map(|constraints| {
diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 62ea95ce07..7585793a01 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -1,9 +1,9 @@
 use self::config::Config;
-use cfg_if::cfg_if;
 #[cfg(not(windows))]
 use super::tun_provider;
 use super::{tun_provider::TunProvider, TunnelEvent, TunnelMetadata};
 use crate::routing::{self, RequiredRoute};
+use cfg_if::cfg_if;
 use futures::future::abortable;
 #[cfg(target_os = "linux")]
 use lazy_static::lazy_static;