Add set_custom_relay API endpoint

6 files changed, 174 insertions, 25 deletions

diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 6953cea498..17b6df06d5 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -43,6 +43,7 @@ mod rpc_info;
 mod settings;
 mod shutdown;
 
+
 use error_chain::ChainedError;
 use futures::Future;
 use jsonrpc_client_http::HttpHandle;
@@ -50,11 +51,13 @@ use jsonrpc_core::futures::sync::oneshot::Sender as OneshotSender;
 use management_interface::{BoxFuture, ManagementInterfaceServer, TunnelCommand};
 use master::AccountsProxy;
 use mullvad_types::account::{AccountData, AccountToken};
+use mullvad_types::relay_endpoint::RelayEndpoint;
 use mullvad_types::states::{DaemonState, SecurityState, TargetState};
 
 use std::io;
-use std::net::Ipv4Addr;
+use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
 use std::path::PathBuf;
+use std::str::FromStr;
 use std::sync::{Arc, Mutex, mpsc};
 use std::thread;
 
@@ -85,12 +88,15 @@ error_chain!{
             description("Invalid settings")
             display("Invalid Settings: {}", msg)
         }
+        NoRelay {
+            description("Found no valid relays to connect to")
+        }
     }
 }
 
 lazy_static! {
-    // Temporary store of hardcoded remotes.
-    static ref REMOTES: [Endpoint; 3] = [
+    // Temporary store of hardcoded relays.
+    static ref RELAYS: [Endpoint; 3] = [
         // se5.mullvad.net
         Endpoint::new(Ipv4Addr::new(193, 138, 219, 240), 1300, TransportProtocol::Udp),
         // se6.mullvad.net
@@ -170,12 +176,12 @@ struct Daemon {
     settings: settings::Settings,
     accounts_proxy: AccountsProxy<HttpHandle>,
     firewall: FirewallProxy,
-    remote_endpoint: Option<Endpoint>,
+    relay_endpoint: Option<Endpoint>,
     tunnel_interface: Option<String>,
 
-    // Just for testing. A cyclic iterator iterating over the hardcoded remotes,
+    // Just for testing. A cyclic iterator iterating over the hardcoded relays,
     // picking a new one for each retry.
-    remote_iter: std::iter::Cycle<std::iter::Cloned<std::slice::Iter<'static, Endpoint>>>,
+    relay_iter: std::iter::Cycle<std::iter::Cloned<std::slice::Iter<'static, Endpoint>>>,
 }
 
 impl Daemon {
@@ -201,9 +207,9 @@ impl Daemon {
                 accounts_proxy: master::create_account_proxy()
                     .chain_err(|| "Unable to bootstrap RPC client")?,
                 firewall: FirewallProxy::new().chain_err(|| ErrorKind::FirewallError)?,
-                remote_endpoint: None,
+                relay_endpoint: None,
                 tunnel_interface: None,
-                remote_iter: REMOTES.iter().cloned().cycle(),
+                relay_iter: RELAYS.iter().cloned().cycle(),
             },
         )
     }
@@ -289,7 +295,7 @@ impl Daemon {
         if let Err(e) = result.chain_err(|| "Tunnel exited in an unexpected way") {
             error!("{}", e.display_chain());
         }
-        self.remote_endpoint = None;
+        self.relay_endpoint = None;
         self.tunnel_interface = None;
         self.reset_security_policy()?;
         self.tunnel_close_handle = None;
@@ -308,6 +314,7 @@ impl Daemon {
             GetAccountData(tx, account_token) => Ok(self.on_get_account_data(tx, account_token)),
             SetAccount(tx, account_token) => self.on_set_account(tx, account_token),
             GetAccount(tx) => Ok(self.on_get_account(tx)),
+            SetCustomRelay(tx, relay_endpoint) => self.on_set_custom_relay(tx, relay_endpoint),
         }
     }
 
@@ -360,6 +367,30 @@ impl Daemon {
         Self::oneshot_send(tx, self.settings.get_account_token(), "current account")
     }
 
+    fn on_set_custom_relay(&mut self,
+                           tx: OneshotSender<()>,
+                           relay_endpoint: RelayEndpoint)
+                           -> Result<()> {
+
+        let save_result = self.settings.set_custom_relay(Some(relay_endpoint));
+        match save_result.chain_err(|| "Unable to save settings") {
+            Ok(servers_changed) => {
+                Self::oneshot_send(tx, (), "set_custom_relay response");
+
+                let tunnel_needs_restart = self.state == TunnelState::Connecting ||
+                                           self.state == TunnelState::Connected;
+
+                if servers_changed && tunnel_needs_restart {
+                    info!("Initiating tunnel restart because a custom relay was selected");
+                    self.kill_tunnel()?;
+                }
+            }
+            Err(e) => error!("{}", e.display_chain()),
+        }
+
+        Ok(())
+    }
+
     fn oneshot_send<T>(tx: OneshotSender<T>, t: T, msg: &'static str) {
         if let Err(_) = tx.send(t) {
             warn!("Unable to send {} to management interface client", msg);
@@ -440,7 +471,7 @@ impl Daemon {
                 debug!("Triggering tunnel start");
                 if let Err(e) = self.start_tunnel().chain_err(|| "Failed to start tunnel") {
                     error!("{}", e.display_chain());
-                    self.remote_endpoint = None;
+                    self.relay_endpoint = None;
                     self.reset_security_policy()?;
                     self.management_interface_broadcaster.notify_error(&e);
                     self.set_target_state(TargetState::Unsecured)?;
@@ -458,26 +489,82 @@ impl Daemon {
             self.state == TunnelState::NotRunning,
             ErrorKind::InvalidState
         );
-        let remote = self.remote_iter.next().unwrap();
+
+        let relay = self.get_relay()
+            .chain_err(|| ErrorKind::NoRelay)?;
+
         let account_token = self.settings
             .get_account_token()
             .ok_or(ErrorKind::InvalidSettings("No account token"))?;
-        self.remote_endpoint = Some(remote);
+
+        self.relay_endpoint = Some(relay);
         self.set_security_policy()?;
-        let tunnel_monitor = self.spawn_tunnel_monitor(remote, &account_token)?;
+
+        let tunnel_monitor = self.spawn_tunnel_monitor(relay, &account_token)?;
         self.tunnel_close_handle = Some(tunnel_monitor.close_handle());
         self.spawn_tunnel_monitor_wait_thread(tunnel_monitor);
+
         self.set_state(TunnelState::Connecting)?;
         Ok(())
     }
 
-    fn spawn_tunnel_monitor(&self, remote: Endpoint, account_token: &str) -> Result<TunnelMonitor> {
+    fn get_relay(&mut self) -> Result<Endpoint> {
+        if let Some(relay_endpoint) = self.settings.get_custom_relay() {
+            self.parse_custom_relay(&relay_endpoint)
+        } else {
+            Ok(self.relay_iter.next().unwrap())
+        }
+    }
+
+    fn parse_custom_relay(&self, relay_endpoint: &RelayEndpoint) -> Result<Endpoint> {
+        let socket_addrs: Vec<SocketAddr> =
+            format!("{}:{}", &relay_endpoint.host, relay_endpoint.port)
+                .to_socket_addrs()
+                .chain_err(
+                    || {
+                        format!(
+                            "Invalid custom server host identifier: {}",
+                            relay_endpoint.host
+                        )
+                    },
+                )?
+                .collect();
+
+        if socket_addrs.len() == 0 {
+            bail!("Unable to resolve {}", relay_endpoint.host)
+        } else {
+
+            let socket_addr = socket_addrs[0];
+
+            if socket_addrs.len() > 1 {
+                info!(
+                    "{} resolved to more than one IP, ignoring all but {}",
+                    relay_endpoint.host,
+                    socket_addr.ip()
+                )
+            }
+
+            let protocol = TransportProtocol::from_str(&relay_endpoint.protocol)
+                .map_err(
+                    |_| {
+                        format!(
+                            "Invalid custom server protocol: {}",
+                            relay_endpoint.protocol
+                        )
+                    },
+                )?;
+
+            Ok(Endpoint::new(socket_addr.ip(), socket_addr.port(), protocol),)
+        }
+    }
+
+    fn spawn_tunnel_monitor(&self, relay: Endpoint, account_token: &str) -> Result<TunnelMonitor> {
         // Must wrap the channel in a Mutex because TunnelMonitor forces the closure to be Sync
         let event_tx = Arc::new(Mutex::new(self.tx.clone()));
         let on_tunnel_event = move |event| {
             let _ = event_tx.lock().unwrap().send(DaemonEvent::TunnelEvent(event));
         };
-        TunnelMonitor::new(remote, account_token, on_tunnel_event)
+        TunnelMonitor::new(relay, account_token, on_tunnel_event)
             .chain_err(|| ErrorKind::TunnelError("Unable to start tunnel monitor"))
     }
 
@@ -515,9 +602,9 @@ impl Daemon {
     }
 
     fn set_security_policy(&mut self) -> Result<()> {
-        let policy = match (self.remote_endpoint, self.tunnel_interface.as_ref()) {
-            (Some(remote), None) => SecurityPolicy::Connecting(remote),
-            (Some(remote), Some(interface)) => SecurityPolicy::Connected(remote, interface.clone()),
+        let policy = match (self.relay_endpoint, self.tunnel_interface.as_ref()) {
+            (Some(relay), None) => SecurityPolicy::Connecting(relay),
+            (Some(relay), Some(interface)) => SecurityPolicy::Connected(relay, interface.clone()),
             _ => bail!(ErrorKind::InvalidState),
         };
         debug!("Set security policy: {:?}", policy);
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index bca4f68f13..15e5132a85 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -10,6 +10,7 @@ use jsonrpc_pubsub::{PubSubHandler, PubSubMetadata, Session, SubscriptionId};
 use jsonrpc_ws_server;
 use mullvad_types::account::{AccountData, AccountToken};
 use mullvad_types::location::{CountryCode, Location};
+use mullvad_types::relay_endpoint::RelayEndpoint;
 use mullvad_types::states::{DaemonState, TargetState};
 
 use serde;
@@ -49,9 +50,9 @@ build_rpc_trait! {
         #[rpc(async, name = "get_account")]
         fn get_account(&self) -> BoxFuture<Option<AccountToken>, Error>;
 
-        /// Set which country to connect to
-        #[rpc(name = "set_country")]
-        fn set_country(&self, CountryCode) -> Result<(), Error>;
+        /// Set which relay to connect to
+        #[rpc(async, name = "set_custom_relay")]
+        fn set_custom_relay(&self, RelayEndpoint) -> BoxFuture<(), Error>;
 
         /// Set if the client should automatically establish a tunnel on start or not.
         #[rpc(name = "set_autoconnect")]
@@ -115,6 +116,8 @@ pub enum TunnelCommand {
     SetAccount(OneshotSender<()>, Option<AccountToken>),
     /// Request the current account token being used.
     GetAccount(OneshotSender<Option<AccountToken>>),
+    /// Set a custom relay instead of the default list of relays
+    SetCustomRelay(OneshotSender<()>, RelayEndpoint),
 }
 
 #[derive(Default)]
@@ -312,9 +315,12 @@ impl<T: From<TunnelCommand> + 'static + Send> ManagementInterfaceApi for Managem
         Box::new(future)
     }
 
-    fn set_country(&self, _country_code: CountryCode) -> Result<(), Error> {
-        trace!("set_country");
-        Ok(())
+    fn set_custom_relay(&self, custom_relay: RelayEndpoint) -> BoxFuture<(), Error> {
+        trace!("set_custom_relay");
+        let (tx, rx) = sync::oneshot::channel();
+        let future = self.send_command_to_daemon(TunnelCommand::SetCustomRelay(tx, custom_relay))
+            .and_then(|_| rx.map_err(|_| Error::internal_error()));
+        Box::new(future)
     }
 
     fn set_autoconnect(&self, _autoconnect: bool) -> Result<(), Error> {
diff --git a/mullvad-daemon/src/settings.rs b/mullvad-daemon/src/settings.rs
index d2d652adf0..ccced02d02 100644
--- a/mullvad-daemon/src/settings.rs
+++ b/mullvad-daemon/src/settings.rs
@@ -3,6 +3,8 @@ extern crate toml;
 
 use self::app_dirs::{AppDataType, AppInfo};
 
+use mullvad_types::relay_endpoint::RelayEndpoint;
+
 use std::fs::File;
 use std::io::{self, Read, Write};
 use std::path::PathBuf;
@@ -36,9 +38,13 @@ static SETTINGS_FILE: &str = "settings.toml";
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Settings {
     account_token: Option<String>,
+    custom_relay: Option<RelayEndpoint>,
 }
 
-const DEFAULT_SETTINGS: Settings = Settings { account_token: None };
+const DEFAULT_SETTINGS: Settings = Settings {
+    account_token: None,
+    custom_relay: None,
+};
 
 impl Settings {
     /// Loads user settings from file. If no file is present it returns the defaults.
@@ -111,4 +117,26 @@ impl Settings {
             None => "[none]".to_owned(),
         }
     }
+
+    pub fn get_custom_relay(&self) -> Option<RelayEndpoint> {
+        self.custom_relay.clone()
+    }
+
+    pub fn set_custom_relay(&mut self, relay_endpoint: Option<RelayEndpoint>) -> Result<bool> {
+        let has_changed = self.custom_relay != relay_endpoint;
+
+        if has_changed {
+            match &relay_endpoint {
+                &Some(ref data) => info!("Setting a custom relay: {}", data),
+                &None => info!("Removing the custom relay"),
+            }
+
+            self.custom_relay = relay_endpoint;
+
+            self.save()
+                .map(|_| true)
+        } else {
+            Ok(false)
+        }
+    }
 }
diff --git a/mullvad-types/src/lib.rs b/mullvad-types/src/lib.rs
index 0886a3c00b..cc4f035c93 100644
--- a/mullvad-types/src/lib.rs
+++ b/mullvad-types/src/lib.rs
@@ -14,3 +14,4 @@ extern crate serde;
 pub mod account;
 pub mod location;
 pub mod states;
+pub mod relay_endpoint;
diff --git a/mullvad-types/src/relay_endpoint.rs b/mullvad-types/src/relay_endpoint.rs
new file mode 100644
index 0000000000..d02f9d9d56
--- /dev/null
+++ b/mullvad-types/src/relay_endpoint.rs
@@ -0,0 +1,14 @@
+use std::fmt;
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+pub struct RelayEndpoint {
+    pub host: String,
+    pub port: u16,
+    pub protocol: String,
+}
+
+impl fmt::Display for RelayEndpoint {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "{}:{} - {}", self.host, self.port, self.protocol)
+    }
+}
diff --git a/talpid-core/src/net.rs b/talpid-core/src/net.rs
index ba8a3fe665..9a75859d61 100644
--- a/talpid-core/src/net.rs
+++ b/talpid-core/src/net.rs
@@ -1,4 +1,5 @@
 use std::net::{IpAddr, SocketAddr};
+use std::str::FromStr;
 
 /// Represents a network layer IP address together with the transport layer protocol and port.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -27,3 +28,15 @@ pub enum TransportProtocol {
     /// Represents the TCP transport protocol.
     Tcp,
 }
+
+impl FromStr for TransportProtocol {
+    type Err = ();
+
+    fn from_str(s: &str) -> ::std::result::Result<TransportProtocol, Self::Err> {
+        match s {
+            "udp" => Ok(TransportProtocol::Udp),
+            "tcp" => Ok(TransportProtocol::Tcp),
+            _ => Err(()),
+        }
+    }
+}