diff --git a/rust/connlib/tunnel/src/tests/sim_net.rs b/rust/connlib/tunnel/src/tests/sim_net.rs index 2e2d96a24..716c9cc37 100644 --- a/rust/connlib/tunnel/src/tests/sim_net.rs +++ b/rust/connlib/tunnel/src/tests/sim_net.rs @@ -2,10 +2,8 @@ use crate::tests::buffered_transmits::BufferedTransmits; use crate::tests::strategies::documentation_ip6s; use connlib_model::{ClientId, GatewayId, RelayId}; use firezone_relay::{AddressFamily, IpStack}; -use ip_network::{IpNetwork, Ipv4Network}; +use ip_network::IpNetwork; use ip_network_table::IpNetworkTable; -use itertools::Itertools as _; -use prop::sample; use proptest::prelude::*; use snownet::Transmit; use std::{ @@ -308,13 +306,10 @@ pub(crate) fn dual_ip_stack() -> impl Strategy { /// /// This uses the `TEST-NET-3` (`203.0.113.0/24`) address space reserved for documentation and examples in [RFC5737](https://datatracker.ietf.org/doc/html/rfc5737). pub(crate) fn host_ip4s() -> impl Strategy { - let ips = Ipv4Network::new(Ipv4Addr::new(203, 0, 113, 0), 24) - .unwrap() - .hosts() - .take(100) - .collect_vec(); + const FIRST: Ipv4Addr = Ipv4Addr::new(203, 0, 113, 0); + const LAST: Ipv4Addr = Ipv4Addr::new(203, 0, 113, 255); - sample::select(ips) + (FIRST.to_bits()..=LAST.to_bits()).prop_map(Ipv4Addr::from_bits) } /// A [`Strategy`] of [`Ipv6Addr`]s used for routing packets between hosts within our test. @@ -323,5 +318,5 @@ pub(crate) fn host_ip4s() -> impl Strategy { pub(crate) fn host_ip6s() -> impl Strategy { const HOST_SUBNET: u16 = 0x1010; - documentation_ip6s(HOST_SUBNET, 100) + documentation_ip6s(HOST_SUBNET) } diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs index 721754ffa..34bda4532 100644 --- a/rust/connlib/tunnel/src/tests/strategies.rs +++ b/rust/connlib/tunnel/src/tests/strategies.rs @@ -278,12 +278,10 @@ pub(crate) fn subdomain_records( /// This uses the `TEST-NET-2` (`198.51.100.0/24`) address space reserved for documentation and examples in [RFC5737](https://datatracker.ietf.org/doc/html/rfc5737). /// `TEST-NET-2` only contains 256 addresses which is small enough to generate overlapping IPs for our DNS resources (i.e. two different domains pointing to the same IP). fn dns_resource_ip4s() -> impl Strategy { - let ips = Ipv4Network::new(Ipv4Addr::new(198, 51, 100, 0), 24) - .unwrap() - .hosts() - .collect_vec(); + const FIRST: Ipv4Addr = Ipv4Addr::new(198, 51, 100, 0); + const LAST: Ipv4Addr = Ipv4Addr::new(198, 51, 100, 255); - sample::select(ips) + (FIRST.to_bits()..=LAST.to_bits()).prop_map(Ipv4Addr::from_bits) } /// A [`Strategy`] of [`Ipv6Addr`]s used for the "real" IPs of DNS resources. @@ -292,21 +290,20 @@ fn dns_resource_ip4s() -> impl Strategy { fn dns_resource_ip6s() -> impl Strategy { const DNS_SUBNET: u16 = 0x2020; - documentation_ip6s(DNS_SUBNET, 256) + documentation_ip6s(DNS_SUBNET) } -pub(crate) fn documentation_ip6s(subnet: u16, num_ips: usize) -> impl Strategy { - let ips = Ipv6Network::new_truncate( +pub(crate) fn documentation_ip6s(subnet: u16) -> impl Strategy { + let network = Ipv6Network::new_truncate( Ipv6Addr::new(0x2001, 0xDB80, subnet, subnet, 0, 0, 0, 0), 32, ) - .unwrap() - .subnets_with_prefix(128) - .map(|n| n.network_address()) - .take(num_ips) - .collect_vec(); + .unwrap(); - sample::select(ips) + let first = network.network_address().to_bits(); + let last = network.last_address().to_bits(); + + (first..=last).prop_map(Ipv6Addr::from_bits) } pub(crate) fn system_dns_servers() -> impl Strategy> {