Add exit location CLI option

3 files changed, 51 insertions, 11 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"))