Select random IP and port if there are additional Shadowsocks addresses

Co-authored-by: Sebastian Holmin <sebastian.holmin@mullvad.net>

5 files changed, 129 insertions, 47 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/detailer.rs b/mullvad-relay-selector/src/relay_selector/detailer.rs
index ca28de36ec..536ea50f26 100644
--- a/mullvad-relay-selector/src/relay_selector/detailer.rs
+++ b/mullvad-relay-selector/src/relay_selector/detailer.rs
@@ -42,7 +42,7 @@ pub enum Error {
     #[error("The selected relay does not support IPv6")]
     NoIPv6(Box<Relay>),
     #[error("Failed to select port")]
-    PortSelectionError(#[source] super::helpers::Error),
+    PortSelectionError,
 }
 
 /// Constructs a [`MullvadWireguardEndpoint`] with details for how to connect to a Wireguard relay.
@@ -159,22 +159,36 @@ fn get_address_for_wireguard_relay(
     query: &WireguardRelayQuery,
     relay: &Relay,
 ) -> Result<IpAddr, Error> {
-    match query.ip_version {
-        Constraint::Any | Constraint::Only(IpVersion::V4) => Ok(relay.ipv4_addr_in.into()),
-        Constraint::Only(IpVersion::V6) => relay
+    match resolve_ip_version(query.ip_version) {
+        IpVersion::V4 => Ok(relay.ipv4_addr_in.into()),
+        IpVersion::V6 => relay
             .ipv6_addr_in
             .map(|addr| addr.into())
             .ok_or(Error::NoIPv6(Box::new(relay.clone()))),
     }
 }
 
+pub fn resolve_ip_version(ip_version: Constraint<IpVersion>) -> IpVersion {
+    match ip_version {
+        Constraint::Any | Constraint::Only(IpVersion::V4) => IpVersion::V4,
+        Constraint::Only(IpVersion::V6) => IpVersion::V6,
+    }
+}
+
 /// Try to pick a valid Wireguard port.
 fn get_port_for_wireguard_relay(
     query: &WireguardRelayQuery,
     data: &WireguardEndpointData,
 ) -> Result<u16, Error> {
-    super::helpers::select_random_port(query.port, &data.port_ranges)
-        .map_err(Error::PortSelectionError)
+    if let Constraint::Only(port) = query.port {
+        if super::helpers::port_in_range(port, &data.port_ranges) {
+            Ok(port)
+        } else {
+            Err(Error::PortSelectionError)
+        }
+    } else {
+        super::helpers::select_random_port(&data.port_ranges).map_err(|_| Error::PortSelectionError)
+    }
 }
 
 /// Read the [`PublicKey`] of a relay. This will only succeed if [relay][`Relay`] is a
diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index 7964bef71f..86afe512db 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -1,6 +1,6 @@
 //! This module contains various helper functions for the relay selector implementation.
 
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 
 use mullvad_types::{
     constraints::Constraint,
@@ -8,17 +8,22 @@ use mullvad_types::{
     relay_constraints::{ShadowsocksSettings, Udp2TcpObfuscationSettings},
     relay_list::Relay,
 };
-use rand::{seq::SliceRandom, thread_rng, Rng};
+use rand::{
+    seq::{IteratorRandom, SliceRandom},
+    thread_rng, Rng,
+};
 use talpid_types::net::obfuscation::ObfuscatorConfig;
 
 use crate::SelectedObfuscator;
 
+/// Port ranges available for WireGuard relays that have extra IPs for Shadowsocks.
+/// For relays that have no additional IPs, only ports provided by the relay list are available.
+const SHADOWSOCKS_EXTRA_PORT_RANGES: &[(u16, u16)] = &[(1, u16::MAX)];
+
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
     #[error("Port selection algorithm is broken")]
     PortSelectionAlgorithm,
-    #[error("Invalid port argument: port {0} is not in any valid Wireguard port range")]
-    PortNotInRange(u16),
 }
 
 /// Picks a relay using [pick_random_relay_weighted], using the `weight` member of each relay
@@ -99,17 +104,69 @@ pub fn get_udp2tcp_obfuscator_port(
 
 pub fn get_shadowsocks_obfuscator(
     settings: &ShadowsocksSettings,
-    port_ranges: &[(u16, u16)],
+    non_extra_port_ranges: &[(u16, u16)],
     relay: Relay,
     endpoint: &MullvadWireguardEndpoint,
 ) -> Option<SelectedObfuscator> {
-    let port = select_random_port(settings.port, port_ranges).ok()?;
+    let port = settings.port;
+    let extra_addrs = match &relay.endpoint_data {
+        mullvad_types::relay_list::RelayEndpointData::Wireguard(wg) => {
+            &wg.shadowsocks_extra_addr_in
+        }
+        _ => panic!("expected wireguard relay"),
+    };
+
+    let endpoint = get_shadowsocks_obfuscator_inner(
+        endpoint.peer.endpoint.ip(),
+        non_extra_port_ranges,
+        extra_addrs,
+        port,
+    )?;
 
-    let config = ObfuscatorConfig::Shadowsocks {
-        endpoint: SocketAddr::new(endpoint.peer.endpoint.ip(), port),
+    Some(SelectedObfuscator {
+        config: ObfuscatorConfig::Shadowsocks { endpoint },
+        relay,
+    })
+}
+
+/// Return an obfuscation config for the wireguard server at `wg_in_addr` or one of `extra_in_addrs`
+/// (unless empty). `wg_in_addr_port_ranges` contains all valid ports for `wg_in_addr`, and
+/// `SHADOWSOCKS_EXTRA_PORT_RANGES` contains valid ports for `extra_in_addrs`.
+fn get_shadowsocks_obfuscator_inner(
+    wg_in_addr: IpAddr,
+    wg_in_addr_port_ranges: &[(u16, u16)],
+    extra_in_addrs: &[IpAddr],
+    desired_port: Constraint<u16>,
+) -> Option<SocketAddr> {
+    // Filter out addresses for the wrong address family
+    let extra_in_addrs: Vec<_> = extra_in_addrs
+        .iter()
+        .filter(|addr| addr.is_ipv4() == wg_in_addr.is_ipv4())
+        .copied()
+        .collect();
+
+    let in_ip = extra_in_addrs
+        .iter()
+        .choose(&mut rand::thread_rng())
+        .copied()
+        .unwrap_or(wg_in_addr);
+
+    let port_ranges = if extra_in_addrs.is_empty() {
+        wg_in_addr_port_ranges
+    } else {
+        SHADOWSOCKS_EXTRA_PORT_RANGES
     };
 
-    Some(SelectedObfuscator { config, relay })
+    let selected_port = match desired_port {
+        // Selected a specific, in-range port
+        Constraint::Only(port) if super::helpers::port_in_range(port, port_ranges) => Some(port),
+        // Selected a specific, out-of-range port
+        Constraint::Only(_port) => None,
+        // Selected no specific port
+        Constraint::Any => super::helpers::select_random_port(port_ranges).ok(),
+    }?;
+
+    Some(SocketAddr::from((in_ip, selected_port)))
 }
 
 /// Selects a random port number from a list of provided port ranges.
@@ -127,24 +184,8 @@ pub fn get_shadowsocks_obfuscator(
 ///
 /// # Panic
 /// - If port ranges contains no ports, this function panics.
-pub fn select_random_port(port: Constraint<u16>, port_ranges: &[(u16, u16)]) -> Result<u16, Error> {
-    match port {
-        Constraint::Any => select_random_port_inner(port_ranges),
-        Constraint::Only(port) => {
-            if port_ranges
-                .iter()
-                .any(|range| (range.0 <= port && port <= range.1))
-            {
-                Ok(port)
-            } else {
-                Err(Error::PortNotInRange(port))
-            }
-        }
-    }
-}
-
-fn select_random_port_inner(port_ranges: &[(u16, u16)]) -> Result<u16, Error> {
-    let get_port_amount = |range: &(u16, u16)| -> u64 { (1 + range.1 - range.0) as u64 };
+pub fn select_random_port(port_ranges: &[(u16, u16)]) -> Result<u16, Error> {
+    let get_port_amount = |range: &(u16, u16)| -> u64 { 1 + range.1 as u64 - range.0 as u64 };
     let port_amount: u64 = port_ranges.iter().map(get_port_amount).sum();
 
     if port_amount < 1 {
@@ -162,3 +203,9 @@ fn select_random_port_inner(port_ranges: &[(u16, u16)]) -> Result<u16, Error> {
     }
     Err(Error::PortSelectionAlgorithm)
 }
+
+pub fn port_in_range(port: u16, port_ranges: &[(u16, u16)]) -> bool {
+    port_ranges
+        .iter()
+        .any(|range| (range.0 <= port && port <= range.1))
+}
diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index 1bd96a0cc4..3d879d8bb6 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/mullvad-relay-selector/src/relay_selector/matcher.rs
@@ -6,15 +6,15 @@ use mullvad_types::{
     custom_list::CustomListsSettings,
     relay_constraints::{
         GeographicLocationConstraint, InternalBridgeConstraints, LocationConstraint, Ownership,
-        Providers, SelectedObfuscation, ShadowsocksSettings,
+        Providers, ShadowsocksSettings,
     },
     relay_list::{Relay, RelayEndpointData, WireguardRelayEndpointData},
 };
-use talpid_types::net::TunnelType;
+use talpid_types::net::{IpVersion, TunnelType};
 
 use super::{
     parsed_relays::ParsedRelays,
-    query::{RelayQuery, WireguardRelayQuery},
+    query::{ObfuscationQuery, RelayQuery, WireguardRelayQuery},
 };
 
 /// Filter a list of relays and their endpoints based on constraints.
@@ -136,13 +136,14 @@ fn filter_on_obfuscation(
     relay_list: &ParsedRelays,
     relay: &Relay,
 ) -> bool {
-    match query.obfuscation {
+    match &query.obfuscation {
         // Shadowsocks has relay-specific constraints
-        SelectedObfuscation::Shadowsocks => {
+        ObfuscationQuery::Shadowsocks(settings) => {
             let wg_data = &relay_list.parsed_list().wireguard;
             filter_on_shadowsocks(
                 &wg_data.shadowsocks_port_ranges,
-                &query.shadowsocks_port,
+                &query.ip_version,
+                settings,
                 relay,
             )
         }
@@ -155,20 +156,31 @@ fn filter_on_obfuscation(
 /// Returns whether `relay` satisfies the Shadowsocks filter posed by `port`.
 fn filter_on_shadowsocks(
     port_ranges: &[(u16, u16)],
-    settings: &Constraint<ShadowsocksSettings>,
+    ip_version: &Constraint<IpVersion>,
+    settings: &ShadowsocksSettings,
     relay: &Relay,
 ) -> bool {
-    match (&settings, &relay.endpoint_data) {
+    let ip_version = super::detailer::resolve_ip_version(*ip_version);
+
+    match (settings, &relay.endpoint_data) {
         // If Shadowsocks is specifically asked for, we must check if the specific relay supports our port.
         // If there are extra addresses, then all ports are available, so we do not need to do this.
         (
-            Constraint::Only(ShadowsocksSettings {
+            ShadowsocksSettings {
                 port: Constraint::Only(desired_port),
-            }),
+            },
             RelayEndpointData::Wireguard(wg_data),
-        ) if wg_data.shadowsocks_extra_addr_in.is_empty() => port_ranges
-            .iter()
-            .any(|(begin, end)| (*begin..=*end).contains(&desired_port)),
+        ) => {
+            let filtered_extra_addrs = wg_data
+                .shadowsocks_extra_addr_in
+                .iter()
+                .find(|&&addr| IpVersion::from(addr) == ip_version);
+
+            filtered_extra_addrs.is_some()
+                || port_ranges
+                    .iter()
+                    .any(|(begin, end)| (*begin..=*end).contains(desired_port))
+        }
 
         // Otherwise, any relay works.
         _ => true,
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 62f8d31ab4..853431f155 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -809,7 +809,7 @@ impl RelaySelector {
                     .wireguard
                     .shadowsocks_port_ranges;
                 let obfuscation = helpers::get_shadowsocks_obfuscator(
-                    &settings,
+                    settings,
                     port_ranges,
                     obfuscator_relay,
                     endpoint,
diff --git a/talpid-types/src/net/mod.rs b/talpid-types/src/net/mod.rs
index 3d6961f50b..5875f33482 100644
--- a/talpid-types/src/net/mod.rs
+++ b/talpid-types/src/net/mod.rs
@@ -474,6 +474,15 @@ pub enum IpVersion {
     V6,
 }
 
+impl From<IpAddr> for IpVersion {
+    fn from(value: IpAddr) -> Self {
+        match value {
+            IpAddr::V4(_) => IpVersion::V4,
+            IpAddr::V6(_) => IpVersion::V6,
+        }
+    }
+}
+
 impl fmt::Display for IpVersion {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         match *self {