Revert misunderstanding and make the JSON a little friendlier

1 files changed, 21 insertions, 12 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index ddb5b8f1a6..68582562f9 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -1,10 +1,11 @@
 use {Command, Result};
 use clap;
 
-use rpc;
-
-use mullvad_types::relay_constraints::{RelayConstraints, HostConstraint, OpenVpnConstraintsUpdate, Port,
+use mullvad_types::relay_constraints::{HostConstraint, OpenVpnConstraintsUpdate, PortConstraint,
+                                       ProtocolConstraint, RelayConstraints,
                                        RelayConstraintsUpdate, TunnelConstraintsUpdate};
+
+use rpc;
 use talpid_types::net::TransportProtocol;
 
 pub struct Relay;
@@ -41,9 +42,7 @@ impl Command for Relay {
                             ),
                     ),
             )
-            .subcommand(
-                clap::SubCommand::with_name("get")
-            )
+            .subcommand(clap::SubCommand::with_name("get"))
     }
 
     fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
@@ -85,8 +84,7 @@ impl Relay {
                 }),
             })
         } else if let Some(protocol_matches) = matches.subcommand_matches("protocol") {
-            let protocol =
-                value_t_or_exit!(protocol_matches.value_of("protocol"), TransportProtocol);
+            let protocol = parse_protocol(protocol_matches.value_of("protocol"))?;
 
             self.update_constraints(RelayConstraintsUpdate {
                 host: None,
@@ -108,14 +106,13 @@ impl Relay {
     }
 }
 
-
-fn parse_port(raw_port: Option<&str>) -> Result<Port> {
+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(Port::Port(num)),
+            Ok(num) => Ok(PortConstraint::Port(num)),
             Err(_) => if s.to_lowercase() == "any" {
-                Ok(Port::Any)
+                Ok(PortConstraint::Any)
             } else {
                 bail!("not 'any' or a short".to_owned())
             },
@@ -124,3 +121,15 @@ fn parse_port(raw_port: Option<&str>) -> Result<Port> {
         bail!("not 'any' or a short".to_owned())
     }
 }
+
+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));
+        }
+    }
+
+    bail!("not 'udp' or 'tcp'".to_owned())
+}