Adjust relay list to support wireguard endpoints

7 files changed, 139 insertions, 89 deletions

diff --git a/Cargo.lock b/Cargo.lock
index f868e60fe9..9da1dc3c68 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1005,6 +1005,7 @@ dependencies = [
  "error-chain 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "fern 0.5.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "futures 0.1.25 (registry+https://github.com/rust-lang/crates.io-index)",
+ "ipnetwork 0.14.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "jsonrpc-core 8.0.2 (git+https://github.com/mullvad/jsonrpc?branch=mullvad-fork)",
  "jsonrpc-ipc-server 8.0.1 (git+https://github.com/mullvad/jsonrpc?branch=mullvad-fork)",
  "jsonrpc-macros 8.0.1 (git+https://github.com/mullvad/jsonrpc?branch=mullvad-fork)",
@@ -1127,7 +1128,6 @@ version = "0.1.0"
 dependencies = [
  "chrono 0.4.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "error-chain 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "ipnetwork 0.14.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "lazy_static 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "mullvad-paths 0.1.0",
diff --git a/mullvad-daemon/Cargo.toml b/mullvad-daemon/Cargo.toml
index 62d23dd5a1..19cccdae85 100644
--- a/mullvad-daemon/Cargo.toml
+++ b/mullvad-daemon/Cargo.toml
@@ -27,6 +27,7 @@ jsonrpc-core = { git = "https://github.com/mullvad/jsonrpc", branch = "mullvad-f
 jsonrpc-macros = { git = "https://github.com/mullvad/jsonrpc", branch = "mullvad-fork" }
 jsonrpc-pubsub = { git = "https://github.com/mullvad/jsonrpc", branch = "mullvad-fork" }
 jsonrpc-ipc-server = { git = "https://github.com/mullvad/jsonrpc", branch = "mullvad-fork" }
+ipnetwork = "0.14"
 uuid = { version = "0.6", features = ["v4"] }
 lazy_static = "1.0"
 rand = "0.5"
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index c11ad2df4c..68c7c4bb11 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -4,22 +4,23 @@ use futures::Future;
 
 use mullvad_rpc::{HttpHandle, RelayListProxy};
 use mullvad_types::{
-    endpoint::{MullvadEndpoint, TunnelEndpointData},
+    endpoint::MullvadEndpoint,
     location::Location,
     relay_constraints::{
         Constraint, LocationConstraint, Match, OpenVpnConstraints, RelayConstraints,
-        TunnelConstraints,
+        TunnelConstraints, WireguardConstraints,
     },
-    relay_list::{Relay, RelayList, RelayTunnels},
+    relay_list::{Relay, RelayList, RelayTunnels, WireguardEndpointData},
 };
 
 use serde_json;
 
-use talpid_types::net::TransportProtocol;
+use talpid_types::net::{all_of_the_internet, wireguard, TransportProtocol};
 
 use std::{
     fs::File,
     io,
+    net::{IpAddr, SocketAddr},
     path::{Path, PathBuf},
     sync::{mpsc, Arc, Mutex, MutexGuard},
     thread,
@@ -278,12 +279,8 @@ impl RelaySelector {
                     "Selected relay {} at {}",
                     selected_relay.hostname, selected_relay.ipv4_addr_in
                 );
-                self.get_random_tunnel(&selected_relay.tunnels)
-                    .map(|tunnel_parameters| {
-                        let endpoint = tunnel_parameters
-                            .to_mullvad_endpoint(selected_relay.ipv4_addr_in.into());
-                        (selected_relay.clone(), endpoint)
-                    })
+                self.get_random_tunnel(&selected_relay, &constraints.tunnel)
+                    .map(|endpoint| (selected_relay.clone(), endpoint))
             })
     }
 
@@ -319,14 +316,24 @@ impl RelaySelector {
             Constraint::Any => relay.clone(),
             Constraint::Only(ref tunnel_constraints) => {
                 let mut relay = relay.clone();
-                relay.tunnels = Self::matching_tunnels(&relay.tunnels, tunnel_constraints);
+                relay.tunnels = Self::matching_tunnels(&mut relay.tunnels, tunnel_constraints);
                 relay
             }
         };
-        if relay.tunnels.openvpn.is_empty() {
-            None
-        } else {
+        let relay_matches = match constraints.tunnel {
+            Constraint::Any => {
+                !relay.tunnels.openvpn.is_empty() || !relay.tunnels.wireguard.is_empty()
+            }
+            Constraint::Only(TunnelConstraints::OpenVpn(_)) => !relay.tunnels.openvpn.is_empty(),
+            Constraint::Only(TunnelConstraints::Wireguard(_)) => {
+                !relay.tunnels.wireguard.is_empty()
+            }
+        };
+
+        if relay_matches {
             Some(relay)
+        } else {
+            None
         }
     }
 
@@ -379,11 +386,85 @@ impl RelaySelector {
         }
     }
 
-    fn get_random_tunnel(&mut self, tunnels: &RelayTunnels) -> Option<TunnelEndpointData> {
-        self.rng
-            .choose(&tunnels.openvpn)
-            .cloned()
-            .map(TunnelEndpointData::OpenVpn)
+    fn get_random_tunnel(
+        &mut self,
+        relay: &Relay,
+        constraints: &Constraint<TunnelConstraints>,
+    ) -> Option<MullvadEndpoint> {
+        match constraints {
+            // TODO: Handle Constraint::Any case by selecting from both openvpn and wireguard
+            // tunnels once wireguard is mature enough
+            Constraint::Only(TunnelConstraints::OpenVpn(_)) | Constraint::Any => self
+                .rng
+                .choose(&relay.tunnels.openvpn)
+                .cloned()
+                .map(|endpoint| endpoint.into_mullvad_endpoint(relay.ipv4_addr_in.into())),
+            Constraint::Only(TunnelConstraints::Wireguard(wg_constraints)) => self
+                .rng
+                .choose(&relay.tunnels.wireguard)
+                .cloned()
+                .and_then(|wg_tunnel| {
+                    self.wg_data_to_endpoint(relay.ipv4_addr_in.into(), wg_tunnel, wg_constraints)
+                }),
+        }
+    }
+
+    fn wg_data_to_endpoint(
+        &mut self,
+        host: IpAddr,
+        data: WireguardEndpointData,
+        constraints: &WireguardConstraints,
+    ) -> Option<MullvadEndpoint> {
+        let port = self.get_port_for_wireguard_relay(&data, constraints)?;
+        let peer_config = wireguard::PeerConfig {
+            public_key: data.public_key,
+            endpoint: SocketAddr::new(host, port),
+            allowed_ips: all_of_the_internet(),
+        };
+        Some(MullvadEndpoint::Wireguard {
+            peer: peer_config,
+            gateway: data.ipv4_gateway.into(),
+        })
+    }
+
+    fn get_port_for_wireguard_relay(
+        &mut self,
+        data: &WireguardEndpointData,
+        constraints: &WireguardConstraints,
+    ) -> Option<u16> {
+        match constraints.port {
+            Constraint::Any => {
+                let get_port_amount =
+                    |range: &(u16, u16)| -> u64 { (1 + range.1 - range.0) as u64 };
+                let port_amount: u64 = data.port_ranges.iter().map(get_port_amount).sum();
+
+                if port_amount < 1 {
+                    return None;
+                }
+
+                let mut port_index = self.rng.gen_range(0, port_amount);
+
+                for range in data.port_ranges.iter() {
+                    let ports_in_range = get_port_amount(range);
+                    if port_index < ports_in_range {
+                        return Some(port_index as u16 + range.0);
+                    }
+                    port_index = port_index - ports_in_range;
+                }
+                panic!("Port selection algorithm is broken")
+            }
+            Constraint::Only(port) => {
+                if data
+                    .port_ranges
+                    .iter()
+                    .any(|range| (range.0 <= port && port <= range.1))
+                {
+                    Some(port)
+                } else {
+                    None
+                }
+            }
+        }
     }
 
     /// Try to read the relays, first from cache and if that fails from the resources.
diff --git a/mullvad-types/Cargo.toml b/mullvad-types/Cargo.toml
index 3ea590c995..8239696b9b 100644
--- a/mullvad-types/Cargo.toml
+++ b/mullvad-types/Cargo.toml
@@ -14,7 +14,6 @@ error-chain = "0.12"
 log = "0.4"
 regex = "1"
 lazy_static = "1.1.0"
-ipnetwork = "0.14"
 
 talpid-types = { path = "../talpid-types" }
 mullvad-paths = { path = "../mullvad-paths" }
diff --git a/mullvad-types/src/endpoint.rs b/mullvad-types/src/endpoint.rs
index da0d5b5394..f6818a8854 100644
--- a/mullvad-types/src/endpoint.rs
+++ b/mullvad-types/src/endpoint.rs
@@ -1,13 +1,10 @@
-use ipnetwork::IpNetwork;
 use serde::{Deserialize, Serialize};
-use std::{
-    fmt,
-    net::{IpAddr, SocketAddr},
-};
+use std::{fmt, net::IpAddr};
 use talpid_types::net::{wireguard, Endpoint, TransportProtocol};
 
 use crate::relay_list::{OpenVpnEndpointData, WireguardEndpointData};
 
+
 /// Contains server data needed to conenct to a single mullvad endpoint
 #[derive(Debug, Clone)]
 pub enum MullvadEndpoint {
@@ -68,44 +65,3 @@ impl fmt::Display for TunnelEndpointData {
         }
     }
 }
-
-impl TunnelEndpointData {
-    pub fn to_mullvad_endpoint(self, host: IpAddr) -> MullvadEndpoint {
-        match self {
-            TunnelEndpointData::OpenVpn(metadata) => {
-                MullvadEndpoint::OpenVpn(Endpoint::new(host, metadata.port, metadata.protocol))
-            }
-            TunnelEndpointData::Wireguard(metadata) => {
-                let peer_config = wireguard::PeerConfig {
-                    public_key: metadata.peer_public_key,
-                    endpoint: SocketAddr::new(host, metadata.port),
-                    allowed_ips: all_of_the_internet(),
-                };
-                MullvadEndpoint::Wireguard {
-                    peer: peer_config,
-                    gateway: metadata.gateway,
-                }
-            }
-        }
-    }
-    pub fn port(&self) -> u16 {
-        match self {
-            TunnelEndpointData::OpenVpn(metadata) => metadata.port,
-            TunnelEndpointData::Wireguard(metadata) => metadata.port,
-        }
-    }
-
-    pub fn transport_protocol(&self) -> TransportProtocol {
-        match self {
-            TunnelEndpointData::OpenVpn(metadata) => metadata.protocol,
-            TunnelEndpointData::Wireguard(_) => TransportProtocol::Udp,
-        }
-    }
-}
-
-pub fn all_of_the_internet() -> Vec<IpNetwork> {
-    vec![
-        "0.0.0.0/0".parse().expect("Failed to parse ipv6 network"),
-        "::0/0".parse().expect("Failed to parse ipv6 network"),
-    ]
-}
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 41305115c2..23f0230c84 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -230,7 +230,13 @@ impl fmt::Display for WireguardConstraints {
 
 impl Match<WireguardEndpointData> for WireguardConstraints {
     fn matches(&self, endpoint: &WireguardEndpointData) -> bool {
-        self.port.matches(&endpoint.port)
+        match self.port {
+            Constraint::Any => true,
+            Constraint::Only(port) => endpoint
+                .port_ranges
+                .iter()
+                .any(|range| (port >= range.0 && port <= range.1)),
+        }
     }
 }
 
diff --git a/mullvad-types/src/relay_list.rs b/mullvad-types/src/relay_list.rs
index 1e16dcb6ba..732a672b5a 100644
--- a/mullvad-types/src/relay_list.rs
+++ b/mullvad-types/src/relay_list.rs
@@ -1,10 +1,14 @@
-use crate::location::{CityCode, CountryCode, Location};
+use crate::{
+    endpoint::MullvadEndpoint,
+    location::{CityCode, CountryCode, Location},
+};
+
 use serde::{Deserialize, Serialize};
 use std::{
     fmt,
-    net::{IpAddr, Ipv4Addr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
 };
-use talpid_types::net::{wireguard, TransportProtocol};
+use talpid_types::net::{wireguard, Endpoint, TransportProtocol};
 
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
@@ -52,7 +56,6 @@ pub struct Relay {
 #[serde(default)]
 pub struct RelayTunnels {
     pub openvpn: Vec<OpenVpnEndpointData>,
-    #[serde(skip)]
     pub wireguard: Vec<WireguardEndpointData>,
 }
 
@@ -73,38 +76,42 @@ pub struct OpenVpnEndpointData {
     pub protocol: TransportProtocol,
 }
 
+impl OpenVpnEndpointData {
+    pub fn into_mullvad_endpoint(self, host: IpAddr) -> MullvadEndpoint {
+        MullvadEndpoint::OpenVpn(Endpoint::new(host, self.port, self.protocol))
+    }
+}
+
 impl fmt::Display for OpenVpnEndpointData {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         write!(f, "{} port {}", self.protocol, self.port)
     }
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Deserialize, Serialize)]
+#[derive(Clone, Eq, PartialEq, Hash, Deserialize, Serialize, Debug)]
 pub struct WireguardEndpointData {
     /// Port to connect to
-    pub port: u16,
-    /// Peer's IP address
-    pub gateway: IpAddr,
+    pub port_ranges: Vec<(u16, u16)>,
+    /// Gateways to be used with the tunnel
+    pub ipv4_gateway: Ipv4Addr,
+    pub ipv6_gateway: Ipv6Addr,
     /// The peer's public key
-    pub peer_public_key: wireguard::PublicKey,
-}
-
-impl fmt::Debug for WireguardEndpointData {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
-        f.debug_struct(&"WireguardEndpointData")
-            .field("port", &self.port)
-            .field("gateway", &self.gateway)
-            .field("peer_public_key", &self.peer_public_key)
-            .finish()
-    }
+    pub public_key: wireguard::PublicKey,
 }
 
 impl fmt::Display for WireguardEndpointData {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         write!(
             f,
-            "gateway {} port {} peer_public_key {}",
-            self.gateway, self.port, self.peer_public_key,
+            "gateways {} - {} port_ranges {{ {} }} public_key {}",
+            self.ipv4_gateway,
+            self.ipv6_gateway,
+            self.port_ranges
+                .iter()
+                .map(|range| format!("[{} - {}]", range.0, range.1))
+                .collect::<Vec<_>>()
+                .join(","),
+            self.public_key,
         )
     }
 }