Add tests for Shadowsocks

3 files changed, 324 insertions, 7 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/helpers.rs b/mullvad-relay-selector/src/relay_selector/helpers.rs
index 86afe512db..fd30c47043 100644
--- a/mullvad-relay-selector/src/relay_selector/helpers.rs
+++ b/mullvad-relay-selector/src/relay_selector/helpers.rs
@@ -209,3 +209,146 @@ pub fn port_in_range(port: u16, port_ranges: &[(u16, u16)]) -> bool {
         .iter()
         .any(|range| (range.0 <= port && port <= range.1))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{get_shadowsocks_obfuscator_inner, port_in_range, SHADOWSOCKS_EXTRA_PORT_RANGES};
+    use mullvad_types::constraints::Constraint;
+    use std::net::IpAddr;
+
+    /// Test whether select ports are available when relay has no extra IPs
+    #[test]
+    fn test_shadowsocks_no_extra_addrs() {
+        const PORT_RANGES: &[(u16, u16)] = &[(100, 200), (1000, 2000)];
+        const WITHIN_RANGE_PORT: u16 = 100;
+        const OUT_OF_RANGE_PORT: u16 = 1;
+        let wg_in_ip: IpAddr = "1.2.3.4".parse().unwrap();
+
+        let selected_addr =
+            get_shadowsocks_obfuscator_inner(wg_in_ip, PORT_RANGES, &[], Constraint::Any)
+                .expect("should find valid port without constraint");
+
+        assert_eq!(selected_addr.ip(), wg_in_ip);
+        assert!(
+            port_in_range(selected_addr.port(), PORT_RANGES),
+            "expected port in port range"
+        );
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            &[],
+            Constraint::Only(WITHIN_RANGE_PORT),
+        )
+        .expect("should find within-range port");
+
+        assert_eq!(selected_addr.ip(), wg_in_ip);
+        assert!(
+            port_in_range(selected_addr.port(), PORT_RANGES),
+            "expected port in port range"
+        );
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            &[],
+            Constraint::Only(OUT_OF_RANGE_PORT),
+        );
+        assert!(
+            selected_addr.is_none(),
+            "expected no relay for port outside range, found {selected_addr:?}"
+        );
+    }
+
+    /// All ports should be available when relay has extra IPs, and only extra IPs should be used
+    #[test]
+    fn test_shadowsocks_extra_addrs() {
+        const PORT_RANGES: &[(u16, u16)] = &[(100, 200), (1000, 2000)];
+        const OUT_OF_RANGE_PORT: u16 = 1;
+        let wg_in_ip: IpAddr = "1.2.3.4".parse().unwrap();
+
+        let extra_in_addrs: &[IpAddr] =
+            &["1.3.3.7".parse().unwrap(), "192.0.2.123".parse().unwrap()];
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            extra_in_addrs,
+            Constraint::Any,
+        )
+        .expect("should find valid port without constraint");
+
+        assert!(
+            extra_in_addrs.contains(&selected_addr.ip()),
+            "expected extra IP to be selected"
+        );
+        assert!(port_in_range(
+            selected_addr.port(),
+            SHADOWSOCKS_EXTRA_PORT_RANGES
+        ));
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            extra_in_addrs,
+            Constraint::Only(OUT_OF_RANGE_PORT),
+        )
+        .expect("expected selected address to be returned");
+        assert!(
+            extra_in_addrs.contains(&selected_addr.ip()),
+            "expected extra IP to be selected, got {selected_addr:?}"
+        );
+        assert_eq!(
+            selected_addr.port(),
+            OUT_OF_RANGE_PORT,
+            "expected selected port, got {selected_addr:?}"
+        );
+    }
+
+    /// Extra addresses that belong to the wrong IP family should be ignored
+    #[test]
+    fn test_shadowsocks_irrelevant_extra_addrs() {
+        const PORT_RANGES: &[(u16, u16)] = &[(100, 200), (1000, 2000)];
+        const IN_RANGE_PORT: u16 = 100;
+        const OUT_OF_RANGE_PORT: u16 = 1;
+        let wg_in_ip: IpAddr = "1.2.3.4".parse().unwrap();
+
+        let extra_in_addrs: &[IpAddr] = &["::2".parse().unwrap()];
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            extra_in_addrs,
+            Constraint::Any,
+        )
+        .expect("should find valid port without constraint");
+
+        assert_eq!(
+            selected_addr.ip(),
+            wg_in_ip,
+            "expected extra IP to be ignored"
+        );
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            extra_in_addrs,
+            Constraint::Only(OUT_OF_RANGE_PORT),
+        );
+        assert!(
+            selected_addr.is_none(),
+            "expected no match for out-of-range port"
+        );
+
+        let selected_addr = get_shadowsocks_obfuscator_inner(
+            wg_in_ip,
+            PORT_RANGES,
+            extra_in_addrs,
+            Constraint::Only(IN_RANGE_PORT),
+        );
+        assert!(
+            selected_addr.is_some(),
+            "expected match for within-range port"
+        );
+    }
+}
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 9069720d2b..38b06666fd 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -1,26 +1,29 @@
 //! Tests for verifying that the relay selector works as expected.
 
 use once_cell::sync::Lazy;
-use std::collections::HashSet;
+use std::{
+    collections::HashSet,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+};
 use talpid_types::net::{
     obfuscation::ObfuscatorConfig,
     wireguard::PublicKey,
-    Endpoint,
+    Endpoint, IpVersion,
     TransportProtocol::{Tcp, Udp},
     TunnelType,
 };
 
 use mullvad_relay_selector::{
     query::{builder::RelayQueryBuilder, BridgeQuery, ObfuscationQuery, OpenVpnRelayQuery},
-    Error, GetRelay, RelaySelector, RuntimeParameters, SelectorConfig, WireguardConfig,
-    RETRY_ORDER,
+    Error, GetRelay, RelaySelector, RuntimeParameters, SelectedObfuscator, SelectorConfig,
+    WireguardConfig, RETRY_ORDER,
 };
 use mullvad_types::{
     constraints::Constraint,
     endpoint::MullvadEndpoint,
     relay_constraints::{
-        BridgeConstraints, BridgeState, GeographicLocationConstraint, Ownership, Providers,
-        TransportPort,
+        BridgeConstraints, BridgeState, GeographicLocationConstraint, LocationConstraint,
+        Ownership, Providers, RelayOverride, TransportPort,
     },
     relay_list::{
         BridgeEndpointData, OpenVpnEndpoint, OpenVpnEndpointData, Relay, RelayEndpointData,
@@ -114,6 +117,7 @@ static RELAYS: Lazy<RelayList> = Lazy::new(|| RelayList {
                     endpoint_data: RelayEndpointData::Bridge,
                     location: None,
                 },
+                SHADOWSOCKS_RELAY.clone(),
             ],
         }],
     }],
@@ -170,6 +174,35 @@ static RELAYS: Lazy<RelayList> = Lazy::new(|| RelayList {
     },
 });
 
+/// A Shadowsocks relay with additional addresses
+static SHADOWSOCKS_RELAY: Lazy<Relay> = Lazy::new(|| Relay {
+    hostname: SHADOWSOCKS_RELAY_LOCATION
+        .get_hostname()
+        .unwrap()
+        .to_owned(),
+    ipv4_addr_in: SHADOWSOCKS_RELAY_IPV4,
+    ipv6_addr_in: Some(SHADOWSOCKS_RELAY_IPV6),
+    include_in_country: true,
+    active: true,
+    owned: true,
+    provider: "provider0".to_string(),
+    weight: 1,
+    endpoint_data: RelayEndpointData::Wireguard(WireguardRelayEndpointData {
+        public_key: PublicKey::from_base64("eaNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=").unwrap(),
+        daita: false,
+        shadowsocks_extra_addr_in: SHADOWSOCKS_RELAY_EXTRA_ADDRS.to_vec(),
+    }),
+    location: None,
+});
+const SHADOWSOCKS_RELAY_IPV4: Ipv4Addr = Ipv4Addr::new(123, 123, 123, 1);
+const SHADOWSOCKS_RELAY_IPV6: Ipv6Addr = Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0, 2);
+const SHADOWSOCKS_RELAY_EXTRA_ADDRS: &[IpAddr; 2] = &[
+    IpAddr::V4(Ipv4Addr::new(123, 123, 123, 2)),
+    IpAddr::V6(Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0, 2)),
+];
+static SHADOWSOCKS_RELAY_LOCATION: Lazy<GeographicLocationConstraint> =
+    Lazy::new(|| GeographicLocationConstraint::hostname("se", "got", "se1337-wireguard"));
+
 // Helper functions
 fn unwrap_relay(get_result: GetRelay) -> Relay {
     match get_result {
@@ -713,6 +746,140 @@ fn test_selecting_any_relay_will_consider_multihop() {
     }
 }
 
+/// Test whether Shadowsocks is always selected as the obfuscation protocol when Shadowsocks is selected.
+#[test]
+fn test_selecting_wireguard_over_shadowsocks() {
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), RELAYS.clone());
+
+    let mut query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    query.wireguard_constraints.use_multihop = Constraint::Only(false);
+
+    let relay = relay_selector.get_relay_by_query(query).unwrap();
+    match relay {
+        GetRelay::Wireguard {
+            obfuscator,
+            inner: WireguardConfig::Singlehop { .. },
+            ..
+        } => {
+            assert!(obfuscator.is_some_and(|obfuscator| matches!(
+                obfuscator.config,
+                ObfuscatorConfig::Shadowsocks { .. }
+            )))
+        }
+        wrong_relay => panic!(
+            "Relay selector should have picked a Wireguard relay with Shadowsocks, instead chose {wrong_relay:?}"
+        ),
+    }
+}
+
+/// Test whether extra Shadowsocks IPs are selected when available
+#[test]
+fn test_selecting_wireguard_over_shadowsocks_extra_ips() {
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), RELAYS.clone());
+
+    let mut query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    query.wireguard_constraints.use_multihop = Constraint::Only(false);
+    query.location = Constraint::Only(LocationConstraint::Location(
+        SHADOWSOCKS_RELAY_LOCATION.clone(),
+    ));
+
+    let relay = relay_selector.get_relay_by_query(query).unwrap();
+    match relay {
+        GetRelay::Wireguard {
+            obfuscator: Some(SelectedObfuscator { config: ObfuscatorConfig::Shadowsocks { endpoint }, .. }),
+            inner: WireguardConfig::Singlehop { .. },
+            ..
+        } => {
+            assert!(SHADOWSOCKS_RELAY_EXTRA_ADDRS.contains(&endpoint.ip()), "{} is not an additional IP", endpoint);
+        }
+        wrong_relay => panic!(
+            "Relay selector should have picked a Wireguard relay with Shadowsocks, instead chose {wrong_relay:?}"
+        ),
+    }
+}
+
+/// Ignore extra IPv4 addresses when overrides are set
+#[test]
+fn test_selecting_wireguard_ignore_extra_ips_override_v4() {
+    const OVERRIDE_IPV4: Ipv4Addr = Ipv4Addr::new(1, 3, 3, 7);
+
+    let config = mullvad_relay_selector::SelectorConfig {
+        relay_overrides: vec![RelayOverride {
+            hostname: SHADOWSOCKS_RELAY_LOCATION
+                .get_hostname()
+                .unwrap()
+                .to_string(),
+            ipv4_addr_in: Some(OVERRIDE_IPV4),
+            ipv6_addr_in: None,
+        }],
+        ..Default::default()
+    };
+
+    let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
+
+    let mut query_v4 = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    query_v4.wireguard_constraints.use_multihop = Constraint::Only(false);
+    query_v4.location = Constraint::Only(LocationConstraint::Location(
+        SHADOWSOCKS_RELAY_LOCATION.clone(),
+    ));
+    query_v4.wireguard_constraints.ip_version = Constraint::Only(IpVersion::V4);
+
+    let relay = relay_selector.get_relay_by_query(query_v4).unwrap();
+    match relay {
+        GetRelay::Wireguard {
+            obfuscator: Some(SelectedObfuscator { config: ObfuscatorConfig::Shadowsocks { endpoint }, .. }),
+            inner: WireguardConfig::Singlehop { .. },
+            ..
+        } => {
+            assert_eq!(endpoint.ip(), IpAddr::from(OVERRIDE_IPV4));
+        }
+        wrong_relay => panic!(
+            "Relay selector should have picked a Wireguard relay with Shadowsocks, instead chose {wrong_relay:?}"
+        ),
+    }
+}
+
+/// Ignore extra IPv6 addresses when overrides are set
+#[test]
+fn test_selecting_wireguard_ignore_extra_ips_override_v6() {
+    const OVERRIDE_IPV6: Ipv6Addr = Ipv6Addr::new(1, 0, 0, 0, 0, 0, 10, 10);
+
+    let config = SelectorConfig {
+        relay_overrides: vec![RelayOverride {
+            hostname: SHADOWSOCKS_RELAY_LOCATION
+                .get_hostname()
+                .unwrap()
+                .to_string(),
+            ipv4_addr_in: None,
+            ipv6_addr_in: Some(OVERRIDE_IPV6),
+        }],
+        ..Default::default()
+    };
+
+    let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
+
+    let mut query_v6 = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    query_v6.wireguard_constraints.use_multihop = Constraint::Only(false);
+    query_v6.location = Constraint::Only(LocationConstraint::Location(
+        SHADOWSOCKS_RELAY_LOCATION.clone(),
+    ));
+    query_v6.wireguard_constraints.ip_version = Constraint::Only(IpVersion::V6);
+
+    let relay = relay_selector.get_relay_by_query(query_v6).unwrap();
+    match relay {
+        GetRelay::Wireguard {
+            obfuscator: Some(SelectedObfuscator { config: ObfuscatorConfig::Shadowsocks { endpoint }, .. }),
+            inner: WireguardConfig::Singlehop { .. },
+            ..
+        } => {
+            assert_eq!(endpoint.ip(), IpAddr::from(OVERRIDE_IPV6));
+        }
+        wrong_relay => panic!(
+            "Relay selector should have picked a Wireguard relay with Shadowsocks, instead chose {wrong_relay:?}"
+        ),
+    }
+}
+
 /// Construct a query for a Wireguard configuration where UDP2TCP obfuscation is selected and
 /// multihop is explicitly turned off. Assert that the relay selector always return an obfuscator
 /// configuration.
@@ -740,7 +907,7 @@ fn test_selecting_wireguard_endpoint_with_udp2tcp_obfuscation() {
     }
 }
 
-/// Construct a query for a Wireguard configuration where UDP2TCP obfuscation is set to "Auto" and
+/// Construct a query for a Wireguard configuration where obfuscation is set to "Auto" and
 /// multihop is explicitly turned off. Assert that the relay selector does *not* return an
 /// obfuscator config.
 ///
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 04e7884519..cbc406eb80 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -198,6 +198,13 @@ impl GeographicLocationConstraint {
     pub fn is_country(&self) -> bool {
         matches!(self, GeographicLocationConstraint::Country(_))
     }
+
+    pub fn get_hostname(&self) -> Option<&Hostname> {
+        match self {
+            GeographicLocationConstraint::Hostname(_, _, hostname) => Some(hostname),
+            _ => None,
+        }
+    }
 }
 
 impl Match<Relay> for GeographicLocationConstraint {