Merge branch 'generalize-constraint'

4 files changed, 68 insertions, 76 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 68582562f9..817ae17b3a 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -1,8 +1,8 @@
-use {Command, Result};
+use {Command, Result, ResultExt};
 use clap;
+use std::str::FromStr;
 
-use mullvad_types::relay_constraints::{HostConstraint, OpenVpnConstraintsUpdate, PortConstraint,
-                                       ProtocolConstraint, RelayConstraints,
+use mullvad_types::relay_constraints::{Constraint, OpenVpnConstraintsUpdate, RelayConstraints,
                                        RelayConstraintsUpdate, TunnelConstraintsUpdate};
 
 use rpc;
@@ -38,7 +38,7 @@ impl Command for Relay {
                             .arg(
                                 clap::Arg::with_name("protocol")
                                     .required(true)
-                                    .possible_values(&["udp", "tcp"]),
+                                    .possible_values(&["any", "udp", "tcp"]),
                             ),
                     ),
             )
@@ -67,14 +67,14 @@ impl Relay {
             let host = value_t_or_exit!(host_matches.value_of("host"), String);
 
             self.update_constraints(RelayConstraintsUpdate {
-                host: Some(HostConstraint::Host(host)),
+                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"))?;
+            let port = parse_port(port_matches.value_of("port").unwrap())?;
 
             self.update_constraints(RelayConstraintsUpdate {
                 host: None,
@@ -84,7 +84,7 @@ impl Relay {
                 }),
             })
         } else if let Some(protocol_matches) = matches.subcommand_matches("protocol") {
-            let protocol = parse_protocol(protocol_matches.value_of("protocol"))?;
+            let protocol = parse_protocol(protocol_matches.value_of("protocol").unwrap());
 
             self.update_constraints(RelayConstraintsUpdate {
                 host: None,
@@ -106,30 +106,23 @@ impl Relay {
     }
 }
 
-fn parse_port(raw_port: Option<&str>) -> Result<PortConstraint> {
-    if let Some(s) = raw_port {
-        let res = u16::from_str_radix(s, 10);
-        match res {
-            Ok(num) => Ok(PortConstraint::Port(num)),
-            Err(_) => if s.to_lowercase() == "any" {
-                Ok(PortConstraint::Any)
-            } else {
-                bail!("not 'any' or a short".to_owned())
-            },
-        }
-    } else {
-        bail!("not 'any' or a short".to_owned())
+
+fn parse_port(raw_port: &str) -> Result<Constraint<u16>> {
+    match raw_port.to_lowercase().as_str() {
+        "any" => Ok(Constraint::Any),
+        port => Ok(Constraint::Only(
+            u16::from_str(port).chain_err(|| "Invalid port")?,
+        )),
     }
 }
 
-fn parse_protocol(raw_protocol: Option<&str>) -> Result<ProtocolConstraint> {
-    if let Some(s) = raw_protocol {
-        if s.to_lowercase() == "any" {
-            return Ok(ProtocolConstraint::Any);
-        } else if ["udp", "tcp"].contains(&s) {
-            return Ok(ProtocolConstraint::Protocol(TransportProtocol::Udp));
-        }
+/// Parses a protocol constraint string. Can be infallible because the possible values are limited
+/// with clap.
+fn parse_protocol(raw_protocol: &str) -> Constraint<TransportProtocol> {
+    match raw_protocol.to_lowercase().as_str() {
+        "any" => Constraint::Any,
+        "udp" => Constraint::Only(TransportProtocol::Udp),
+        "tcp" => Constraint::Only(TransportProtocol::Tcp),
+        _ => unreachable!(),
     }
-
-    bail!("not 'udp' or 'tcp'".to_owned())
 }
diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index e4edb30527..73a77e2008 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -49,8 +49,7 @@ use jsonrpc_core::futures::sync::oneshot::Sender as OneshotSender;
 use management_interface::{BoxFuture, ManagementInterfaceServer, TunnelCommand};
 use mullvad_rpc::{AccountsProxy, HttpHandle};
 use mullvad_types::account::{AccountData, AccountToken};
-use mullvad_types::relay_constraints::{HostConstraint, OpenVpnConstraints, PortConstraint,
-                                       ProtocolConstraint, RelayConstraints,
+use mullvad_types::relay_constraints::{Constraint, OpenVpnConstraints, RelayConstraints,
                                        RelayConstraintsUpdate, TunnelConstraints};
 use mullvad_types::relay_endpoint::RelayEndpoint;
 use mullvad_types::states::{DaemonState, SecurityState, TargetState};
@@ -532,8 +531,8 @@ impl Daemon {
         let relay_constraints = self.settings.get_relay_constraints();
 
         let host = match relay_constraints.host {
-            HostConstraint::Any => format!("{}", self.relay_iter.next().unwrap().address),
-            HostConstraint::Host(host) => host,
+            Constraint::Any => format!("{}", self.relay_iter.next().unwrap().address),
+            Constraint::Only(host) => host,
         };
 
         match relay_constraints.tunnel {
@@ -547,13 +546,12 @@ impl Daemon {
         constraints: OpenVpnConstraints,
     ) -> Result<Endpoint> {
         let protocol = match constraints.protocol {
-            ProtocolConstraint::Any => TransportProtocol::Udp,
-            ProtocolConstraint::Protocol(protocol) => protocol,
+            Constraint::Any => TransportProtocol::Udp,
+            Constraint::Only(protocol) => protocol,
         };
-
         let port = match constraints.port {
-            PortConstraint::Any => randomize_port(protocol),
-            PortConstraint::Port(port) => port,
+            Constraint::Any => randomize_port(protocol),
+            Constraint::Only(port) => port,
         };
 
         RelayEndpoint {
diff --git a/mullvad-daemon/src/settings.rs b/mullvad-daemon/src/settings.rs
index d1a852c11d..10cf63123a 100644
--- a/mullvad-daemon/src/settings.rs
+++ b/mullvad-daemon/src/settings.rs
@@ -3,8 +3,7 @@ extern crate serde_json;
 
 use self::app_dirs::{AppDataType, AppInfo};
 
-use mullvad_types::relay_constraints::{HostConstraint, OpenVpnConstraints, PortConstraint,
-                                       ProtocolConstraint, RelayConstraints,
+use mullvad_types::relay_constraints::{Constraint, OpenVpnConstraints, RelayConstraints,
                                        RelayConstraintsUpdate, TunnelConstraints};
 use std::fs::File;
 use std::io;
@@ -52,10 +51,10 @@ impl Default for Settings {
 const DEFAULT_SETTINGS: Settings = Settings {
     account_token: None,
     relay_constraints: RelayConstraints {
-        host: HostConstraint::Any,
+        host: Constraint::Any,
         tunnel: TunnelConstraints::OpenVpn(OpenVpnConstraints {
-            port: PortConstraint::Any,
-            protocol: ProtocolConstraint::Any,
+            port: Constraint::Any,
+            protocol: Constraint::Any,
         }),
     },
 };
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 26a3a04f84..aa55e55c22 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -1,9 +1,27 @@
+use std::fmt;
 use talpid_types::net::TransportProtocol;
 
-#[derive(Debug, Clone, Deserialize, Serialize, Eq, PartialEq)]
+
+#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Constraint<T: fmt::Debug + Clone + Eq + PartialEq> {
+    Any,
+    Only(T),
+}
+
+impl<T: fmt::Debug + Clone + Eq + PartialEq> Default for Constraint<T> {
+    fn default() -> Self {
+        Constraint::Any
+    }
+}
+
+impl<T: Copy + fmt::Debug + Clone + Eq + PartialEq> Copy for Constraint<T> {}
+
+
+#[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct RelayConstraints {
-    pub host: HostConstraint,
+    pub host: Constraint<String>,
     pub tunnel: TunnelConstraints,
 }
 
@@ -17,12 +35,18 @@ impl RelayConstraints {
 }
 
 
-#[derive(Debug, Clone, Deserialize, Serialize, Eq, PartialEq)]
+#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum TunnelConstraints {
     #[serde(rename = "openvpn")] OpenVpn(OpenVpnConstraints),
 }
 
+impl Default for TunnelConstraints {
+    fn default() -> Self {
+        TunnelConstraints::OpenVpn(OpenVpnConstraints::default())
+    }
+}
+
 impl TunnelConstraints {
     pub fn merge(&mut self, update: TunnelConstraintsUpdate) -> Self {
         match *self {
@@ -35,18 +59,18 @@ impl TunnelConstraints {
     }
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct OpenVpnConstraints {
-    pub port: PortConstraint,
-    pub protocol: ProtocolConstraint,
+    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_else(|| self.protocol.clone()),
+            protocol: update.protocol.unwrap_or(self.protocol),
         }
     }
 }
@@ -55,7 +79,7 @@ impl OpenVpnConstraints {
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct RelayConstraintsUpdate {
-    pub host: Option<HostConstraint>,
+    pub host: Option<Constraint<String>>,
     pub tunnel: TunnelConstraintsUpdate,
 }
 
@@ -65,31 +89,9 @@ pub enum TunnelConstraintsUpdate {
     #[serde(rename = "openvpn")] OpenVpn(OpenVpnConstraintsUpdate),
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Default, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct OpenVpnConstraintsUpdate {
-    pub port: Option<PortConstraint>,
-    pub protocol: Option<ProtocolConstraint>,
-}
-
-
-#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
-#[serde(rename_all = "snake_case")]
-pub enum HostConstraint {
-    Any,
-    Host(String),
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
-#[serde(rename_all = "snake_case")]
-pub enum PortConstraint {
-    Any,
-    Port(u16),
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
-#[serde(rename_all = "snake_case")]
-pub enum ProtocolConstraint {
-    Any,
-    Protocol(TransportProtocol),
+    pub port: Option<Constraint<u16>>,
+    pub protocol: Option<Constraint<TransportProtocol>>,
 }