fixup: Simplify shadowsocks tests

dlon · dlon · commit 5e9878954e9d · 2024-08-15T10:57:01.000+02:00
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
@@ -117,6 +117,7 @@ static RELAYS: Lazy<RelayList> = Lazy::new(|| RelayList {
                     endpoint_data: RelayEndpointData::Bridge,
                     location: None,
                 },
+                SHADOWSOCKS_RELAY.clone(),
             ],
         }],
     }],
@@ -173,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 {
@@ -720,10 +750,7 @@ fn test_selecting_any_relay_will_consider_multihop() {
 /// selected.
 #[test]
 fn test_selecting_wireguard_over_shadowsocks() {
-    let RelayListWithExtraShadowsocksIps { relay_list, .. } =
-        get_relay_list_with_extra_shadowsocks_addresses();
-
-    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list);
+    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);
@@ -749,17 +776,13 @@ 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 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_extra_relay));
+    query.location = Constraint::Only(LocationConstraint::Location(
+        SHADOWSOCKS_RELAY_LOCATION.clone(),
+    ));
 
     let relay = relay_selector.get_relay_by_query(query).unwrap();
     match relay {
@@ -768,7 +791,7 @@ fn test_selecting_wireguard_over_shadowsocks_extra_ips() {
             inner: WireguardConfig::Singlehop { .. },
             ..
         } => {
-            assert!(shadowsocks_extra_addr_in.contains(&endpoint.ip()), "{} is not an additional IP", endpoint);
+            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:?}"
@@ -779,28 +802,27 @@ fn test_selecting_wireguard_over_shadowsocks_extra_ips() {
 /// 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 config = mullvad_relay_selector::SelectorConfig {
         relay_overrides: vec![RelayOverride {
-            hostname: shadowsocks_extra_relay.get_hostname().unwrap().to_string(),
+            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, relay_list);
+    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_extra_relay));
+    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();
@@ -821,28 +843,27 @@ fn test_selecting_wireguard_ignore_extra_ips_override_v4() {
 /// 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 config = SelectorConfig {
         relay_overrides: vec![RelayOverride {
-            hostname: shadowsocks_extra_relay.get_hostname().unwrap().to_string(),
+            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, relay_list);
+    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_extra_relay));
+    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();
@@ -860,102 +881,6 @@ fn test_selecting_wireguard_ignore_extra_ips_override_v6() {
     }
 }
 
-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 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.
