From fd337dd4653d4eb292572e50afdbb2271c25a5be Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 28 Nov 2024 20:45:02 +0000
Subject: [PATCH] test: reduce number of local rejects for generating IPs
 (#7401)

When generating random input data in property-based tests, we have to
ensure that the data conforms to certain criteria. For example, IP
addresses must not be multicast or unspecified addresses and they must
not be within our reserved IP ranges.

Currently, we ensure this using "filtering" which is a pretty poor
technique [0]. To improve on this, we refactor the generation of IPs to
automatically exclude all IPs within certain ranges. On very big
test-runs (i.e. > 30000 test cases), too many local rejections lead to
the test suite being aborted early.

[0]:
https://proptest-rs.github.io/proptest/proptest/tutorial/filtering.html
---
 rust/connlib/tunnel/src/tests/strategies.rs | 59 ++++++++++++---------
 1 file changed, 35 insertions(+), 24 deletions(-)

diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs
index a5418f60e..721754ffa 100644
--- a/rust/connlib/tunnel/src/tests/strategies.rs
+++ b/rust/connlib/tunnel/src/tests/strategies.rs
@@ -135,19 +135,11 @@ pub(crate) fn relays(
 /// We make sure to always have at least 1 IPv4 and 1 IPv6 DNS server.
 pub(crate) fn dns_servers() -> impl Strategy<Value = BTreeSet<SocketAddr>> {
     let ip4_dns_servers = collection::btree_set(
-        non_reserved_ipv4()
-            .prop_filter("must be addressable IP", |ip| {
-                !ip.is_unspecified() && !ip.is_multicast() && !ip.is_broadcast()
-            })
-            .prop_map(|ip| SocketAddr::from((ip, 53))),
+        non_reserved_ipv4().prop_map(|ip| SocketAddr::from((ip, 53))),
         1..4,
     );
     let ip6_dns_servers = collection::btree_set(
-        non_reserved_ipv6()
-            .prop_filter("must be addressable IP", |ip| {
-                !ip.is_unspecified() && !ip.is_multicast()
-            })
-            .prop_map(|ip| SocketAddr::from((ip, 53))),
+        non_reserved_ipv6().prop_map(|ip| SocketAddr::from((ip, 53))),
         1..4,
     );
 
@@ -165,23 +157,42 @@ pub(crate) fn non_reserved_ip() -> impl Strategy<Value = IpAddr> {
 }
 
 fn non_reserved_ipv4() -> impl Strategy<Value = Ipv4Addr> {
-    any::<Ipv4Addr>()
-        .prop_filter("must not be in sentinel IP range", |ip| {
-            !DNS_SENTINELS_V4.contains(*ip)
-        })
-        .prop_filter("must not be in IPv4 resources range", |ip| {
-            !IPV4_RESOURCES.contains(*ip)
-        })
+    let undesired_ranges = [
+        Ipv4Network::new(Ipv4Addr::BROADCAST, 32).unwrap(),
+        Ipv4Network::new(Ipv4Addr::UNSPECIFIED, 32).unwrap(),
+        Ipv4Network::new(Ipv4Addr::new(224, 0, 0, 0), 4).unwrap(), // Multicast
+        DNS_SENTINELS_V4,
+        IPV4_RESOURCES,
+    ];
+
+    any::<Ipv4Addr>().prop_map(move |mut ip| {
+        while let Some(range) = undesired_ranges.iter().find(|range| range.contains(ip)) {
+            ip = Ipv4Addr::from(u32::from(range.broadcast_address()).wrapping_add(1));
+        }
+
+        debug_assert!(undesired_ranges.iter().all(|range| !range.contains(ip)));
+
+        ip
+    })
 }
 
 fn non_reserved_ipv6() -> impl Strategy<Value = Ipv6Addr> {
-    any::<Ipv6Addr>()
-        .prop_filter("must not be in sentinel IP range", |ip| {
-            !DNS_SENTINELS_V6.contains(*ip)
-        })
-        .prop_filter("must not be in IPv6 resources range", |ip| {
-            !IPV6_RESOURCES.contains(*ip)
-        })
+    let undesired_ranges = [
+        Ipv6Network::new(Ipv6Addr::UNSPECIFIED, 32).unwrap(),
+        DNS_SENTINELS_V6,
+        IPV6_RESOURCES,
+        Ipv6Network::new(Ipv6Addr::new(0xff00, 0, 0, 0, 0, 0, 0, 0), 8).unwrap(), // Multicast
+    ];
+
+    any::<Ipv6Addr>().prop_map(move |mut ip| {
+        while let Some(range) = undesired_ranges.iter().find(|range| range.contains(ip)) {
+            ip = Ipv6Addr::from(u128::from(range.last_address()).wrapping_add(1));
+        }
+
+        debug_assert!(undesired_ranges.iter().all(|range| !range.contains(ip)));
+
+        ip
+    })
 }
 
 fn any_site(sites: BTreeSet<Site>) -> impl Strategy<Value = Site> {