Add exit location CLI option

2 files changed, 32 insertions, 3 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(),
+    ))
+}