From 582e919929ff2e653fa9f0b7878bdbc4dc3759ea Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Wed, 23 Oct 2024 17:42:40 +1100
Subject: [PATCH] test(connlib): don't generate IPs in sentinel range (#7139)

Fixes: #7137.
---
 .../tunnel/proptest-regressions/tests.txt     |  1 +
 rust/connlib/tunnel/src/tests/strategies.rs   | 50 +++++++++++++------
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/rust/connlib/tunnel/proptest-regressions/tests.txt b/rust/connlib/tunnel/proptest-regressions/tests.txt
index 4a24be945..e5f492f7c 100644
--- a/rust/connlib/tunnel/proptest-regressions/tests.txt
+++ b/rust/connlib/tunnel/proptest-regressions/tests.txt
@@ -125,3 +125,4 @@ cc d9793b201ec425bd77f9849ea48e63677014aeb4a91a55be9371b81e644b7a24
 cc 8fcbd19c41f0483d9b81aac2ab7440bb23d7796ef9f6bf346f73f0d633f65baa
 cc 4494e475d22ff9a318d676f10c79f545982b7787d145925c3719fe47e9868acc
 cc bafb7db795d394d1771ef07f4dd36db8ac1333dd852653900480d7ed03307853
+cc 9226be75db567f1d205a36f95cf348eb4aacebbc87c2a4778c52a573b51f0ee2
diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs
index 62d286bb2..77e4564a2 100644
--- a/rust/connlib/tunnel/src/tests/strategies.rs
+++ b/rust/connlib/tunnel/src/tests/strategies.rs
@@ -1,5 +1,8 @@
 use super::{sim_net::Host, sim_relay::ref_relay_host, stub_portal::StubPortal};
-use crate::client::{CidrResource, DnsResource, InternetResource, IPV4_RESOURCES, IPV6_RESOURCES};
+use crate::client::{
+    CidrResource, DnsResource, InternetResource, DNS_SENTINELS_V4, DNS_SENTINELS_V6,
+    IPV4_RESOURCES, IPV6_RESOURCES,
+};
 use crate::proptest::*;
 use crate::{messages::DnsServer, DomainName};
 use connlib_model::{RelayId, Site};
@@ -17,7 +20,7 @@ pub(crate) fn global_dns_records() -> impl Strategy<Value = BTreeMap<DomainName,
 {
     collection::btree_map(
         domain_name(2..4).prop_map(|d| d.parse().unwrap()),
-        collection::btree_set(any::<IpAddr>(), 1..6),
+        collection::btree_set(non_reserved_ip(), 1..6),
         0..5,
     )
 }
@@ -108,13 +111,7 @@ 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(
-        any::<Ipv4Addr>()
-            .prop_filter("must not be in sentinel IP range", |ip| {
-                !crate::client::DNS_SENTINELS_V4.contains(*ip)
-            })
-            .prop_filter("must not be in IPv4 resources range", |ip| {
-                !crate::client::IPV4_RESOURCES.contains(*ip)
-            })
+        non_reserved_ipv4()
             .prop_filter("must be addressable IP", |ip| {
                 !ip.is_unspecified() && !ip.is_multicast() && !ip.is_broadcast()
             })
@@ -122,13 +119,7 @@ pub(crate) fn dns_servers() -> impl Strategy<Value = BTreeSet<SocketAddr>> {
         1..4,
     );
     let ip6_dns_servers = collection::btree_set(
-        any::<Ipv6Addr>()
-            .prop_filter("must not be in sentinel IP range", |ip| {
-                !crate::client::DNS_SENTINELS_V6.contains(*ip)
-            })
-            .prop_filter("must not be in IPv6 resources range", |ip| {
-                !crate::client::IPV6_RESOURCES.contains(*ip)
-            })
+        non_reserved_ipv6()
             .prop_filter("must be addressable IP", |ip| {
                 !ip.is_unspecified() && !ip.is_multicast()
             })
@@ -142,6 +133,33 @@ pub(crate) fn dns_servers() -> impl Strategy<Value = BTreeSet<SocketAddr>> {
     })
 }
 
+fn non_reserved_ip() -> impl Strategy<Value = IpAddr> {
+    prop_oneof![
+        non_reserved_ipv4().prop_map_into(),
+        non_reserved_ipv6().prop_map_into(),
+    ]
+}
+
+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)
+        })
+}
+
+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)
+        })
+}
+
 fn any_site(sites: BTreeSet<Site>) -> impl Strategy<Value = Site> {
     sample::select(Vec::from_iter(sites))
 }