Add transport protocol constraint for WireGuard

5 files changed, 58 insertions, 19 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 245eda9e91..d119617802 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -167,6 +167,14 @@ impl Command for Relay {
                                             .required(true)
                                     )
                                     .arg(
+                                        clap::Arg::with_name("transport protocol")
+                                            .help("Transport protocol. If TCP is selected, traffic is \
+                                                   sent over TCP using a udp-over-tcp proxy")
+                                            .long("protocol")
+                                            .default_value("any")
+                                            .possible_values(&["any", "udp", "tcp"]),
+                                    )
+                                    .arg(
                                         clap::Arg::with_name("ip version")
                                             .long("ipv")
                                             .default_value("any")
@@ -528,6 +536,7 @@ impl Relay {
 
     async fn set_wireguard_constraints(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let port = parse_port_constraint(matches.value_of("port").unwrap())?;
+        let protocol = parse_protocol_constraint(matches.value_of("transport protocol").unwrap());
         let ip_version = parse_ip_version_constraint(matches.value_of("ip version").unwrap());
         let entry_location =
             parse_entry_location_constraint(matches.values_of("entry location").unwrap());
@@ -537,6 +546,11 @@ impl Relay {
                 NormalRelaySettingsUpdate {
                     wireguard_constraints: Some(WireguardConstraints {
                         port: port.unwrap_or(0) as u32,
+                        protocol: protocol
+                            .option()
+                            .map(|protocol| TransportProtocolConstraint {
+                                protocol: protocol as i32,
+                            }),
                         ip_version: ip_version.option().map(|protocol| IpVersionConstraint {
                             protocol: protocol as i32,
                         }),
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 54b56cde5f..0e1111e74e 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -54,6 +54,7 @@ const EXPONENTIAL_BACKOFF_FACTOR: u32 = 8;
 const DEFAULT_WIREGUARD_PORT: u16 = 51820;
 const WIREGUARD_EXIT_CONSTRAINTS: WireguardConstraints = WireguardConstraints {
     port: Constraint::Only(DEFAULT_WIREGUARD_PORT),
+    protocol: Constraint::Only(TransportProtocol::Udp),
     ip_version: Constraint::Only(IpVersion::V4),
     entry_location: None,
 };
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 33c4294db7..1d3b40258b 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -332,10 +332,10 @@ message IpVersionConstraint {
 }
 
 message WireguardConstraints {
-	// NOTE: optional
 	uint32 port = 1;
-	IpVersionConstraint ip_version = 2;
-	RelayLocation entry_location = 3;
+	TransportProtocolConstraint protocol = 2;
+	IpVersionConstraint ip_version = 3;
+	RelayLocation entry_location = 4;
 }
 
 message CustomRelaySettings {
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index 11547887a4..79a450e2c3 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -492,6 +492,13 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for RelaySettings {
 
                     wireguard_constraints: Some(WireguardConstraints {
                         port: u32::from(constraints.wireguard_constraints.port.unwrap_or(0)),
+                        protocol: constraints
+                            .wireguard_constraints
+                            .protocol
+                            .as_ref()
+                            .option()
+                            .map(|protocol| TransportProtocol::from(*protocol))
+                            .map(TransportProtocolConstraint::from),
                         ip_version: constraints
                             .wireguard_constraints
                             .ip_version
@@ -870,21 +877,36 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                     None
                 };
 
-                let transport_protocol = if let Some(ref constraints) = settings.openvpn_constraints
-                {
-                    match &constraints.protocol {
-                        Some(constraint) => Some(
-                            TransportProtocol::from_i32(constraint.protocol)
-                                .ok_or(FromProtobufTypeError::InvalidArgument(
-                                    "invalid transport protocol",
-                                ))?
-                                .into(),
-                        ),
-                        None => None,
-                    }
-                } else {
-                    None
-                };
+                let openvpn_transport_protocol =
+                    if let Some(ref constraints) = settings.openvpn_constraints {
+                        match &constraints.protocol {
+                            Some(constraint) => Some(
+                                TransportProtocol::from_i32(constraint.protocol)
+                                    .ok_or(FromProtobufTypeError::InvalidArgument(
+                                        "invalid transport protocol",
+                                    ))?
+                                    .into(),
+                            ),
+                            None => None,
+                        }
+                    } else {
+                        None
+                    };
+                let wireguard_transport_protocol =
+                    if let Some(ref constraints) = settings.wireguard_constraints {
+                        match &constraints.protocol {
+                            Some(constraint) => Some(
+                                TransportProtocol::from_i32(constraint.protocol)
+                                    .ok_or(FromProtobufTypeError::InvalidArgument(
+                                        "invalid transport protocol",
+                                    ))?
+                                    .into(),
+                            ),
+                            None => None,
+                        }
+                    } else {
+                        None
+                    };
 
                 let providers = if let Some(ref provider_update) = settings.providers {
                     if !provider_update.providers.is_empty() {
@@ -933,6 +955,7 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                                 } else {
                                     Constraint::Any
                                 },
+                                protocol: Constraint::from(wireguard_transport_protocol),
                                 ip_version: Constraint::from(ip_version),
                                 entry_location: constraints.entry_location.map(
                                     Constraint::<
@@ -948,7 +971,7 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                                 } else {
                                     Constraint::Any
                                 },
-                                protocol: Constraint::from(transport_protocol),
+                                protocol: Constraint::from(openvpn_transport_protocol),
                             }
                         }),
                     },
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 53b618a2bf..1549034e2a 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -477,6 +477,7 @@ impl Match<OpenVpnEndpointData> for OpenVpnConstraints {
 #[serde(default)]
 pub struct WireguardConstraints {
     pub port: Constraint<u16>,
+    pub protocol: Constraint<TransportProtocol>,
     pub ip_version: Constraint<IpVersion>,
     pub entry_location: Option<Constraint<LocationConstraint>>,
 }