From d7698d50fb1a70e17480750aeff0eb0d4efb63ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B6nnhager?= Date: Mon, 29 Jul 2024 11:38:43 +0200 Subject: [PATCH] Add tests for Shadowsocks and IP overrides --- .../tests/relay_selector.rs | 227 +++++++++++++++++- mullvad-types/src/relay_constraints.rs | 7 + 2 files changed, 227 insertions(+), 7 deletions(-) diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs index 3327e1b4f034..6fcc09735f18 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, @@ -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,212 @@ 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, +} + +// 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 shadowsocks_extra_relay = + GeographicLocationConstraint::hostname("se", "got", "se9-wireguard"); + + 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: shadowsocks_extra_relay.get_hostname().unwrap().to_owned(), + 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, + } +} + /// 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. diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs index ba29f7b3d0a4..b348727685bd 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 for GeographicLocationConstraint {