Add IP version constraint for WireGuard

5 files changed, 93 insertions, 39 deletions

diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index c87b846387..3866f51014 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -27,7 +27,7 @@ use std::{
     sync::{mpsc, Arc},
 };
 use talpid_types::{
-    net::{TransportProtocol, TunnelType},
+    net::{IpVersion, TransportProtocol, TunnelType},
     ErrorExt,
 };
 
@@ -949,6 +949,20 @@ fn convert_relay_settings_update(
             } else {
                 None
             };
+            let ip_protocol = if let Some(ref constraints) = settings.wireguard_constraints {
+                match &constraints.ip_protocol {
+                    Some(constraint) => match types::IpVersion::from_i32(constraint.protocol) {
+                        Some(types::IpVersion::V4) => Some(IpVersion::V4),
+                        Some(types::IpVersion::V6) => Some(IpVersion::V6),
+                        None => {
+                            return Err(Status::invalid_argument("unknown ip protocol version"))
+                        }
+                    },
+                    None => None,
+                }
+            } else {
+                None
+            };
 
             Ok(RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
                 location,
@@ -961,6 +975,7 @@ fn convert_relay_settings_update(
                         } else {
                             Constraint::Any
                         },
+                        ip_protocol: Constraint::from(ip_protocol),
                     }
                 }),
                 openvpn_constraints: settings.openvpn_constraints.map(|constraints| {
@@ -1003,6 +1018,17 @@ fn convert_relay_settings(settings: &RelaySettings) -> types::RelaySettings {
 
                 wireguard_constraints: Some(types::WireguardConstraints {
                     port: u32::from(constraints.wireguard_constraints.port.unwrap_or(0)),
+                    ip_protocol: constraints
+                        .wireguard_constraints
+                        .ip_protocol
+                        .option()
+                        .map(|version| match version {
+                            IpVersion::V4 => types::IpVersion::V4,
+                            IpVersion::V6 => types::IpVersion::V6,
+                        })
+                        .map(|version| types::IpVersionConstraint {
+                            protocol: i32::from(version),
+                        }),
                 }),
 
                 openvpn_constraints: Some(types::OpenvpnConstraints {
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 26b855b3b4..dd3d2e12fd 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -30,7 +30,10 @@ use std::{
 };
 use talpid_core::future_retry::{retry_future_with_backoff, ExponentialBackoff, Jittered};
 use talpid_types::{
-    net::{all_of_the_internet, openvpn::ProxySettings, wireguard, TransportProtocol, TunnelType},
+    net::{
+        all_of_the_internet, openvpn::ProxySettings, wireguard, IpVersion, TransportProtocol,
+        TunnelType,
+    },
     ErrorExt,
 };
 use tokio::fs::File;
@@ -464,12 +467,13 @@ impl RelaySelector {
 
         self.pick_random_relay(&matching_relays)
             .and_then(|selected_relay| {
-                info!(
-                    "Selected relay {} at {}",
-                    selected_relay.hostname, selected_relay.ipv4_addr_in
-                );
-                self.get_random_tunnel(&selected_relay, &constraints)
-                    .map(|endpoint| (selected_relay.clone(), endpoint))
+                let endpoint = self.get_random_tunnel(&selected_relay, &constraints);
+                let addr_in = endpoint
+                    .as_ref()
+                    .map(|endpoint| endpoint.to_endpoint().address.ip())
+                    .unwrap_or(IpAddr::from(selected_relay.ipv4_addr_in));
+                info!("Selected relay {} at {}", selected_relay.hostname, addr_in);
+                endpoint.map(|endpoint| (selected_relay.clone(), endpoint))
             })
     }
 
@@ -638,52 +642,38 @@ impl RelaySelector {
         relay: &Relay,
         constraints: &RelayConstraints,
     ) -> Option<MullvadEndpoint> {
-        match constraints.tunnel_protocol {
-            // TODO: Handle Constraint::Any case by selecting from both openvpn and wireguard
-            // tunnels once wireguard is mature enough
-            #[cfg(not(target_os = "android"))]
-            Constraint::Only(TunnelType::OpenVpn) | Constraint::Any => relay
-                .tunnels
-                .openvpn
-                .choose(&mut self.rng)
-                .cloned()
-                .map(|endpoint| endpoint.into_mullvad_endpoint(relay.ipv4_addr_in.into())),
-            Constraint::Only(TunnelType::Wireguard) => relay
+        let mut new_wg_endpoint = || {
+            relay
                 .tunnels
                 .wireguard
                 .choose(&mut self.rng)
                 .cloned()
                 .and_then(|wg_tunnel| {
-                    self.wg_data_to_endpoint(
-                        relay.ipv4_addr_in.into(),
-                        wg_tunnel,
-                        constraints.wireguard_constraints,
-                    )
-                }),
-            #[cfg(target_os = "android")]
-            Constraint::Any => relay
+                    self.wg_data_to_endpoint(relay, wg_tunnel, constraints.wireguard_constraints)
+                })
+        };
+
+        #[cfg(not(target_os = "android"))]
+        match constraints.tunnel_protocol {
+            Constraint::Only(TunnelType::OpenVpn) | Constraint::Any => relay
                 .tunnels
-                .wireguard
+                .openvpn
                 .choose(&mut self.rng)
                 .cloned()
-                .and_then(|wg_tunnel| {
-                    self.wg_data_to_endpoint(
-                        relay.ipv4_addr_in.into(),
-                        wg_tunnel,
-                        WireguardConstraints::default(),
-                    )
-                }),
-            #[cfg(target_os = "android")]
-            Constraint::Only(TunnelType::OpenVpn) => None,
+                .map(|endpoint| endpoint.into_mullvad_endpoint(relay.ipv4_addr_in.into())),
+            Constraint::Only(TunnelType::Wireguard) => new_wg_endpoint(),
         }
+        #[cfg(target_os = "android")]
+        new_wg_endpoint()
     }
 
     fn wg_data_to_endpoint(
         &mut self,
-        host: IpAddr,
+        relay: &Relay,
         data: WireguardEndpointData,
         constraints: WireguardConstraints,
     ) -> Option<MullvadEndpoint> {
+        let host = self.get_address_for_wireguard_relay(relay, constraints)?;
         let port = self.get_port_for_wireguard_relay(&data, constraints)?;
         let peer_config = wireguard::PeerConfig {
             public_key: data.public_key,
@@ -697,6 +687,17 @@ impl RelaySelector {
         })
     }
 
+    fn get_address_for_wireguard_relay(
+        &mut self,
+        relay: &Relay,
+        constraints: WireguardConstraints,
+    ) -> Option<IpAddr> {
+        match constraints.ip_protocol {
+            Constraint::Any | Constraint::Only(IpVersion::V4) => Some(relay.ipv4_addr_in.into()),
+            Constraint::Only(IpVersion::V6) => relay.ipv6_addr_in.map(|addr| addr.into()),
+        }
+    }
+
     fn get_port_for_wireguard_relay(
         &mut self,
         data: &WireguardEndpointData,
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 08b8a708d9..f864665766 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -302,9 +302,20 @@ message OpenvpnConstraints {
 	uint32 port = 1;
 	TransportProtocolConstraint protocol = 2;
 }
+
+enum IpVersion {
+	V4 = 0;
+	V6 = 1;
+}
+
+message IpVersionConstraint {
+	IpVersion protocol = 1;
+}
+
 message WireguardConstraints {
 	// NOTE: optional
 	uint32 port = 1;
+	IpVersionConstraint ip_protocol = 2;
 }
 
 message CustomRelaySettings {
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index c9042c8f49..0fd252b234 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -10,7 +10,7 @@ use crate::{
 use jnix::{FromJava, IntoJava};
 use serde::{Deserialize, Serialize};
 use std::{collections::HashSet, fmt};
-use talpid_types::net::{openvpn::ProxySettings, TransportProtocol, TunnelType};
+use talpid_types::net::{openvpn::ProxySettings, IpVersion, TransportProtocol, TunnelType};
 
 
 pub trait Match<T> {
@@ -434,8 +434,10 @@ impl Match<OpenVpnEndpointData> for OpenVpnConstraints {
 
 /// [`Constraint`]s applicable to WireGuard relay servers.
 #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(default)]
 pub struct WireguardConstraints {
     pub port: Constraint<u16>,
+    pub ip_protocol: Constraint<IpVersion>,
 }
 
 impl fmt::Display for WireguardConstraints {
diff --git a/talpid-types/src/net/mod.rs b/talpid-types/src/net/mod.rs
index 29a871ee07..724c05c4a9 100644
--- a/talpid-types/src/net/mod.rs
+++ b/talpid-types/src/net/mod.rs
@@ -156,6 +156,20 @@ impl fmt::Display for Endpoint {
     }
 }
 
+/// IP protocol version.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum IpVersion {
+    V4,
+    V6,
+}
+
+impl Default for IpVersion {
+    fn default() -> IpVersion {
+        IpVersion::V4
+    }
+}
+
 /// Representation of a transport protocol, either UDP or TCP.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]