2 files changed, 78 insertions, 26 deletions

diff --git a/mullvad-daemon/src/relays/matcher.rs b/mullvad-daemon/src/relays/matcher.rs
index 08924a78e9..69e98cfeff 100644
--- a/mullvad-daemon/src/relays/matcher.rs
+++ b/mullvad-daemon/src/relays/matcher.rs
@@ -258,17 +258,20 @@ impl From<WireguardConstraints> for WireguardMatcher {
 
 impl Match<WireguardEndpointData> for WireguardMatcher {
     fn matches(&self, endpoint: &WireguardEndpointData) -> bool {
-        match self
-            .port
-            .as_ref()
-            .map(|port| port.port)
-            .unwrap_or(Constraint::Any)
-        {
+        match self.port {
             Constraint::Any => true,
-            Constraint::Only(port) => endpoint
-                .port_ranges
-                .iter()
-                .any(|range| (port >= range.0 && port <= range.1)),
+            Constraint::Only(TransportPort { port, protocol }) => {
+                if protocol != endpoint.protocol {
+                    return false;
+                }
+                match port {
+                    Constraint::Any => true,
+                    Constraint::Only(port) => endpoint
+                        .port_ranges
+                        .iter()
+                        .any(|range| (port >= range.0 && port <= range.1)),
+                }
+            }
         }
     }
 }
diff --git a/mullvad-daemon/src/relays/mod.rs b/mullvad-daemon/src/relays/mod.rs
index 1fead05f77..984b2280ae 100644
--- a/mullvad-daemon/src/relays/mod.rs
+++ b/mullvad-daemon/src/relays/mod.rs
@@ -1350,33 +1350,82 @@ mod test {
         );
     }
 
+    const WIREGUARD_MULTIHOP_CONSTRAINTS: RelayConstraints = RelayConstraints {
+        location: Constraint::Any,
+        providers: Constraint::Any,
+        wireguard_constraints: WireguardConstraints {
+            use_multihop: true,
+            port: Constraint::Any,
+            ip_version: Constraint::Any,
+            entry_location: Constraint::Any,
+        },
+        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
+        openvpn_constraints: OpenVpnConstraints {
+            port: Constraint::Any,
+        },
+    };
+
     #[test]
     fn test_selecting_wireguard_location_will_consider_multihop() {
-        let wireguard_specific_location = LocationConstraint::Hostname(
-            "se".to_string(),
-            "got".to_string(),
-            "se9-wireguard".to_string(),
-        );
+        let relay_selector = new_relay_selector();
+
+        let result = relay_selector.get_tunnel_endpoint(&WIREGUARD_MULTIHOP_CONSTRAINTS, BridgeState::Off, 0, true)
+
+            .expect("Failed to get relay when tunnel constraints are set to Any and retrying the selection");
+
+        assert!(result.entry_relay.is_some());
+        let endpoint = result.endpoint.unwrap_wireguard();
+        assert!(matches!(endpoint.peer.protocol, TransportProtocol::Udp));
+        assert!(matches!(
+            endpoint.exit_peer.as_ref().unwrap().protocol,
+            TransportProtocol::Udp
+        ));
+    }
+
+    #[test]
+    fn test_selecting_wg_multihop_tcp() {
+        let mut relay_constraints = WIREGUARD_MULTIHOP_CONSTRAINTS.clone();
+        relay_constraints.wireguard_constraints.port = Constraint::Only(TransportPort {
+            port: Constraint::Any,
+            protocol: TransportProtocol::Tcp,
+        });
 
+        let relay_selector = new_relay_selector();
+
+        let result = relay_selector
+            .get_tunnel_endpoint(&relay_constraints, BridgeState::Off, 0, true)
+            .expect("Failed to get WireGuard TCP multihop relay");
+
+        assert!(result.entry_relay.is_some());
+        let endpoint = result.endpoint.unwrap_wireguard();
+        assert!(matches!(endpoint.peer.protocol, TransportProtocol::Tcp));
+        assert!(matches!(
+            endpoint.exit_peer.as_ref().unwrap().protocol,
+            TransportProtocol::Udp
+        ));
+    }
+
+    #[test]
+    fn test_selecting_wg_tcp() {
         let relay_constraints = RelayConstraints {
-            location: Constraint::Only(wireguard_specific_location),
             wireguard_constraints: WireguardConstraints {
-                use_multihop: true,
+                port: Constraint::Only(TransportPort {
+                    port: Constraint::Any,
+                    protocol: TransportProtocol::Tcp,
+                }),
                 ..WireguardConstraints::default()
             },
-            // This has to be explicit otherwise Android will chose WireGuard when default
-            // constructing.
-            tunnel_protocol: Constraint::Any,
+            tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
             ..RelayConstraints::default()
         };
 
         let relay_selector = new_relay_selector();
 
-        let result = relay_selector.get_tunnel_endpoint(&relay_constraints, BridgeState::Off, 0, true)
-            .expect("Failed to get relay when tunnel constraints are set to Any and retrying the selection");
-        assert!(
-            matches!(result.endpoint, MullvadEndpoint::Wireguard(_))
-                && result.entry_relay.is_some()
-        );
+        let result = relay_selector
+            .get_tunnel_endpoint(&relay_constraints, BridgeState::Off, 0, true)
+            .expect("Failed to get WireGuard TCP relay");
+        let endpoint = result.endpoint.unwrap_wireguard();
+        assert!(matches!(endpoint.peer.protocol, TransportProtocol::Tcp));
+        assert!(endpoint.exit_peer.is_none());
     }
 }