Add tests for Shadowsocks and IP overrides

mullvad-relay-selector/tests/relay_selector.rs

@@ -1,26 +1,29 @@
 //! Tests for verifying that the relay selector works as expected.
 
+use std::{
+    collections::HashSet,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+};
 use once_cell::sync::Lazy;
-use std::collections::HashSet;
 use talpid_types::net::{
     obfuscation::ObfuscatorConfig,
     wireguard::PublicKey,
-    Endpoint,
+    Endpoint, IpVersion,
     TransportProtocol::{Tcp, Udp},
     TunnelType,
 };
 
 use mullvad_relay_selector::{
     query::{builder::RelayQueryBuilder, BridgeQuery, 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,
-        SelectedObfuscation, TransportPort,
+        BridgeConstraints, BridgeState, GeographicLocationConstraint, LocationConstraint,
+        Ownership, Providers, RelayOverride, SelectedObfuscation, TransportPort,
     },
     relay_list::{
         BridgeEndpointData, OpenVpnEndpoint, OpenVpnEndpointData, Relay, RelayEndpointData,
@@ -719,7 +722,11 @@ fn test_selecting_any_relay_will_consider_multihop() {
 /// selected.
 #[test]
 fn test_selecting_wireguard_over_shadowsocks() {
-    let relay_selector = default_relay_selector();
+    let RelayListWithExtraShadowsocksIps { relay_list, .. } =
+        get_relay_list_with_extra_shadowsocks_addresses();
+
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list);
+
     let mut query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
     query.wireguard_constraints.use_multihop = Constraint::Only(false);
 
@@ -741,6 +748,214 @@ fn test_selecting_wireguard_over_shadowsocks() {
     }
 }
 
+/// Test whether extra Shadowsocks IPs are selected when available
+#[test]
+fn test_selecting_wireguard_over_shadowsocks_extra_ips() {
+    let RelayListWithExtraShadowsocksIps {
+        relay_list,
+        shadowsocks_extra_relay,
+        shadowsocks_extra_addr_in,
+    } = get_relay_list_with_extra_shadowsocks_addresses();
+
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list);
+
+    let mut query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+    query.wireguard_constraints.use_multihop = Constraint::Only(false);
+    query.location = Constraint::Only(LocationConstraint::Location(shadowsocks_extra_relay));
+
+    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_extra_addr_in.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() {
+    let RelayListWithExtraShadowsocksIps {
+        relay_list,
+        shadowsocks_extra_relay,
+        ..
+    } = get_relay_list_with_extra_shadowsocks_addresses();
+
+    const OVERRIDE_IPV4: Ipv4Addr = Ipv4Addr::new(1, 3, 3, 7);
+
+    let mut config = SelectorConfig::default();
+    config.relay_overrides = vec![RelayOverride {
+        hostname: shadowsocks_extra_relay.get_hostname().unwrap().to_string(),
+        ipv4_addr_in: Some(OVERRIDE_IPV4),
+        ipv6_addr_in: None,
+    }];
+
+    let relay_selector = RelaySelector::from_list(config, relay_list);
+
+    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_extra_relay));
+    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() {
+    let RelayListWithExtraShadowsocksIps {
+        relay_list,
+        shadowsocks_extra_relay,
+        ..
+    } = get_relay_list_with_extra_shadowsocks_addresses();
+
+    const OVERRIDE_IPV6: Ipv6Addr = Ipv6Addr::new(1, 0, 0, 0, 0, 0, 10, 10);
+
+    let mut config = SelectorConfig::default();
+    config.relay_overrides = vec![RelayOverride {
+        hostname: shadowsocks_extra_relay.get_hostname().unwrap().to_string(),
+        ipv4_addr_in: None,
+        ipv6_addr_in: Some(OVERRIDE_IPV6),
+    }];
+
+    let relay_selector = RelaySelector::from_list(config, relay_list);
+
+    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_extra_relay));
+    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:?}"
+        ),
+    }
+}
+
+struct RelayListWithExtraShadowsocksIps {
+    relay_list: RelayList,
+    shadowsocks_extra_relay: GeographicLocationConstraint,
+    shadowsocks_extra_addr_in: Vec<IpAddr>,
+}
+
+// Return a relay list with a single relay
+fn get_relay_list_with_extra_shadowsocks_addresses() -> RelayListWithExtraShadowsocksIps {
+    static SHADOWSOCKS_EXTRA_IPS_RELAY_NORMAL_IPV4: Ipv4Addr = Ipv4Addr::new(123, 123, 123, 1);
+    static SHADOWSOCKS_EXTRA_IPS_RELAY_NORMAL_IPV6: Ipv6Addr =
+        Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0, 2);
+    let shadowsocks_extra_addr_in = vec![
+        Ipv4Addr::new(123, 123, 123, 2).into(),
+        Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0, 2).into(),
+    ];
+
+    let relay_list = RelayList {
+        etag: None,
+        countries: vec![RelayListCountry {
+            name: "Sweden".to_string(),
+            code: "se".to_string(),
+            cities: vec![RelayListCity {
+                name: "Gothenburg".to_string(),
+                code: "got".to_string(),
+                latitude: 57.70887,
+                longitude: 11.97456,
+                relays: vec![
+                    Relay {
+                        hostname: "se9-wireguard".to_string(),
+                        ipv4_addr_in: SHADOWSOCKS_EXTRA_IPS_RELAY_NORMAL_IPV4,
+                        ipv6_addr_in: Some(SHADOWSOCKS_EXTRA_IPS_RELAY_NORMAL_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(
+                                "BLNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=",
+                            )
+                            .unwrap(),
+                            daita: false,
+                            shadowsocks_extra_addr_in: shadowsocks_extra_addr_in.clone(),
+                        }),
+                        location: None,
+                    },
+                    Relay {
+                        hostname: "se10-wireguard".to_string(),
+                        ipv4_addr_in: Ipv4Addr::new(123, 123, 123, 3),
+                        ipv6_addr_in: Some(Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0, 3)),
+                        include_in_country: true,
+                        active: true,
+                        owned: false,
+                        provider: "provider1".to_string(),
+                        weight: 1,
+                        endpoint_data: RelayEndpointData::Wireguard(WireguardRelayEndpointData {
+                            public_key: PublicKey::from_base64(
+                                "BLNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=",
+                            )
+                            .unwrap(),
+                            daita: false,
+                            shadowsocks_extra_addr_in: vec![],
+                        }),
+                        location: None,
+                    },
+                ],
+            }],
+        }],
+        openvpn: OpenVpnEndpointData { ports: vec![] },
+        bridge: BridgeEndpointData {
+            shadowsocks: vec![],
+        },
+        wireguard: WireguardEndpointData {
+            port_ranges: vec![
+                (53, 53),
+                (443, 443),
+                (4000, 33433),
+                (33565, 51820),
+                (52000, 60000),
+            ],
+            ipv4_gateway: "10.64.0.1".parse().unwrap(),
+            ipv6_gateway: "fc00:bbbb:bbbb:bb01::1".parse().unwrap(),
+            udp2tcp_ports: vec![],
+            shadowsocks_port_ranges: vec![(100, 200), (1000, 2000)],
+        },
+    };
+
+    RelayListWithExtraShadowsocksIps {
+        relay_list,
+        shadowsocks_extra_addr_in,
+        shadowsocks_extra_relay: GeographicLocationConstraint::hostname(
+            "se",
+            "got",
+            "se9-wireguard",
+        ),
+    }
+}
+
 /// 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.

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 {