Merge branch 'cli-set-location-hostname'

4 files changed, 184 insertions, 65 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8dd6eff2cf..7bcce16a46 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,8 +23,12 @@ Line wrap the file at 100 chars.                                              Th
 
 
 ## [Unreleased]
+### Added
+- Add CLI command to set the location constraint via `mullvad relay set relay HOSTNAME`.
+
 ### Changed
 - Use gRPC for communication between frontends and the backend instead of JSON-RPC.
+- Show a warning in the CLI if the provided location constraints don't match any known relay.
 
 ### Fixed
 #### Android
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index b23f8638d3..7071db4d45 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -9,9 +9,9 @@ use std::{
 use mullvad_management_interface::types::{
     connection_config::{self, OpenvpnConfig, WireguardConfig},
     relay_settings, relay_settings_update, ConnectionConfig, CustomRelaySettings,
-    NormalRelaySettingsUpdate, OpenvpnConstraints, RelaySettingsUpdate, TransportProtocol,
-    TransportProtocolConstraint, TunnelType, TunnelTypeConstraint, TunnelTypeUpdate,
-    WireguardConstraints,
+    NormalRelaySettingsUpdate, OpenvpnConstraints, RelayListCountry, RelayLocation,
+    RelaySettingsUpdate, TransportProtocol, TransportProtocolConstraint, TunnelType,
+    TunnelTypeConstraint, TunnelTypeUpdate, WireguardConstraints,
 };
 use mullvad_types::relay_constraints::Constraint;
 use talpid_types::net::all_of_the_internet;
@@ -117,6 +117,16 @@ impl Command for Relay {
                                    command to show available alternatives.")
                     )
                     .subcommand(
+                        clap::SubCommand::with_name("relay")
+                            .about("Set the exact relay to use via its hostname. Shortcut for \
+                                'location <country> <city> <hostname>'.")
+                            .arg(
+                                clap::Arg::with_name("hostname")
+                                    .help("The hostname")
+                                    .required(true),
+                            ),
+                    )
+                    .subcommand(
                         clap::SubCommand::with_name("tunnel")
                             .about("Set individual tunnel constraints")
                             .arg(
@@ -183,6 +193,8 @@ impl Relay {
             self.set_custom(custom_matches).await
         } else if let Some(location_matches) = matches.subcommand_matches("location") {
             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(tunnel_matches) = matches.subcommand_matches("tunnel") {
             self.set_tunnel(tunnel_matches).await
         } else if let Some(tunnel_matches) = matches.subcommand_matches("tunnel-protocol") {
@@ -310,8 +322,98 @@ impl Relay {
         key
     }
 
+    async fn set_relay(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
+        let hostname = matches.value_of("hostname").unwrap();
+        let countries = Self::get_filtered_relays().await?;
+
+        let find_relay = || {
+            for country in &countries {
+                for city in &country.cities {
+                    for relay in &city.relays {
+                        if relay.hostname == hostname {
+                            return Some((country, city, relay));
+                        }
+                    }
+                }
+            }
+            None
+        };
+
+        if let Some(location) = find_relay() {
+            println!(
+                "Setting location constraint to {} in {}, {}",
+                location.2.hostname, location.1.name, location.0.name
+            );
+
+            let location_constraint = RelayLocation {
+                country: location.0.code.clone(),
+                city: location.1.code.clone(),
+                hostname: location.2.hostname.clone(),
+            };
+
+            self.update_constraints(RelaySettingsUpdate {
+                r#type: Some(relay_settings_update::Type::Normal(
+                    NormalRelaySettingsUpdate {
+                        location: Some(location_constraint),
+                        ..Default::default()
+                    },
+                )),
+            })
+            .await
+        } else {
+            clap::Error::with_description(
+                "No matching server found",
+                clap::ErrorKind::ValueValidation,
+            )
+            .exit()
+        }
+    }
+
     async fn set_location(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let location_constraint = location::get_constraint(matches);
+        let mut found = false;
+
+        if !location_constraint.country.is_empty() {
+            // TODO: `mullvad_types::relay_constraints::LocationConstraint::matches(&relay)`
+            //       could be used to guarantee consistency with the daemon.
+            let countries = Self::get_filtered_relays().await?;
+            for country in &countries {
+                if country.code != location_constraint.country {
+                    continue;
+                }
+
+                if location_constraint.city.is_empty() {
+                    found = true;
+                    break;
+                }
+
+                for city in &country.cities {
+                    if city.code != location_constraint.city {
+                        continue;
+                    }
+
+                    if location_constraint.hostname.is_empty() {
+                        found = true;
+                        break;
+                    }
+
+                    for relay in &city.relays {
+                        if relay.hostname != location_constraint.hostname {
+                            continue;
+                        }
+                        found = true;
+                        break;
+                    }
+
+                    break;
+                }
+                break;
+            }
+
+            if !found {
+                eprintln!("Warning: No matching relay was found.");
+            }
+        }
 
         self.update_constraints(RelaySettingsUpdate {
             r#type: Some(relay_settings_update::Type::Normal(
@@ -459,34 +561,7 @@ impl Relay {
     }
 
     async fn list(&self) -> Result<()> {
-        let mut rpc = new_rpc_client().await?;
-        let mut locations = rpc.get_relay_locations(()).await?.into_inner();
-
-        let mut countries = Vec::new();
-
-        while let Some(mut country) = locations.message().await? {
-            country.cities = country
-                .cities
-                .into_iter()
-                .filter_map(|mut city| {
-                    city.relays.retain(|relay| {
-                        relay.active
-                            && relay.tunnels.is_some()
-                            && !(relay.tunnels.as_ref().unwrap().openvpn.is_empty()
-                                && relay.tunnels.as_ref().unwrap().wireguard.is_empty())
-                    });
-                    if !city.relays.is_empty() {
-                        Some(city)
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-            if !country.cities.is_empty() {
-                countries.push(country);
-            }
-        }
-
+        let mut countries = Self::get_filtered_relays().await?;
         countries.sort_by(|c1, c2| natord::compare_ignore_case(&c1.name, &c2.name));
         for mut country in countries {
             country
@@ -567,6 +642,38 @@ impl Relay {
             "any port".to_string()
         }
     }
+
+    async fn get_filtered_relays() -> Result<Vec<RelayListCountry>> {
+        let mut rpc = new_rpc_client().await?;
+        let mut locations = rpc.get_relay_locations(()).await?.into_inner();
+
+        let mut countries = Vec::new();
+
+        while let Some(mut country) = locations.message().await? {
+            country.cities = country
+                .cities
+                .into_iter()
+                .filter_map(|mut city| {
+                    city.relays.retain(|relay| {
+                        relay.active
+                            && relay.tunnels.is_some()
+                            && !(relay.tunnels.as_ref().unwrap().openvpn.is_empty()
+                                && relay.tunnels.as_ref().unwrap().wireguard.is_empty())
+                    });
+                    if !city.relays.is_empty() {
+                        Some(city)
+                    } else {
+                        None
+                    }
+                })
+                .collect();
+            if !country.cities.is_empty() {
+                countries.push(country);
+            }
+        }
+
+        Ok(countries)
+    }
 }
 
 
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 4f10f6e3fd..8640b9069e 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -388,7 +388,7 @@ impl RelaySelector {
                 self.parsed_relays.lock().relays().iter().any(|relay| {
                     relay.active
                         && !relay.tunnels.wireguard.is_empty()
-                        && Self::relay_matches_location(relay, &location_constraint)
+                        && location_constraint.matches(relay)
                 });
             // If location does not support WireGuard, defer to preferred OpenVPN tunnel
             // constraints
@@ -472,7 +472,7 @@ impl RelaySelector {
     /// Takes a `Relay` and a corresponding `RelayConstraints` and returns a new `Relay` if the
     /// given relay matches the constraints.
     fn matching_relay(relay: &Relay, constraints: &RelayConstraints) -> Option<Relay> {
-        if !Self::relay_matches_location(relay, &constraints.location) {
+        if !constraints.location.matches(relay) {
             return None;
         }
 
@@ -533,36 +533,11 @@ impl RelaySelector {
         }
     }
 
-    fn relay_matches_location(relay: &Relay, location: &Constraint<LocationConstraint>) -> bool {
-        match location {
-            Constraint::Any => true,
-            Constraint::Only(LocationConstraint::Country(ref country)) => {
-                relay
-                    .location
-                    .as_ref()
-                    .map_or(false, |loc| loc.country_code == *country)
-                    && relay.include_in_country
-            }
-            Constraint::Only(LocationConstraint::City(ref country, ref city)) => {
-                relay.location.as_ref().map_or(false, |loc| {
-                    loc.country_code == *country && loc.city_code == *city
-                })
-            }
-            Constraint::Only(LocationConstraint::Hostname(ref country, ref city, ref hostname)) => {
-                relay.location.as_ref().map_or(false, |loc| {
-                    loc.country_code == *country
-                        && loc.city_code == *city
-                        && relay.hostname == *hostname
-                })
-            }
-        }
-    }
-
     fn matching_bridge_relay(
         relay: &Relay,
         constraints: &InternalBridgeConstraints,
     ) -> Option<Relay> {
-        if !Self::relay_matches_location(relay, &constraints.location) {
+        if !constraints.location.matches(relay) {
             return None;
         }
 
@@ -570,7 +545,7 @@ impl RelaySelector {
         filtered_relay
             .bridges
             .shadowsocks
-            .retain(|bridge| constraints.transport_protocol.matches(&bridge.protocol));
+            .retain(|bridge| constraints.transport_protocol.matches_eq(&bridge.protocol));
         if filtered_relay.bridges.shadowsocks.is_empty() {
             return None;
         }
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index c5d58f621f..857b8747fb 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -3,7 +3,7 @@
 
 use crate::{
     location::{CityCode, CountryCode, Hostname},
-    relay_list::{OpenVpnEndpointData, WireguardEndpointData},
+    relay_list::{OpenVpnEndpointData, Relay, WireguardEndpointData},
     CustomTunnelEndpoint,
 };
 #[cfg(target_os = "android")]
@@ -79,6 +79,13 @@ impl<T: fmt::Debug + Clone + Eq + PartialEq> Constraint<T> {
             Constraint::Only(value) => Some(value),
         }
     }
+
+    pub fn matches_eq(&self, other: &T) -> bool {
+        match self {
+            Constraint::Any => true,
+            Constraint::Only(ref value) => value == other,
+        }
+    }
 }
 
 impl<T: fmt::Debug + Clone + Eq + PartialEq> Default for Constraint<T> {
@@ -89,11 +96,11 @@ impl<T: fmt::Debug + Clone + Eq + PartialEq> Default for Constraint<T> {
 
 impl<T: Copy + fmt::Debug + Clone + Eq + PartialEq> Copy for Constraint<T> {}
 
-impl<T: fmt::Debug + Clone + Eq + PartialEq> Match<T> for Constraint<T> {
-    fn matches(&self, other: &T) -> bool {
+impl<T: fmt::Debug + Clone + Eq + Match<U>, U> Match<U> for Constraint<T> {
+    fn matches(&self, other: &U) -> bool {
         match *self {
             Constraint::Any => true,
-            Constraint::Only(ref value) => value == other,
+            Constraint::Only(ref value) => value.matches(other),
         }
     }
 }
@@ -258,6 +265,32 @@ pub enum LocationConstraint {
     Hostname(CountryCode, CityCode, Hostname),
 }
 
+impl Match<Relay> for LocationConstraint {
+    fn matches(&self, relay: &Relay) -> bool {
+        match self {
+            LocationConstraint::Country(ref country) => {
+                relay
+                    .location
+                    .as_ref()
+                    .map_or(false, |loc| loc.country_code == *country)
+                    && relay.include_in_country
+            }
+            LocationConstraint::City(ref country, ref city) => {
+                relay.location.as_ref().map_or(false, |loc| {
+                    loc.country_code == *country && loc.city_code == *city
+                })
+            }
+            LocationConstraint::Hostname(ref country, ref city, ref hostname) => {
+                relay.location.as_ref().map_or(false, |loc| {
+                    loc.country_code == *country
+                        && loc.city_code == *city
+                        && relay.hostname == *hostname
+                })
+            }
+        }
+    }
+}
+
 impl fmt::Display for LocationConstraint {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         match self {
@@ -335,7 +368,7 @@ impl fmt::Display for OpenVpnConstraints {
 
 impl Match<OpenVpnEndpointData> for OpenVpnConstraints {
     fn matches(&self, endpoint: &OpenVpnEndpointData) -> bool {
-        self.port.matches(&endpoint.port) && self.protocol.matches(&endpoint.protocol)
+        self.port.matches_eq(&endpoint.port) && self.protocol.matches_eq(&endpoint.protocol)
     }
 }