refactor(connlib): simplify sampling of initial state (#6194)

Instead of having one giant, composed strategy, we introduce a dedicated
`stub_portal` strategy. That one samples what is defined in the portal
in production: sites, gateways and resources.

Based on a sampled portal, we can then sample gateways, a client and DNS
records for our resources.
This commit is contained in:
Thomas Eizinger
2024-08-07 07:07:39 +01:00
committed by GitHub
parent 423d70854b
commit 376900ca4e
3 changed files with 150 additions and 112 deletions

View File

@@ -57,21 +57,32 @@ impl ReferenceStateMachine for ReferenceState {
type Transition = Transition;
fn init_state() -> BoxedStrategy<Self::State> {
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::<bool>();
(
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::<bool>(),
)
(
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,

View File

@@ -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<Value = Ipv6Ad
]
}
/// An [`Iterator`] over the possible IPv4 addresses of a tunnel interface.
///
/// We use the CG-NAT range for IPv4.
/// See <https://github.com/firezone/firezone/blob/81dfa90f38299595e14ce9e022d1ee919909f124/elixir/apps/domain/lib/domain/network.ex#L7>.
pub(crate) fn tunnel_ip4s() -> impl Iterator<Item = Ipv4Addr> {
Ipv4Network::new(Ipv4Addr::new(100, 64, 0, 0), 11)
.unwrap()
.hosts()
}
/// An [`Iterator`] over the possible IPv6 addresses of a tunnel interface.
///
/// See <https://github.com/firezone/firezone/blob/81dfa90f38299595e14ce9e022d1ee919909f124/elixir/apps/domain/lib/domain/network.ex#L8>.
pub(crate) fn tunnel_ip6s() -> impl Iterator<Item = Ipv6Addr> {
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<Value = Duration> {
(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<GatewayId, Host<RefGateway>>,
StubPortal,
HashMap<DomainName, HashSet<IpAddr>>,
),
> {
pub(crate) fn stub_portal() -> impl Strategy<Value = StubPortal> {
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::<SiteId, HashSet<GatewayId>>::default();
for (gid, sid) in gateway_site {
gateways_by_site.entry(sid).or_default().insert(gid);
}
gateways_by_site
});
let gateway_selector = any::<sample::Selector>();
(
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::<GatewayId, _>::default(),
HashMap::<SiteId, HashSet<GatewayId>>::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::<Vec<_>>()
.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<Value = HashSet<IpAddr>> {
pub(crate) fn resolved_ips() -> impl Strategy<Value = HashSet<IpAddr>> {
collection::hash_set(
prop_oneof![
dns_resource_ip4s().prop_map_into(),
@@ -276,7 +208,7 @@ fn resolved_ips() -> impl Strategy<Value = HashSet<IpAddr>> {
}
/// 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<Value = String>,
) -> impl Strategy<Value = HashMap<DomainName, HashSet<IpAddr>>> {

View File

@@ -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<Value = BTreeMap<GatewayId, Host<RefGateway>>> {
self.gateways_by_site
.values()
.flatten()
.map(|gid| (Just(*gid), ref_gateway_host())) // Map each ID to a strategy that samples a gateway.
.collect::<Vec<_>>() // A `Vec<Strategy>` implements `Strategy<Value = Vec<_>>`
.prop_map(BTreeMap::from_iter)
}
pub(crate) fn client(&self) -> impl Strategy<Value = Host<RefClient>> {
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<Value = HashMap<DomainName, HashSet<IpAddr>>> {
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::<Vec<_>>()
.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 <https://github.com/firezone/firezone/blob/81dfa90f38299595e14ce9e022d1ee919909f124/elixir/apps/domain/lib/domain/network.ex#L7>.
fn tunnel_ip4s() -> impl Iterator<Item = Ipv4Addr> {
IPV4_TUNNEL.hosts()
}
/// An [`Iterator`] over the possible IPv6 addresses of a tunnel interface.
///
/// See <https://github.com/firezone/firezone/blob/81dfa90f38299595e14ce9e022d1ee919909f124/elixir/apps/domain/lib/domain/network.ex#L8>.
fn tunnel_ip6s() -> impl Iterator<Item = Ipv6Addr> {
IPV6_TUNNEL
.subnets_with_prefix(128)
.map(|n| n.network_address())
}