Set the tunnel port and protocol over the API

3 files changed, 88 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..1787dc0bee
--- /dev/null
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -0,0 +1,85 @@
+use {Command, Result};
+use clap;
+
+use rpc;
+
+use mullvad_types::relay_constraints::{OpenVpnConstraintsUpdate, Port, RelayConstraintsUpdate,
+                                       TunnelConstraintsUpdate};
+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("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"]),
+                    ),
+            )
+    }
+
+    fn run(&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(Some(host)),
+                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
+                    port: None,
+                    protocol: None,
+                }),
+            })
+        } else if let Some(port_matches) = matches.subcommand_matches("port") {
+            let port = value_t_or_exit!(port_matches.value_of("port"), 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 =
+                value_t_or_exit!(protocol_matches.value_of("protocol"), TransportProtocol);
+
+            self.update_constraints(RelayConstraintsUpdate {
+                host: None,
+                tunnel: TunnelConstraintsUpdate::OpenVpn(OpenVpnConstraintsUpdate {
+                    port: None,
+                    protocol: Some(protocol),
+                }),
+            })
+        } 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"))
+    }
+}