Update relay constraints to contain location etc

3 files changed, 278 insertions, 103 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 817ae17b3a..4f81af9192 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -2,11 +2,15 @@ use {Command, Result, ResultExt};
 use clap;
 use std::str::FromStr;
 
-use mullvad_types::relay_constraints::{Constraint, OpenVpnConstraintsUpdate, RelayConstraints,
-                                       RelayConstraintsUpdate, TunnelConstraintsUpdate};
+use mullvad_types::CustomTunnelEndpoint;
+use mullvad_types::relay_constraints::{Constraint, LocationConstraint, OpenVpnConstraints,
+                                       RelayConstraintsUpdate, RelaySettings, RelaySettingsUpdate,
+                                       TunnelConstraints};
+use mullvad_types::relay_list::RelayList;
 
 use rpc;
-use talpid_types::net::TransportProtocol;
+use talpid_types::net::{OpenVpnParameters, TransportProtocol, TunnelParameters,
+                        WireguardParameters};
 
 pub struct Relay;
 
@@ -23,26 +27,77 @@ impl Command for Relay {
                 clap::SubCommand::with_name("set")
                     .setting(clap::AppSettings::SubcommandRequired)
                     .subcommand(
-                        clap::SubCommand::with_name("host")
-                            .about("Set host")
-                            .arg(clap::Arg::with_name("host").required(true)),
+                        clap::SubCommand::with_name("custom")
+                            .about("Set a custom VPN relay")
+                            .arg(
+                                clap::Arg::with_name("tunnel")
+                                    .required(true)
+                                    .index(1)
+                                    .possible_values(&["openvpn", "wireguard"]),
+                            )
+                            .arg(
+                                clap::Arg::with_name("host")
+                                    .help("Hostname or IP")
+                                    .required(true)
+                                    .index(2),
+                            )
+                            .arg(
+                                clap::Arg::with_name("port")
+                                    .help("Remote network port")
+                                    .required(true)
+                                    .index(3),
+                            )
+                            .arg(
+                                clap::Arg::with_name("protocol")
+                                    .help("Transport protocol. For Wireguard this is ignored.")
+                                    .index(4)
+                                    .default_value("udp")
+                                    .possible_values(&["udp", "tcp"]),
+                            ),
                     )
                     .subcommand(
-                        clap::SubCommand::with_name("port")
-                            .about("Set port")
-                            .arg(clap::Arg::with_name("port").required(true)),
+                        clap::SubCommand::with_name("location")
+                            .about(
+                                "Set country or city to select relays from. Use the 'list' \
+                                 command to show available alternatives.",
+                            )
+                            .arg(
+                                clap::Arg::with_name("country")
+                                    .help(
+                                        "The two letter country code, or 'any' for no preference.",
+                                    )
+                                    .required(true)
+                                    .index(1)
+                                    .validator(country_code_validator),
+                            )
+                            .arg(
+                                clap::Arg::with_name("city")
+                                    .help("The three letter city code")
+                                    .index(2)
+                                    .validator(city_code_validator),
+                            ),
                     )
                     .subcommand(
-                        clap::SubCommand::with_name("protocol")
-                            .about("Set protocol")
+                        clap::SubCommand::with_name("tunnel")
+                            .about("Set tunnel constraints")
+                            .arg(clap::Arg::with_name("port").required(true).index(1))
                             .arg(
                                 clap::Arg::with_name("protocol")
                                     .required(true)
+                                    .index(2)
                                     .possible_values(&["any", "udp", "tcp"]),
                             ),
                     ),
             )
             .subcommand(clap::SubCommand::with_name("get"))
+            .subcommand(
+                clap::SubCommand::with_name("list")
+                    .setting(clap::AppSettings::SubcommandRequired)
+                    .subcommand(
+                        clap::SubCommand::with_name("locations")
+                            .about("List available countries and cities"),
+                    ),
+            )
     }
 
     fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
@@ -50,6 +105,8 @@ impl Command for Relay {
             self.set(set_matches)
         } else if let Some(_) = matches.subcommand_matches("get") {
             self.get()
+        } else if let Some(list_matches) = matches.subcommand_matches("list") {
+            self.list(list_matches)
         } else {
             unreachable!("No relay command given");
         }
@@ -57,57 +114,98 @@ impl Command for Relay {
 }
 
 impl Relay {
-    fn update_constraints(&self, constraints_update: RelayConstraintsUpdate) -> Result<()> {
-        rpc::call("update_relay_constraints", &[constraints_update])
+    fn update_constraints(&self, update: RelaySettingsUpdate) -> Result<()> {
+        rpc::call("update_relay_settings", &[update])
             .map(|_: Option<()>| println!("Relay constraints updated"))
     }
 
     fn set(&self, matches: &clap::ArgMatches) -> Result<()> {
-        if let Some(host_matches) = matches.subcommand_matches("host") {
-            let host = value_t_or_exit!(host_matches.value_of("host"), String);
-
-            self.update_constraints(RelayConstraintsUpdate {
-                host: Some(Constraint::Only(host)),
-                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
-                    port: None,
-                    protocol: None,
-                }),
-            })
-        } else if let Some(port_matches) = matches.subcommand_matches("port") {
-            let port = parse_port(port_matches.value_of("port").unwrap())?;
-
-            self.update_constraints(RelayConstraintsUpdate {
-                host: None,
-                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
-                    port: Some(port),
-                    protocol: None,
-                }),
-            })
-        } else if let Some(protocol_matches) = matches.subcommand_matches("protocol") {
-            let protocol = parse_protocol(protocol_matches.value_of("protocol").unwrap());
-
-            self.update_constraints(RelayConstraintsUpdate {
-                host: None,
-                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
-                    port: None,
-                    protocol: Some(protocol),
-                }),
-            })
+        if let Some(custom_matches) = matches.subcommand_matches("custom") {
+            self.set_custom(custom_matches)
+        } else if let Some(location_matches) = matches.subcommand_matches("location") {
+            self.set_location(location_matches)
+        } else if let Some(tunnel_matches) = matches.subcommand_matches("tunnel") {
+            self.set_tunnel(tunnel_matches)
         } else {
             unreachable!("No set relay command given");
         }
     }
 
+    fn set_custom(&self, matches: &clap::ArgMatches) -> Result<()> {
+        let host = value_t!(matches.value_of("host"), String).unwrap_or_else(|e| e.exit());
+        let port = value_t!(matches.value_of("port"), u16).unwrap_or_else(|e| e.exit());
+        let tunnel = match matches.value_of("tunnel").unwrap() {
+            "openvpn" => TunnelParameters::OpenVpn(OpenVpnParameters {
+                port,
+                protocol: value_t!(matches.value_of("protocol"), TransportProtocol).unwrap(),
+            }),
+            "wireguard" => TunnelParameters::Wireguard(WireguardParameters { port }),
+            _ => unreachable!("Invalid tunnel protocol"),
+        };
+        self.update_constraints(RelaySettingsUpdate::CustomTunnelEndpoint(
+            CustomTunnelEndpoint { host, tunnel },
+        ))
+    }
+
+    fn set_location(&self, matches: &clap::ArgMatches) -> Result<()> {
+        let country = matches.value_of("country").unwrap();
+        let city = matches.value_of("city");
+
+        let location_constraint = match (country, city) {
+            ("any", None) => Constraint::Any,
+            ("any", _) => clap::Error::with_description(
+                "City can't be given when selecting 'any' country",
+                clap::ErrorKind::InvalidValue,
+            ).exit(),
+            (country, None) => Constraint::Only(LocationConstraint::Country(country.to_owned())),
+            (country, Some(city)) => Constraint::Only(LocationConstraint::City(
+                country.to_owned(),
+                city.to_owned(),
+            )),
+        };
+
+        self.update_constraints(RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
+            location: Some(location_constraint),
+            tunnel: None,
+        }))
+    }
+
+    fn set_tunnel(&self, matches: &clap::ArgMatches) -> Result<()> {
+        let port = parse_port_constraint(matches.value_of("port").unwrap())?;
+        let protocol = parse_protocol_constraint(matches.value_of("protocol").unwrap());
+
+        self.update_constraints(RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
+            location: None,
+            tunnel: Some(Constraint::Only(TunnelConstraints::OpenVpn(
+                OpenVpnConstraints { port, protocol },
+            ))),
+        }))
+    }
+
     fn get(&self) -> Result<()> {
-        let constraints: RelayConstraints = rpc::call("get_relay_constraints", &[] as &[u8; 0])?;
-        println!("Current constraints: {:?}", constraints);
+        let constraints: RelaySettings = rpc::call("get_relay_settings", &[] as &[u8; 0])?;
+        println!("Current constraints: {:#?}", constraints);
 
         Ok(())
     }
+
+    fn list(&self, _matches: &clap::ArgMatches) -> Result<()> {
+        let mut locations: RelayList = rpc::call("get_relay_locations", &[] as &[u8; 0])?;
+        locations.countries.sort_by(|c1, c2| c1.name.cmp(&c2.name));
+        for mut country in locations.countries {
+            country.cities.sort_by(|c1, c2| c1.name.cmp(&c2.name));
+            println!("{} ({})", country.name, country.code);
+            for city in &country.cities {
+                println!("\t{} ({}) @ {:?}", city.name, city.code, city.position);
+            }
+            println!("");
+        }
+        Ok(())
+    }
 }
 
 
-fn parse_port(raw_port: &str) -> Result<Constraint<u16>> {
+fn parse_port_constraint(raw_port: &str) -> Result<Constraint<u16>> {
     match raw_port.to_lowercase().as_str() {
         "any" => Ok(Constraint::Any),
         port => Ok(Constraint::Only(
@@ -118,7 +216,7 @@ fn parse_port(raw_port: &str) -> Result<Constraint<u16>> {
 
 /// Parses a protocol constraint string. Can be infallible because the possible values are limited
 /// with clap.
-fn parse_protocol(raw_protocol: &str) -> Constraint<TransportProtocol> {
+fn parse_protocol_constraint(raw_protocol: &str) -> Constraint<TransportProtocol> {
     match raw_protocol.to_lowercase().as_str() {
         "any" => Constraint::Any,
         "udp" => Constraint::Only(TransportProtocol::Udp),
@@ -126,3 +224,19 @@ fn parse_protocol(raw_protocol: &str) -> Constraint<TransportProtocol> {
         _ => unreachable!(),
     }
 }
+
+fn country_code_validator(code: String) -> ::std::result::Result<(), String> {
+    if code.len() == 2 || code == "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 {
+        Ok(())
+    } else {
+        Err(String::from("City codes must be three letters"))
+    }
+}
diff --git a/mullvad-daemon/src/settings.rs b/mullvad-daemon/src/settings.rs
index 4053c36c19..3bb120c544 100644
--- a/mullvad-daemon/src/settings.rs
+++ b/mullvad-daemon/src/settings.rs
@@ -1,9 +1,10 @@
 extern crate serde_json;
 
-use app_dirs::{self, AppDataType};
+use app_dirs;
+
+use mullvad_types::relay_constraints::{Constraint, RelayConstraints, RelaySettings,
+                                       RelaySettingsUpdate};
 
-use mullvad_types::relay_constraints::{Constraint, OpenVpnConstraints, RelayConstraints,
-                                       RelayConstraintsUpdate, TunnelConstraints};
 use std::fs::File;
 use std::io;
 use std::path::PathBuf;
@@ -33,7 +34,7 @@ static SETTINGS_FILE: &str = "settings.json";
 #[serde(default)]
 pub struct Settings {
     account_token: Option<String>,
-    relay_constraints: RelayConstraints,
+    relay_settings: RelaySettings,
 }
 
 impl Default for Settings {
@@ -44,13 +45,10 @@ impl Default for Settings {
 
 const DEFAULT_SETTINGS: Settings = Settings {
     account_token: None,
-    relay_constraints: RelayConstraints {
-        host: Constraint::Any,
-        tunnel: TunnelConstraints::OpenVpn(OpenVpnConstraints {
-            port: Constraint::Any,
-            protocol: Constraint::Any,
-        }),
-    },
+    relay_settings: RelaySettings::Normal(RelayConstraints {
+        location: Constraint::Any,
+        tunnel: Constraint::Any,
+    }),
 };
 
 impl Settings {
@@ -84,7 +82,7 @@ impl Settings {
     }
 
     fn get_settings_path() -> Result<PathBuf> {
-        let dir = app_dirs::app_root(AppDataType::UserConfig, &::APP_INFO)
+        let dir = app_dirs::app_root(app_dirs::AppDataType::UserConfig, &::APP_INFO)
             .chain_err(|| ErrorKind::DirectoryError)?;
         Ok(dir.join(SETTINGS_FILE))
     }
@@ -119,20 +117,20 @@ impl Settings {
         }
     }
 
-    pub fn get_relay_constraints(&self) -> RelayConstraints {
-        self.relay_constraints.clone()
+    pub fn get_relay_settings(&self) -> RelaySettings {
+        self.relay_settings.clone()
     }
 
-    pub fn update_relay_constraints(&mut self, update: RelayConstraintsUpdate) -> Result<bool> {
-        let new_constraints = self.relay_constraints.merge(update);
-        if self.relay_constraints != new_constraints {
+    pub fn update_relay_settings(&mut self, update: RelaySettingsUpdate) -> Result<bool> {
+        let new_settings = self.relay_settings.merge(update);
+        if self.relay_settings != new_settings {
             debug!(
-                "changing relay constraints from {:?} to {:?}",
-                self.relay_constraints,
-                new_constraints
+                "changing relay settings from {:?} to {:?}",
+                self.relay_settings,
+                new_settings
             );
 
-            self.relay_constraints = new_constraints;
+            self.relay_settings = new_settings;
             self.save().map(|_| true)
         } else {
             Ok(false)
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index aa55e55c22..c38c695fc8 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -1,6 +1,14 @@
+use CustomTunnelEndpoint;
+use location::{CityCode, CountryCode};
+
 use std::fmt;
-use talpid_types::net::TransportProtocol;
 
+use talpid_types::net::{OpenVpnParameters, TransportProtocol, WireguardParameters};
+
+
+pub trait Match<T> {
+    fn matches(&self, other: &T) -> bool;
+}
 
 #[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
@@ -9,6 +17,15 @@ pub enum Constraint<T: fmt::Debug + Clone + Eq + PartialEq> {
     Only(T),
 }
 
+impl<T: fmt::Debug + Clone + Eq + PartialEq> Constraint<T> {
+    pub fn unwrap_or(self, other: T) -> T {
+        match self {
+            Constraint::Any => other,
+            Constraint::Only(value) => value,
+        }
+    }
+}
+
 impl<T: fmt::Debug + Clone + Eq + PartialEq> Default for Constraint<T> {
     fn default() -> Self {
         Constraint::Any
@@ -17,81 +34,127 @@ 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 {
+        match *self {
+            Constraint::Any => true,
+            Constraint::Only(ref value) => value == other,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum RelaySettings {
+    CustomTunnelEndpoint(CustomTunnelEndpoint),
+    Normal(RelayConstraints),
+}
+
+impl Default for RelaySettings {
+    fn default() -> Self {
+        RelaySettings::Normal(RelayConstraints::default())
+    }
+}
+
+impl RelaySettings {
+    pub fn merge(&mut self, update: RelaySettingsUpdate) -> Self {
+        match update {
+            RelaySettingsUpdate::CustomTunnelEndpoint(relay) => {
+                RelaySettings::CustomTunnelEndpoint(relay)
+            }
+            RelaySettingsUpdate::Normal(constraint_update) => RelaySettings::Normal(match *self {
+                RelaySettings::CustomTunnelEndpoint(_) => {
+                    RelayConstraints::default().merge(constraint_update)
+                }
+                RelaySettings::Normal(ref constraint) => constraint.merge(constraint_update),
+            }),
+        }
+    }
+}
 
 #[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct RelayConstraints {
-    pub host: Constraint<String>,
-    pub tunnel: TunnelConstraints,
+    pub location: Constraint<LocationConstraint>,
+    pub tunnel: Constraint<TunnelConstraints>,
 }
 
 impl RelayConstraints {
-    pub fn merge(&mut self, update: RelayConstraintsUpdate) -> Self {
+    pub fn merge(&self, update: RelayConstraintsUpdate) -> Self {
         RelayConstraints {
-            host: update.host.unwrap_or_else(|| self.host.clone()),
-            tunnel: self.tunnel.merge(update.tunnel),
+            location: update.location.unwrap_or_else(|| self.location.clone()),
+            tunnel: update.tunnel.unwrap_or_else(|| self.tunnel.clone()),
         }
     }
 }
 
 
 #[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
-#[serde(rename_all = "snake_case")]
+pub enum LocationConstraint {
+    /// A country is represented by its two letter country code.
+    Country(CountryCode),
+    /// A city is composed of a country code and a city code.
+    City(CountryCode, CityCode),
+}
+
+
+#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
 pub enum TunnelConstraints {
     #[serde(rename = "openvpn")] OpenVpn(OpenVpnConstraints),
+    #[serde(rename = "wireguard")] Wireguard(WireguardConstraints),
 }
 
-impl Default for TunnelConstraints {
-    fn default() -> Self {
-        TunnelConstraints::OpenVpn(OpenVpnConstraints::default())
+impl Match<OpenVpnParameters> for TunnelConstraints {
+    fn matches(&self, endpoint: &OpenVpnParameters) -> bool {
+        match *self {
+            TunnelConstraints::OpenVpn(ref constraints) => constraints.matches(endpoint),
+            _ => false,
+        }
     }
 }
 
-impl TunnelConstraints {
-    pub fn merge(&mut self, update: TunnelConstraintsUpdate) -> Self {
+impl Match<WireguardParameters> for TunnelConstraints {
+    fn matches(&self, endpoint: &WireguardParameters) -> bool {
         match *self {
-            TunnelConstraints::OpenVpn(ref mut current) => match update {
-                TunnelConstraintsUpdate::OpenVpn(openvpn_update) => {
-                    TunnelConstraints::OpenVpn(current.merge(openvpn_update))
-                }
-            },
+            TunnelConstraints::Wireguard(ref constraints) => constraints.matches(endpoint),
+            _ => false,
         }
     }
 }
 
 #[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
-#[serde(rename_all = "snake_case")]
 pub struct OpenVpnConstraints {
     pub port: Constraint<u16>,
     pub protocol: Constraint<TransportProtocol>,
 }
 
-impl OpenVpnConstraints {
-    pub fn merge(&mut self, update: OpenVpnConstraintsUpdate) -> Self {
-        OpenVpnConstraints {
-            port: update.port.unwrap_or_else(|| self.port.clone()),
-            protocol: update.protocol.unwrap_or(self.protocol),
-        }
+impl Match<OpenVpnParameters> for OpenVpnConstraints {
+    fn matches(&self, endpoint: &OpenVpnParameters) -> bool {
+        self.port.matches(&endpoint.port) && self.protocol.matches(&endpoint.protocol)
     }
 }
 
+#[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
+pub struct WireguardConstraints {
+    pub port: Constraint<u16>,
+}
 
-#[derive(Debug, Deserialize, Serialize)]
-#[serde(rename_all = "snake_case")]
-pub struct RelayConstraintsUpdate {
-    pub host: Option<Constraint<String>>,
-    pub tunnel: TunnelConstraintsUpdate,
+impl Match<WireguardParameters> for WireguardConstraints {
+    fn matches(&self, endpoint: &WireguardParameters) -> bool {
+        self.port.matches(&endpoint.port)
+    }
 }
 
+
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
-pub enum TunnelConstraintsUpdate {
-    #[serde(rename = "openvpn")] OpenVpn(OpenVpnConstraintsUpdate),
+pub enum RelaySettingsUpdate {
+    CustomTunnelEndpoint(CustomTunnelEndpoint),
+    Normal(RelayConstraintsUpdate),
 }
 
 #[derive(Debug, Default, Deserialize, Serialize)]
-#[serde(rename_all = "snake_case")]
-pub struct OpenVpnConstraintsUpdate {
-    pub port: Option<Constraint<u16>>,
-    pub protocol: Option<Constraint<TransportProtocol>>,
+pub struct RelayConstraintsUpdate {
+    pub location: Option<Constraint<LocationConstraint>>,
+    pub tunnel: Option<Constraint<TunnelConstraints>>,
 }