Merge branch 'tcp-tunnel'

3 files changed, 138 insertions, 83 deletions

diff --git a/mullvad-cli/src/cmds/custom_relay.rs b/mullvad-cli/src/cmds/custom_relay.rs
deleted file mode 100644
index d37169d52a..0000000000
--- a/mullvad-cli/src/cmds/custom_relay.rs
+++ /dev/null
@@ -1,80 +0,0 @@
-pub struct CustomRelay;
-
-use {Command, Result};
-use clap;
-use mullvad_types::relay_endpoint::RelayEndpoint;
-
-use rpc;
-
-use talpid_types::net::TransportProtocol;
-
-impl Command for CustomRelay {
-    fn name(&self) -> &'static str {
-        "relay"
-    }
-
-    fn clap_subcommand(&self) -> clap::App<'static, 'static> {
-        clap::SubCommand::with_name(self.name())
-            .about("Set or remove custom relay")
-            .setting(clap::AppSettings::SubcommandRequired)
-            .subcommand(
-                clap::SubCommand::with_name("set")
-                    .about("Set a custom relay")
-                    .arg(
-                        clap::Arg::with_name("host")
-                            .help("The host name or IP of the relay")
-                            .required(true),
-                    )
-                    .arg(
-                        clap::Arg::with_name("port")
-                            .help("The port of the relay")
-                            .required(true),
-                    )
-                    .arg(
-                        clap::Arg::with_name("protocol")
-                            .help(
-                                "The transport protocol. UDP is recommended as it usually results in
-                           higher throughput than TCP",
-                            )
-                            .possible_values(&["udp", "tcp"])
-                            .default_value("udp"),
-                    ),
-            )
-            .subcommand(
-                clap::SubCommand::with_name("remove")
-                    .about("Remove the custom relay and use the default relays instead"),
-            )
-    }
-
-    fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
-        if let Some(set_matches) = matches.subcommand_matches("set") {
-            let host = value_t_or_exit!(set_matches.value_of("host"), String);
-            let port = value_t_or_exit!(set_matches.value_of("port"), u16);
-            let protocol = value_t_or_exit!(set_matches.value_of("protocol"), TransportProtocol);
-
-            self.set(host, port, protocol)
-        } else if let Some(_) = matches.subcommand_matches("remove") {
-            self.remove()
-        } else {
-            unreachable!("No sub command given");
-        }
-    }
-}
-
-impl CustomRelay {
-    fn set(&self, host: String, port: u16, protocol: TransportProtocol) -> Result<()> {
-        let relay_endpoint = RelayEndpoint {
-            host,
-            port,
-            protocol,
-        };
-
-        rpc::call("set_custom_relay", &[relay_endpoint])
-            .map(|_: Option<()>| println!("Custom relay set"))
-    }
-
-    fn remove(&self) -> Result<()> {
-        rpc::call("remove_custom_relay", &[] as &[u8; 0])
-            .map(|_: Option<()>| println!("Custom relay removed"))
-    }
-}
diff --git a/mullvad-cli/src/cmds/mod.rs b/mullvad-cli/src/cmds/mod.rs
index 0104185543..1c10fa81c2 100644
--- a/mullvad-cli/src/cmds/mod.rs
+++ b/mullvad-cli/src/cmds/mod.rs
@@ -13,8 +13,8 @@ pub use self::connect::Connect;
 mod disconnect;
 pub use self::disconnect::Disconnect;
 
-mod custom_relay;
-pub use self::custom_relay::CustomRelay;
+mod relay;
+pub use self::relay::Relay;
 
 mod shutdown;
 pub use self::shutdown::Shutdown;
@@ -26,8 +26,8 @@ pub fn get_commands() -> HashMap<&'static str, Box<Command>> {
         Box::new(Status),
         Box::new(Connect),
         Box::new(Disconnect),
-        Box::new(CustomRelay),
         Box::new(Shutdown),
+        Box::new(Relay),
     ];
     let mut map = HashMap::new();
     for cmd in commands {
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
new file mode 100644
index 0000000000..68582562f9
--- /dev/null
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -0,0 +1,135 @@
+use {Command, Result};
+use clap;
+
+use mullvad_types::relay_constraints::{HostConstraint, OpenVpnConstraintsUpdate, PortConstraint,
+                                       ProtocolConstraint, RelayConstraints,
+                                       RelayConstraintsUpdate, TunnelConstraintsUpdate};
+
+use rpc;
+use talpid_types::net::TransportProtocol;
+
+pub struct Relay;
+
+impl Command for Relay {
+    fn name(&self) -> &'static str {
+        "relay"
+    }
+
+    fn clap_subcommand(&self) -> clap::App<'static, 'static> {
+        clap::SubCommand::with_name(self.name())
+            .about("Manage relay and tunnel constraints")
+            .setting(clap::AppSettings::SubcommandRequired)
+            .subcommand(
+                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)),
+                    )
+                    .subcommand(
+                        clap::SubCommand::with_name("port")
+                            .about("Set port")
+                            .arg(clap::Arg::with_name("port").required(true)),
+                    )
+                    .subcommand(
+                        clap::SubCommand::with_name("protocol")
+                            .about("Set protocol")
+                            .arg(
+                                clap::Arg::with_name("protocol")
+                                    .required(true)
+                                    .possible_values(&["udp", "tcp"]),
+                            ),
+                    ),
+            )
+            .subcommand(clap::SubCommand::with_name("get"))
+    }
+
+    fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
+        if let Some(set_matches) = matches.subcommand_matches("set") {
+            self.set(set_matches)
+        } else if let Some(_) = matches.subcommand_matches("get") {
+            self.get()
+        } else {
+            unreachable!("No relay command given");
+        }
+    }
+}
+
+impl Relay {
+    fn update_constraints(&self, constraints_update: RelayConstraintsUpdate) -> Result<()> {
+        rpc::call("update_relay_constraints", &[constraints_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(HostConstraint::Host(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"))?;
+
+            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"))?;
+
+            self.update_constraints(RelayConstraintsUpdate {
+                host: None,
+                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
+                    port: None,
+                    protocol: Some(protocol),
+                }),
+            })
+        } else {
+            unreachable!("No set relay command given");
+        }
+    }
+
+    fn get(&self) -> Result<()> {
+        let constraints: RelayConstraints = rpc::call("get_relay_constraints", &[] as &[u8; 0])?;
+        println!("Current constraints: {:?}", constraints);
+
+        Ok(())
+    }
+}
+
+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_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())
+}