diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs index 51f62c06c..9e4ca709d 100644 --- a/rust/connlib/tunnel/src/tests/reference.rs +++ b/rust/connlib/tunnel/src/tests/reference.rs @@ -57,21 +57,32 @@ impl ReferenceStateMachine for ReferenceState { type Transition = Transition; fn init_state() -> BoxedStrategy { - let client_tunnel_ip4 = tunnel_ip4s().next().unwrap(); - let client_tunnel_ip6 = tunnel_ip6s().next().unwrap(); + stub_portal() + .prop_flat_map(move |portal| { + let gateways = portal.gateways(); + let dns_resource_records = portal.dns_resource_records(); + let client = portal.client(); + let relays = relays(); + let global_dns_records = global_dns_records(); // Start out with a set of global DNS records so we have something to resolve outside of DNS resources. + let drop_direct_client_traffic = any::(); - ( - ref_client_host(Just(client_tunnel_ip4), Just(client_tunnel_ip6)), - gateways_and_portal(), - relays(), - global_dns_records(), // Start out with a set of global DNS records so we have something to resolve outside of DNS resources. - any::(), - ) + ( + client, + gateways, + Just(portal), + dns_resource_records, + relays, + global_dns_records, + drop_direct_client_traffic, + ) + }) .prop_filter_map( "network IPs must be unique", |( c, - (gateways, portal, records), + gateways, + portal, + records, relays, mut global_dns, drop_direct_client_traffic, diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs index 385d1c98f..bf7e7cafc 100644 --- a/rust/connlib/tunnel/src/tests/strategies.rs +++ b/rust/connlib/tunnel/src/tests/strategies.rs @@ -1,9 +1,4 @@ -use super::{ - sim_gateway::{ref_gateway_host, RefGateway}, - sim_net::Host, - sim_relay::ref_relay_host, - stub_portal::StubPortal, -}; +use super::{sim_net::Host, sim_relay::ref_relay_host, stub_portal::StubPortal}; use crate::client::{IPV4_RESOURCES, IPV6_RESOURCES}; use connlib_shared::{ messages::{ @@ -14,8 +9,7 @@ use connlib_shared::{ DnsServer, GatewayId, RelayId, }, proptest::{ - any_ip_network, cidr_resource, dns_resource, domain_label, domain_name, gateway_id, - relay_id, site, + any_ip_network, cidr_resource, dns_resource, domain_name, gateway_id, relay_id, site, }, DomainName, }; @@ -78,43 +72,15 @@ pub(crate) fn packet_source_v6(client: Ipv6Addr) -> impl Strategy. -pub(crate) fn tunnel_ip4s() -> impl Iterator { - Ipv4Network::new(Ipv4Addr::new(100, 64, 0, 0), 11) - .unwrap() - .hosts() -} - -/// An [`Iterator`] over the possible IPv6 addresses of a tunnel interface. -/// -/// See . -pub(crate) fn tunnel_ip6s() -> impl Iterator { - Ipv6Network::new(Ipv6Addr::new(0xfd00, 0x2021, 0x1111, 0, 0, 0, 0, 0), 107) - .unwrap() - .subnets_with_prefix(128) - .map(|n| n.network_address()) -} - pub(crate) fn latency(max: u64) -> impl Strategy { (10..max).prop_map(Duration::from_millis) } -/// A [`Strategy`] for sampling a set of gateways and a corresponding [`StubPortal`] that has a set of [`Site`]s configured with those gateways. +/// A [`Strategy`] for sampling a [`StubPortal`] that is configured with various [`Site`]s and gateways within those sites. /// /// Similar as in production, the portal holds a list of DNS and CIDR resources (those are also sampled from the given sites). /// Via this site mapping, these resources are implicitly assigned to a gateway. -/// -/// Lastly, we also sample a set of DNS records for the DNS resources that we created. -pub(crate) fn gateways_and_portal() -> impl Strategy< - Value = ( - BTreeMap>, - StubPortal, - HashMap>, - ), -> { +pub(crate) fn stub_portal() -> impl Strategy { collection::hash_set(site(), 1..=3) .prop_flat_map(|sites| { let gateway_site = any_site(sites.clone()).prop_map(|s| s.id); @@ -132,78 +98,44 @@ pub(crate) fn gateways_and_portal() -> impl Strategy< ); let internet_resource = internet_resource(any_site(sites)); - let gateways = - collection::hash_map(gateway_id(), (ref_gateway_host(), gateway_site), 1..=3); + // Gateways are unique across sites. + // Generate a map with `GatewayId`s as keys and then flip it into a map of site -> set(gateways). + let gateways_by_site = collection::hash_map(gateway_id(), gateway_site, 1..=3) + .prop_map(|gateway_site| { + let mut gateways_by_site = HashMap::>::default(); + + for (gid, sid) in gateway_site { + gateways_by_site.entry(sid).or_default().insert(gid); + } + + gateways_by_site + }); + let gateway_selector = any::(); ( - gateways, + gateways_by_site, cidr_resources, dns_resources, internet_resource, gateway_selector, ) }) - .prop_flat_map( - |(gateways, cidr_resources, dns_resources, internet_resource, gateway_selector)| { - let (gateways, gateways_by_site) = gateways.into_iter().fold( - ( - BTreeMap::::default(), - HashMap::>::default(), - ), - |(mut gateways, mut sites), (gid, (gateway, site))| { - sites.entry(site).or_default().insert(gid); - gateways.insert(gid, gateway); - - (gateways, sites) - }, - ); - - // For each DNS resource, we need to generate a set of DNS records. - let dns_resource_records = dns_resources - .clone() - .into_iter() - .map(|resource| { - let address = resource.address; - - match address.chars().next().unwrap() { - '*' => subdomain_records( - address.trim_start_matches("*.").to_owned(), - domain_name(1..3), - ) - .boxed(), - '?' => subdomain_records( - address.trim_start_matches("?.").to_owned(), - domain_label(), - ) - .boxed(), - _ => resolved_ips() - .prop_map(move |resolved_ips| { - HashMap::from([(address.parse().unwrap(), resolved_ips)]) - }) - .boxed(), - } - }) - .collect::>() - .prop_map(|records| { - let mut map = HashMap::default(); - - for record in records { - map.extend(record) - } - - map - }); - - let portal = StubPortal::new( + .prop_map( + |( + gateways_by_site, + cidr_resources, + dns_resources, + internet_resource, + gateway_selector, + )| { + StubPortal::new( gateways_by_site, gateway_selector, cidr_resources, dns_resources, internet_resource, - ); - - (Just(gateways), Just(portal), dns_resource_records) + ) }, ) } @@ -265,7 +197,7 @@ fn question_mark_wildcard_dns_resource( }) } -fn resolved_ips() -> impl Strategy> { +pub(crate) fn resolved_ips() -> impl Strategy> { collection::hash_set( prop_oneof![ dns_resource_ip4s().prop_map_into(), @@ -276,7 +208,7 @@ fn resolved_ips() -> impl Strategy> { } /// A strategy for generating a set of DNS records all nested under the provided base domain. -fn subdomain_records( +pub(crate) fn subdomain_records( base: String, subdomains: impl Strategy, ) -> impl Strategy>> { diff --git a/rust/connlib/tunnel/src/tests/stub_portal.rs b/rust/connlib/tunnel/src/tests/stub_portal.rs index 75294c07d..961f87aa4 100644 --- a/rust/connlib/tunnel/src/tests/stub_portal.rs +++ b/rust/connlib/tunnel/src/tests/stub_portal.rs @@ -1,10 +1,24 @@ -use connlib_shared::messages::{client, gateway, GatewayId, ResourceId}; +use super::{ + sim_client::{ref_client_host, RefClient}, + sim_gateway::{ref_gateway_host, RefGateway}, + sim_net::Host, + strategies::{resolved_ips, subdomain_records}, +}; +use connlib_shared::{ + messages::{client, gateway, GatewayId, ResourceId}, + proptest::{domain_label, domain_name}, + DomainName, +}; +use ip_network::{Ipv4Network, Ipv6Network}; use itertools::Itertools; -use proptest::sample::Selector; +use proptest::{ + sample::Selector, + strategy::{Just, Strategy}, +}; use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeMap, HashMap, HashSet}, iter, - net::IpAddr, + net::{IpAddr, Ipv4Addr, Ipv6Addr}, }; /// Stub implementation of the portal. @@ -172,4 +186,85 @@ impl StubPortal { Some(gid) } + + pub(crate) fn gateways(&self) -> impl Strategy>> { + self.gateways_by_site + .values() + .flatten() + .map(|gid| (Just(*gid), ref_gateway_host())) // Map each ID to a strategy that samples a gateway. + .collect::>() // A `Vec` implements `Strategy>` + .prop_map(BTreeMap::from_iter) + } + + pub(crate) fn client(&self) -> impl Strategy> { + let client_tunnel_ip4 = tunnel_ip4s().next().unwrap(); + let client_tunnel_ip6 = tunnel_ip6s().next().unwrap(); + + ref_client_host(Just(client_tunnel_ip4), Just(client_tunnel_ip6)) + } + + pub(crate) fn dns_resource_records( + &self, + ) -> impl Strategy>> { + self.dns_resources + .values() + .map(|resource| { + let address = resource.address.clone(); + + match address.chars().next().unwrap() { + '*' => subdomain_records( + address.trim_start_matches("*.").to_owned(), + domain_name(1..3), + ) + .boxed(), + '?' => subdomain_records( + address.trim_start_matches("?.").to_owned(), + domain_label(), + ) + .boxed(), + _ => resolved_ips() + .prop_map(move |resolved_ips| { + HashMap::from([(address.parse().unwrap(), resolved_ips)]) + }) + .boxed(), + } + }) + .collect::>() + .prop_map(|records| { + let mut map = HashMap::default(); + + for record in records { + map.extend(record) + } + + map + }) + } +} + +const IPV4_TUNNEL: Ipv4Network = match Ipv4Network::new(Ipv4Addr::new(100, 64, 0, 0), 11) { + Ok(n) => n, + Err(_) => unreachable!(), +}; +const IPV6_TUNNEL: Ipv6Network = + match Ipv6Network::new(Ipv6Addr::new(0xfd00, 0x2021, 0x1111, 0, 0, 0, 0, 0), 107) { + Ok(n) => n, + Err(_) => unreachable!(), + }; + +/// An [`Iterator`] over the possible IPv4 addresses of a tunnel interface. +/// +/// We use the CG-NAT range for IPv4. +/// See . +fn tunnel_ip4s() -> impl Iterator { + IPV4_TUNNEL.hosts() +} + +/// An [`Iterator`] over the possible IPv6 addresses of a tunnel interface. +/// +/// See . +fn tunnel_ip6s() -> impl Iterator { + IPV6_TUNNEL + .subnets_with_prefix(128) + .map(|n| n.network_address()) }