Review fixes

10 files changed, 116 insertions, 69 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 4258357e9f..fc8263b78d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -681,6 +681,8 @@ name = "mullvad-types"
 version = "0.1.0"
 dependencies = [
  "chrono 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "error-chain 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "log 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.9 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_derive 1.0.9 (registry+https://github.com/rust-lang/crates.io-index)",
  "talpid-types 0.1.0",
diff --git a/mullvad-cli/src/cmds/account.rs b/mullvad-cli/src/cmds/account.rs
index 553458e11a..f7cb106fc3 100644
--- a/mullvad-cli/src/cmds/account.rs
+++ b/mullvad-cli/src/cmds/account.rs
@@ -1,5 +1,4 @@
-use Command;
-use Result;
+use {Command, Result};
 use clap;
 
 use mullvad_types::account::{AccountData, AccountToken};
diff --git a/mullvad-cli/src/cmds/custom_relay.rs b/mullvad-cli/src/cmds/custom_relay.rs
index 90b4b5bd0e..172040322e 100644
--- a/mullvad-cli/src/cmds/custom_relay.rs
+++ b/mullvad-cli/src/cmds/custom_relay.rs
@@ -1,7 +1,6 @@
 pub struct CustomRelay;
 
-use Command;
-use Result;
+use {Command, Result};
 use clap;
 use mullvad_types::relay_endpoint::RelayEndpoint;
 
@@ -16,22 +15,22 @@ impl Command for CustomRelay {
 
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
         clap::SubCommand::with_name(self.name())
-            .about("Set or remove custom remote relay")
+            .about("Set or remove custom relay")
             .setting(clap::AppSettings::SubcommandRequired)
             .subcommand(clap::SubCommand::with_name("set")
-                .about("Set a custom remote relay")
+                .about("Set a custom relay")
                 .arg(clap::Arg::with_name("host")
-                     .help("The host name or IP of the remote relay")
+                     .help("The host name or IP of the relay")
                      .required(true))
                 .arg(clap::Arg::with_name("port")
-                     .help("The port of the remote relay")
+                     .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 that TCP")
+                     .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 remote server and use the default remotes instead"))
+                .about("Remove the custom relay and use the default relays instead"))
     }
 
     fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
@@ -61,7 +60,7 @@ impl CustomRelay {
             "set_custom_relay",
             &[relay_endpoint],
         )
-                .map(|_: Option<()>| println!("Custom remote relay set"))
+                .map(|_: Option<()>| println!("Custom relay set"))
     }
 
     fn remove(&self) -> Result<()> {
diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 764427f1e1..54b1ac1747 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -56,7 +56,7 @@ use mullvad_types::relay_endpoint::RelayEndpoint;
 use mullvad_types::states::{DaemonState, SecurityState, TargetState};
 
 use std::io;
-use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
+use std::net::Ipv4Addr;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex, mpsc};
 use std::thread;
@@ -375,13 +375,13 @@ impl Daemon {
         let save_result = self.settings.set_custom_relay(relay_endpoint);
 
         match save_result.chain_err(|| "Unable to save settings") {
-            Ok(servers_changed) => {
+            Ok(relays_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 {
+                if relays_changed && tunnel_needs_restart {
                     info!("Initiating tunnel restart because a custom relay was selected");
                     self.kill_tunnel()?;
                 }
@@ -511,50 +511,14 @@ impl Daemon {
 
     fn get_relay(&mut self) -> Result<Endpoint> {
         if let Some(relay_endpoint) = self.settings.get_custom_relay() {
-            self.parse_custom_relay(&relay_endpoint)
+            relay_endpoint
+                .to_endpoint()
+                .chain_err(|| "Invalid custom relay")
         } 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()
-                )
-            }
-
-            Ok(
-                Endpoint::new(
-                    socket_addr.ip(),
-                    socket_addr.port(),
-                    relay_endpoint.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()));
diff --git a/mullvad-daemon/src/settings.rs b/mullvad-daemon/src/settings.rs
index ccced02d02..63a4325d5e 100644
--- a/mullvad-daemon/src/settings.rs
+++ b/mullvad-daemon/src/settings.rs
@@ -123,9 +123,7 @@ impl Settings {
     }
 
     pub fn set_custom_relay(&mut self, relay_endpoint: Option<RelayEndpoint>) -> Result<bool> {
-        let has_changed = self.custom_relay != relay_endpoint;
-
-        if has_changed {
+        if self.custom_relay != relay_endpoint {
             match &relay_endpoint {
                 &Some(ref data) => info!("Setting a custom relay: {}", data),
                 &None => info!("Removing the custom relay"),
diff --git a/mullvad-types/Cargo.toml b/mullvad-types/Cargo.toml
index 3f6c13f0dc..e710e67d87 100644
--- a/mullvad-types/Cargo.toml
+++ b/mullvad-types/Cargo.toml
@@ -8,5 +8,7 @@ license = "GPL-3.0"
 chrono = { version = "0.4", features = ["serde"] }
 serde_derive = "1.0"
 serde = "1.0"
+error-chain = "0.11"
+log = "0.3"
 
 talpid-types = { path = "../talpid-types" }
diff --git a/mullvad-types/src/lib.rs b/mullvad-types/src/lib.rs
index 2801efefad..608e3eee80 100644
--- a/mullvad-types/src/lib.rs
+++ b/mullvad-types/src/lib.rs
@@ -13,6 +13,12 @@ extern crate serde;
 
 extern crate talpid_types;
 
+#[macro_use]
+extern crate log;
+
+#[macro_use]
+extern crate error_chain;
+
 pub mod account;
 pub mod location;
 pub mod states;
diff --git a/mullvad-types/src/relay_endpoint.rs b/mullvad-types/src/relay_endpoint.rs
index d5a80a5879..e5335eb810 100644
--- a/mullvad-types/src/relay_endpoint.rs
+++ b/mullvad-types/src/relay_endpoint.rs
@@ -1,6 +1,17 @@
+use std::cmp::Ordering;
 use std::fmt;
+use std::net::{SocketAddr, ToSocketAddrs};
+use talpid_types;
 use talpid_types::net::TransportProtocol;
 
+error_chain!{
+    errors {
+        InvalidHost(host: String) {
+            display("Invalid host: {}", host)
+        }
+    }
+}
+
 #[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
 pub struct RelayEndpoint {
     pub host: String,
@@ -8,8 +19,57 @@ pub struct RelayEndpoint {
     pub protocol: TransportProtocol,
 }
 
+impl RelayEndpoint {
+    pub fn to_endpoint(&self) -> Result<talpid_types::net::Endpoint> {
+
+        let mut socket_addrs = to_socket_addrs(self.host.as_str(), self.port)?;
+        ensure!(
+            socket_addrs.len() > 0,
+            ErrorKind::InvalidHost(self.host.clone())
+        );
+
+        let socket_addr = choose_ip(&mut socket_addrs);
+
+        if socket_addrs.len() > 1 {
+            info!(
+                "{} resolved to more than one IP, ignoring all but {}",
+                self.host,
+                socket_addr.ip()
+            )
+        }
+
+        Ok(talpid_types::net::Endpoint::new(socket_addr.ip(), socket_addr.port(), self.protocol),)
+    }
+}
+
+/// Does a DNS lookup if the host isn't an IP.
+fn to_socket_addrs(host: &str, port: u16) -> Result<Vec<SocketAddr>> {
+    Ok(
+        (host, port)
+            .to_socket_addrs()
+            .chain_err(|| ErrorKind::InvalidHost(host.to_owned()))?
+            .collect(),
+    )
+}
+
+fn choose_ip(socket_addrs: &mut Vec<SocketAddr>) -> SocketAddr {
+
+    // We want to prefer IPv4 addresses so we sort the addresses putting
+    // the IPv4 addresses at the start of the vector.
+    socket_addrs.sort_by(
+        |a, _| if a.is_ipv4() {
+            Ordering::Less
+        } else {
+            Ordering::Equal
+        },
+    );
+
+    // If there are many IP:s, we simply ignore the rest
+    socket_addrs[0]
+}
+
 impl fmt::Display for RelayEndpoint {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        write!(fmt, "{}:{} - {:?}", self.host, self.port, self.protocol)
+        write!(fmt, "{}:{} - {}", self.host, self.port, self.protocol)
     }
 }
diff --git a/talpid-core/src/firewall/macos.rs b/talpid-core/src/firewall/macos.rs
index 8f0ba7a897..e3bfa9870f 100644
--- a/talpid-core/src/firewall/macos.rs
+++ b/talpid-core/src/firewall/macos.rs
@@ -76,13 +76,12 @@ impl PacketFilter {
     }
 
     fn get_relay_rule(relay_endpoint: net::Endpoint) -> Result<pfctl::FilterRule> {
-        let pfctl_endpoint = as_pfctl_endpoint(relay_endpoint);
         let pfctl_proto = as_pfctl_proto(relay_endpoint.protocol);
 
         pfctl::FilterRuleBuilder::default()
             .action(pfctl::FilterRuleAction::Pass)
             .direction(pfctl::Direction::Out)
-            .to(pfctl_endpoint)
+            .to(relay_endpoint.address)
             .proto(pfctl_proto)
             .keep_state(pfctl::StatePolicy::Keep)
             .tcp_flags(Self::get_tcp_flags())
@@ -169,17 +168,9 @@ impl PacketFilter {
     }
 }
 
-fn as_pfctl_endpoint(relay_endpoint: net::Endpoint) -> pfctl::Endpoint {
-    pfctl::Endpoint::new(
-        pfctl::Ip::from(relay_endpoint.address.ip()),
-        pfctl::Port::from(relay_endpoint.address.port())
-    )
-}
-
 fn as_pfctl_proto(protocol: net::TransportProtocol) -> pfctl::Proto {
     match protocol {
         net::TransportProtocol::Udp => pfctl::Proto::Udp,
         net::TransportProtocol::Tcp => pfctl::Proto::Tcp,
     }
 }
-
diff --git a/talpid-types/src/net.rs b/talpid-types/src/net.rs
index 18d9d1a768..d6e71da9be 100644
--- a/talpid-types/src/net.rs
+++ b/talpid-types/src/net.rs
@@ -1,3 +1,5 @@
+use std::error::Error;
+use std::fmt;
 use std::net::{IpAddr, SocketAddr};
 use std::str::FromStr;
 
@@ -30,13 +32,37 @@ pub enum TransportProtocol {
 }
 
 impl FromStr for TransportProtocol {
-    type Err = ();
+    type Err = TransportProtocolParseError;
 
     fn from_str(s: &str) -> ::std::result::Result<TransportProtocol, Self::Err> {
         match s {
             "udp" => Ok(TransportProtocol::Udp),
             "tcp" => Ok(TransportProtocol::Tcp),
-            _ => Err(()),
+            _ => Err(TransportProtocolParseError),
         }
     }
 }
+
+impl fmt::Display for TransportProtocol {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            TransportProtocol::Udp => "UDP".fmt(fmt),
+            TransportProtocol::Tcp => "TCP".fmt(fmt),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct TransportProtocolParseError;
+
+impl fmt::Display for TransportProtocolParseError {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.write_str(self.description())
+    }
+}
+
+impl Error for TransportProtocolParseError {
+    fn description(&self) -> &str {
+        "Not a valid transport protocol"
+    }
+}