Merge branch 'fix-wg-relay-matching'

2 files changed, 50 insertions, 8 deletions

diff --git a/mullvad-daemon/src/relays/matcher.rs b/mullvad-daemon/src/relays/matcher.rs
index 69e98cfeff..fb708d05fb 100644
--- a/mullvad-daemon/src/relays/matcher.rs
+++ b/mullvad-daemon/src/relays/matcher.rs
@@ -8,7 +8,7 @@ use mullvad_types::{
 };
 use rand::{seq::SliceRandom, Rng};
 use std::net::{IpAddr, SocketAddr};
-use talpid_types::net::{all_of_the_internet, wireguard, IpVersion, TransportProtocol, TunnelType};
+use talpid_types::net::{all_of_the_internet, wireguard, IpVersion, TunnelType};
 
 #[derive(Clone)]
 pub struct RelayMatcher<T: TunnelMatcher> {
@@ -183,10 +183,7 @@ impl WireguardMatcher {
             public_key: data.public_key,
             endpoint: SocketAddr::new(host, port),
             allowed_ips: all_of_the_internet(),
-            protocol: self
-                .port
-                .map(|port| port.protocol)
-                .unwrap_or(TransportProtocol::Udp),
+            protocol: data.protocol,
         };
         Some(MullvadEndpoint::Wireguard(MullvadWireguardEndpoint {
             peer: peer_config,
@@ -306,11 +303,17 @@ impl TunnelMatcher for WireguardMatcher {
     }
 
     fn mullvad_endpoint(&self, relay: &Relay) -> Option<MullvadEndpoint> {
-        relay
+        let valid_relays = relay
             .tunnels
             .wireguard
+            .iter()
+            .filter(|tunnel| match self.port {
+                Constraint::Any => true,
+                Constraint::Only(port) => port.protocol == tunnel.protocol,
+            })
+            .collect::<Vec<_>>();
+        valid_relays
             .choose(&mut rand::thread_rng())
-            .cloned()
-            .and_then(|wg_tunnel| self.wg_data_to_endpoint(relay, wg_tunnel))
+            .and_then(|wg_tunnel| self.wg_data_to_endpoint(relay, (*wg_tunnel).clone()))
     }
 }
diff --git a/mullvad-daemon/src/relays/mod.rs b/mullvad-daemon/src/relays/mod.rs
index 984b2280ae..0c2aec38fe 100644
--- a/mullvad-daemon/src/relays/mod.rs
+++ b/mullvad-daemon/src/relays/mod.rs
@@ -389,6 +389,10 @@ impl RelaySelector {
         }
 
         entry_relay_matcher.location = wireguard_constraints.entry_location.clone();
+        entry_relay_matcher.tunnel.port = entry_relay_matcher
+            .tunnel
+            .port
+            .or(Self::preferred_wireguard_port(retry_attempt));
         self.get_wireguard_multi_hop_endpoint(entry_relay_matcher, location.clone())
     }
 
@@ -1428,4 +1432,39 @@ mod test {
         assert!(matches!(endpoint.peer.protocol, TransportProtocol::Tcp));
         assert!(endpoint.exit_peer.is_none());
     }
+
+    #[test]
+    fn test_selecting_wg_multihop_ports() {
+        let mut relay_constraints = WIREGUARD_MULTIHOP_CONSTRAINTS.clone();
+        let relay_selector = new_relay_selector();
+
+        const INVALID_UDP_PORTS: [u16; 2] = [80, 443];
+        for attempt in 0..1000 {
+            let result = relay_selector
+                .get_tunnel_endpoint(&relay_constraints, BridgeState::Off, attempt, true)
+                .expect("Failed to get WireGuard TCP multihop relay");
+            assert!(!INVALID_UDP_PORTS.contains(&result.endpoint.to_endpoint().address.port()));
+            assert_eq!(
+                result.endpoint.unwrap_wireguard().peer.protocol,
+                TransportProtocol::Udp
+            );
+        }
+
+        relay_constraints.wireguard_constraints.port = Constraint::Only(TransportPort {
+            port: Constraint::Any,
+            protocol: TransportProtocol::Tcp,
+        });
+
+        const VALID_TCP_PORTS: [u16; 3] = [80, 443, 5001];
+        for attempt in 0..1000 {
+            let result = relay_selector
+                .get_tunnel_endpoint(&relay_constraints, BridgeState::Off, attempt, true)
+                .expect("Failed to get WireGuard TCP multihop relay");
+            assert!(VALID_TCP_PORTS.contains(&result.endpoint.to_endpoint().address.port()));
+            assert_eq!(
+                result.endpoint.unwrap_wireguard().peer.protocol,
+                TransportProtocol::Tcp
+            );
+        }
+    }
 }