github-personal/firezone
1
0
Fork 0
mirror of https://github.com/outbackdingo/firezone.git synced 2026-01-27 10:18:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(connlib): introduce dynamic number of gateways to tunnel_test (#5823)

Browse Source 
Currently, `tunnel_test` exercises a lot of code paths within connlib
already by adding & removing resources, roaming the client and sending
ICMP packets. Yet, it does all of this with just a single gateway
whereas in production, we are very likely using more than one gateway.

To capture these other code-paths, we now sample between 1 and 3
gateways and randomly assign the added resources to one of them, which
makes us hit the codepaths that select between different gateways.

Most importantly, the reference implementation has barely any knowledge
about those individual connections. Instead, it is implementation in
terms of connectivity to resources.
This commit is contained in:
Thomas Eizinger
2024-07-12 09:42:46 +10:00
committed by GitHub
parent 960ce80680
commit d95193be7d
15 changed files with 403 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
rust/Cargo.lock generated
View File

@@ -3416,7 +3416,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19"
dependencies = [
"cfg-if",
"windows-targets 0.52.5",
"windows-targets 0.48.5",
]
[[package]]
@@ -4825,9 +4825,8 @@ dependencies = [
[[package]]
name = "proptest"
version = "1.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31b476131c3c86cb68032fdc5cb6d5a1045e3e42d96b69fa599fd77701e1f5bf"
version = "1.5.0"
source = "git+https://github.com/thomaseizinger/proptest?branch=fix/always-check-acceptable-current-state#26c036b5ca832726d7f0c8438a750c8efa2f0f8b"
dependencies = [
"bit-set",
"bit-vec",
@@ -4846,8 +4845,7 @@ dependencies = [
[[package]]
name = "proptest-state-machine"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28278d6a11102264b0c569c33dbe5286ba00d2dd6d96ff2a94296e0e5b3d1e04"
source = "git+https://github.com/thomaseizinger/proptest?branch=fix/always-check-acceptable-current-state#26c036b5ca832726d7f0c8438a750c8efa2f0f8b"
dependencies = [
"proptest",
]

2
rust/Cargo.toml
View File

@@ -69,6 +69,8 @@ rustdoc.private-intra-doc-links = "allow" # We don't publish any of our docs but
boringtun = { git = "https://github.com/cloudflare/boringtun", branch = "master" }
str0m = { git = "https://github.com/firezone/str0m", branch = "main" }
ip_network_table = { git = "https://github.com/edmonds/ip_network_table", branch = "some-useful-traits" } # For `Debug` and `Clone`
proptest = { git = "https://github.com/thomaseizinger/proptest", branch = "fix/always-check-acceptable-current-state" }
proptest-state-machine = { git = "https://github.com/thomaseizinger/proptest", branch = "fix/always-check-acceptable-current-state" }
[profile.release]
strip = true

2
rust/connlib/shared/Cargo.toml
View File

@@ -32,7 +32,7 @@ domain = { workspace = true }
libc = "0.2"
phoenix-channel = { workspace = true }
ip-packet = { workspace = true }
proptest = { version = "1.4.0", optional = true }
proptest = { version = "1", optional = true }
itertools = "0.13"
# Needed for Android logging until tracing is working

2
rust/connlib/tunnel/Cargo.toml
View File

@@ -30,7 +30,7 @@ socket2 = { version = "0.5" }
snownet = { workspace = true }
quinn-udp = { git = "https://github.com/quinn-rs/quinn", branch = "main"}
hex = "0.4.3"
proptest = { version = "1.4.0", optional = true }
proptest = { version = "1", optional = true }
ip-packet = { workspace = true }
rangemap = "1.5.1"
anyhow = "1.0"

2
rust/connlib/tunnel/proptest-regressions/tests.txt
View File

File diff suppressed because one or more lines are too long

102
rust/connlib/tunnel/src/tests/assertions.rs
View File

@@ -3,7 +3,7 @@ use super::{
sim_gateway::SimGateway,
};
use crate::tests::reference::ResourceDst;
use connlib_shared::DomainName;
use connlib_shared::{messages::GatewayId, DomainName};
use ip_packet::IpPacket;
use pretty_assertions::assert_eq;
use std::{
@@ -20,11 +20,15 @@ use std::{
pub(crate) fn assert_icmp_packets_properties(
ref_client: &RefClient,
sim_client: &SimClient,
sim_gateway: &SimGateway,
sim_gateways: HashMap<GatewayId, &SimGateway>,
global_dns_records: &BTreeMap<DomainName, HashSet<IpAddr>>,
) {
let unexpected_icmp_replies = find_unexpected_entries(
&ref_client.expected_icmp_handshakes,
&ref_client
.expected_icmp_handshakes
.values()
.flatten()
.collect(),
&sim_client.received_icmp_replies,
|(_, seq_a, id_a), (seq_b, id_b)| seq_a == seq_b && id_a == id_b,
);
@@ -34,56 +38,66 @@ pub(crate) fn assert_icmp_packets_properties(
"Unexpected ICMP replies on client"
);
assert_eq!(
ref_client.expected_icmp_handshakes.len(),
sim_gateway.received_icmp_requests.len(),
"Unexpected ICMP requests on gateway"
);
for (id, expected_icmp_handshakes) in ref_client.expected_icmp_handshakes.iter() {
let gateway = sim_gateways.get(id).unwrap();
tracing::info!(target: "assertions", "✅ Performed the expected {} ICMP handshakes", sim_gateway.received_icmp_requests.len());
assert_eq!(
expected_icmp_handshakes.len(),
gateway.received_icmp_requests.len(),
"Unexpected ICMP requests on gateway {id}"
);
tracing::info!(target: "assertions", "✅ Performed the expected {} ICMP handshakes with gateway {id}", expected_icmp_handshakes.len());
}
let mut mapping = HashMap::new();
for ((resource_dst, seq, identifier), gateway_received_request) in ref_client
.expected_icmp_handshakes
.iter()
.zip(sim_gateway.received_icmp_requests.iter())
{
let _guard = tracing::info_span!(target: "assertions", "icmp", %seq, %identifier).entered();
// Assert properties of the individual ICMP handshakes per gateway.
// Due to connlib's implementation of NAT64, we cannot match the packets sent by the client to the packets arriving at the resource by port or ICMP identifier.
// Thus, we rely on the _order_ here which is why the packets are indexed by gateway in the `RefClient`.
for (gateway, expected_icmp_handshakes) in &ref_client.expected_icmp_handshakes {
let received_icmp_requests = &sim_gateways.get(gateway).unwrap().received_icmp_requests;
let client_sent_request = &sim_client
.sent_icmp_requests
.get(&(*seq, *identifier))
.expect("to have ICMP request on client");
let client_received_reply = &sim_client
.received_icmp_replies
.get(&(*seq, *identifier))
.expect("to have ICMP reply on client");
for ((resource_dst, seq, identifier), gateway_received_request) in
expected_icmp_handshakes.iter().zip(received_icmp_requests)
{
let _guard =
tracing::info_span!(target: "assertions", "icmp", %seq, %identifier).entered();
assert_correct_src_and_dst_ips(client_sent_request, client_received_reply);
let client_sent_request = &sim_client
.sent_icmp_requests
.get(&(*seq, *identifier))
.expect("to have ICMP request on client");
let client_received_reply = &sim_client
.received_icmp_replies
.get(&(*seq, *identifier))
.expect("to have ICMP reply on client");
assert_eq!(
gateway_received_request.source(),
ref_client.tunnel_ip_for(gateway_received_request.source()),
"ICMP request on gateway to originate from client"
);
assert_correct_src_and_dst_ips(client_sent_request, client_received_reply);
match resource_dst {
ResourceDst::Cidr(resource_dst) => {
assert_destination_is_cdir_resource(gateway_received_request, resource_dst)
}
ResourceDst::Dns(domain) => {
assert_destination_is_dns_resource(
gateway_received_request,
global_dns_records,
domain,
);
assert_eq!(
gateway_received_request.source(),
ref_client.tunnel_ip_for(gateway_received_request.source()),
"ICMP request on gateway to originate from client"
);
assert_proxy_ip_mapping_is_stable(
client_sent_request,
gateway_received_request,
&mut mapping,
)
match resource_dst {
ResourceDst::Cidr(resource_dst) => {
assert_destination_is_cdir_resource(gateway_received_request, resource_dst)
}
ResourceDst::Dns(domain) => {
assert_destination_is_dns_resource(
gateway_received_request,
global_dns_records,
domain,
);
assert_proxy_ip_mapping_is_stable(
client_sent_request,
gateway_received_request,
&mut mapping,
)
}
}
}
}

134
rust/connlib/tunnel/src/tests/reference.rs
View File

@@ -3,14 +3,18 @@ use super::{
strategies::*, transition::*,
};
use chrono::{DateTime, Utc};
use connlib_shared::{messages::RelayId, proptest::*, DomainName, StaticSecret};
use connlib_shared::{
messages::{GatewayId, RelayId},
proptest::*,
DomainName, StaticSecret,
};
use hickory_proto::rr::RecordType;
use prop::collection;
use proptest::{prelude::*, sample};
use proptest_state_machine::ReferenceStateMachine;
use std::{
collections::{BTreeMap, HashMap, HashSet},
fmt,
fmt, iter,
net::IpAddr,
time::{Duration, Instant},
};
@@ -23,7 +27,7 @@ pub(crate) struct ReferenceState {
pub(crate) now: Instant,
pub(crate) utc_now: DateTime<Utc>,
pub(crate) client: Host<RefClient>,
pub(crate) gateway: Host<RefGateway>,
pub(crate) gateways: HashMap<GatewayId, Host<RefGateway>>,
pub(crate) relays: HashMap<RelayId, Host<u64>>,
/// All IP addresses a domain resolves to in our test.
@@ -49,29 +53,31 @@ impl ReferenceStateMachine for ReferenceState {
type State = Self;
type Transition = Transition;
fn init_state() -> proptest::prelude::BoxedStrategy<Self::State> {
fn init_state() -> BoxedStrategy<Self::State> {
let mut tunnel_ip4s = tunnel_ip4s();
let mut tunnel_ip6s = tunnel_ip6s();
(
ref_client_host(&mut tunnel_ip4s, &mut tunnel_ip6s),
ref_gateway_host(),
collection::hash_map(relay_id(), relay_prototype(), 2),
collection::hash_map(gateway_id(), ref_gateway_host(), 1..=3),
collection::hash_map(relay_id(), relay_prototype(), 1..=2),
global_dns_records(), // Start out with a set of global DNS records so we have something to resolve outside of DNS resources.
Just(Instant::now()),
Just(Utc::now()),
)
.prop_filter_map(
"network IPs must be unique",
|(c, g, relays, global_dns, now, utc_now)| {
|(c, gateways, relays, global_dns, now, utc_now)| {
let mut routing_table = RoutingTable::default();
if !routing_table.add_host(c.inner().id, &c) {
return None;
}
if !routing_table.add_host(g.inner().id, &g) {
return None;
};
for (id, gateway) in &gateways {
if !routing_table.add_host(*id, gateway) {
return None;
};
}
for (id, relay) in &relays {
if !routing_table.add_host(*id, relay) {
@@ -79,19 +85,27 @@ impl ReferenceStateMachine for ReferenceState {
};
}
Some((c, g, relays, global_dns, now, utc_now, routing_table))
Some((c, gateways, relays, global_dns, now, utc_now, routing_table))
},
)
.prop_filter(
"client and gateway priv key must be different",
|(c, g, _, _, _, _, _)| c.inner().key != g.inner().key,
"private keys must be unique",
|(c, gateways, _, _, _, _, _)| {
let different_keys = gateways
.iter()
.map(|(_, g)| g.inner().key)
.chain(iter::once(c.inner().key))
.collect::<HashSet<_>>();
different_keys.len() == gateways.len() + 1
},
)
.prop_map(
|(client, gateway, relays, global_dns_records, now, utc_now, network)| Self {
|(client, gateways, relays, global_dns_records, now, utc_now, network)| Self {
now,
utc_now,
client,
gateway,
gateways,
relays,
global_dns_records,
network,
@@ -104,7 +118,7 @@ impl ReferenceStateMachine for ReferenceState {
///
/// This is invoked by proptest repeatedly to explore further state transitions.
/// Here, we should only generate [`Transition`]s that make sense for the current state.
fn transitions(state: &Self::State) -> proptest::prelude::BoxedStrategy<Self::Transition> {
fn transitions(state: &Self::State) -> BoxedStrategy<Self::Transition> {
CompositeStrategy::default()
.with(
1,
@@ -120,14 +134,19 @@ impl ReferenceStateMachine for ReferenceState {
upstream_dns_servers()
.prop_map(|servers| Transition::UpdateUpstreamDnsServers { servers }),
)
.with(1, cidr_resource(8).prop_map(Transition::AddCidrResource))
.with(
1,
(cidr_resource(8), sample::select(state.all_gateways())).prop_map(
|(resource, gateway)| Transition::AddCidrResource { resource, gateway },
),
)
.with(1, roam_client())
.with(
1,
prop_oneof![
non_wildcard_dns_resource(),
star_wildcard_dns_resource(),
question_mark_wildcard_dns_resource(),
non_wildcard_dns_resource(sample::select(state.all_gateways())),
star_wildcard_dns_resource(sample::select(state.all_gateways())),
question_mark_wildcard_dns_resource(sample::select(state.all_gateways())),
],
)
.with_if_not_empty(
@@ -229,10 +248,13 @@ impl ReferenceStateMachine for ReferenceState {
/// Here is where we implement the "expected" logic.
fn apply(mut state: Self::State, transition: &Self::Transition) -> Self::State {
match transition {
Transition::AddCidrResource(r) => {
state
.client
.exec_mut(|client| client.cidr_resources.insert(r.address, r.clone()));
Transition::AddCidrResource { resource, gateway } => {
state.client.exec_mut(|client| {
client
.cidr_resources
.insert(resource.address, resource.clone());
client.gateways_by_resource.insert(resource.id, *gateway);
});
}
Transition::RemoveResource(id) => {
state
@@ -248,8 +270,12 @@ impl ReferenceStateMachine for ReferenceState {
Transition::AddDnsResource {
resource: new_resource,
records,
gateway,
} => {
let existing_resource = state.client.exec_mut(|client| {
client
.gateways_by_resource
.insert(new_resource.id, *gateway);
client
.dns_resources
.insert(new_resource.id, new_resource.clone())
@@ -371,28 +397,38 @@ impl ReferenceStateMachine for ReferenceState {
/// Any additional checks on whether a particular [`Transition`] can be applied to a certain state.
fn preconditions(state: &Self::State, transition: &Self::Transition) -> bool {
match transition {
Transition::AddCidrResource(r) => {
Transition::AddCidrResource { resource, gateway } => {
let Some(gateway) = state.gateways.get(gateway) else {
return false;
};
// TODO: PRODUCTION CODE DOES NOT HANDLE THIS!
if r.address.is_ipv6() && state.gateway.ip6.is_none() {
if resource.address.is_ipv6() && gateway.ip6.is_none() {
return false;
}
if r.address.is_ipv4() && state.gateway.ip4.is_none() {
if resource.address.is_ipv4() && gateway.ip4.is_none() {
return false;
}
// TODO: PRODUCTION CODE DOES NOT HANDLE THIS!
for dns_resolved_ip in state.global_dns_records.values().flat_map(|ip| ip.iter()) {
// If the CIDR resource overlaps with an IP that a DNS record resolved to, we have problems ...
if r.address.contains(*dns_resolved_ip) {
if resource.address.contains(*dns_resolved_ip) {
return false;
}
}
true
}
Transition::AddDnsResource { records, .. } => {
Transition::AddDnsResource {
records, gateway, ..
} => {
if !state.gateways.contains_key(gateway) {
return false;
};
// TODO: Should we allow adding a DNS resource if we don't have an DNS resolvers?
// TODO: For these tests, we assign the resolved IP of a DNS resource as part of this transition.
@@ -439,8 +475,18 @@ impl ReferenceStateMachine for ReferenceState {
is_valid_icmp_packet && !is_cidr_resource
}
Transition::SendICMPPacketToCidrResource {
seq, identifier, ..
} => state.client.inner().is_valid_icmp_packet(seq, identifier),
seq,
identifier,
dst,
..
} => {
let ref_client = state.client.inner();
ref_client.is_valid_icmp_packet(seq, identifier)
&& ref_client
.gateway_by_cidr_resource_ip(*dst)
.is_some_and(|g| state.gateways.contains_key(&g))
}
Transition::SendICMPPacketToDnsResource {
seq,
identifier,
@@ -448,16 +494,16 @@ impl ReferenceStateMachine for ReferenceState {
src,
..
} => {
state.client.inner().is_valid_icmp_packet(seq, identifier)
&& state
.client
.inner()
.dns_records
.get(dst)
.is_some_and(|r| match src {
IpAddr::V4(_) => r.contains(&RecordType::A),
IpAddr::V6(_) => r.contains(&RecordType::AAAA),
})
let ref_client = state.client.inner();
ref_client.is_valid_icmp_packet(seq, identifier)
&& ref_client.dns_records.get(dst).is_some_and(|r| match src {
IpAddr::V4(_) => r.contains(&RecordType::A),
IpAddr::V6(_) => r.contains(&RecordType::AAAA),
})
&& ref_client
.gateway_by_domain_name(dst)
.is_some_and(|g| state.gateways.contains_key(&g))
}
Transition::UpdateSystemDnsServers { servers } => {
// TODO: PRODUCTION CODE DOES NOT HANDLE THIS!
@@ -521,6 +567,10 @@ impl ReferenceState {
)
.collect()
}
fn all_gateways(&self) -> Vec<GatewayId> {
self.gateways.keys().copied().collect()
}
}
fn matches_domain(resource_address: &str, domain: &DomainName) -> bool {
@@ -539,7 +589,7 @@ pub(crate) fn private_key() -> impl Strategy<Value = PrivateKey> {
any::<[u8; 32]>().prop_map(PrivateKey)
}
#[derive(Clone, Copy, PartialEq)]
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct PrivateKey([u8; 32]);
impl From<PrivateKey> for StaticSecret {

53
rust/connlib/tunnel/src/tests/sim_client.rs
View File

@@ -10,7 +10,7 @@ use bimap::BiMap;
use connlib_shared::{
messages::{
client::{ResourceDescriptionCidr, ResourceDescriptionDns},
ClientId, DnsServer, Interface, ResourceId,
ClientId, DnsServer, GatewayId, Interface, ResourceId,
},
proptest::{client_id, domain_name},
DomainName,
@@ -83,10 +83,12 @@ impl SimClient {
dns_server: SocketAddr,
now: Instant,
) -> Option<Transmit<'static>> {
let dns_server = *self
.dns_by_sentinel
.get_by_right(&dns_server)
.expect("to have a sentinel DNS server for the sampled one");
let Some(dns_server) = self.dns_by_sentinel.get_by_right(&dns_server).copied() else {
tracing::error!(%dns_server, "Unknown DNS server");
return None;
};
tracing::debug!(%dns_server, %domain, "Sending DNS query");
let name = domain_to_hickory_name(domain);
@@ -208,6 +210,9 @@ impl SimClient {
}
/// Reference state for a particular client.
///
/// The reference state machine is designed to be as abstract as possible over connlib's functionality.
/// For example, we try to model connectivity to _resources_ and don't really care, which gateway is being used to route us there.
#[derive(Debug, Clone)]
pub struct RefClient {
pub(crate) id: ClientId,
@@ -239,9 +244,14 @@ pub struct RefClient {
pub(crate) connected_dns_resources: HashSet<(ResourceId, DomainName)>,
/// The expected ICMP handshakes.
pub(crate) expected_icmp_handshakes: VecDeque<(ResourceDst, IcmpSeq, IcmpIdentifier)>,
///
/// This is indexed by gateway because our assertions rely on the order of the sent packets.
pub(crate) expected_icmp_handshakes:
HashMap<GatewayId, VecDeque<(ResourceDst, IcmpSeq, IcmpIdentifier)>>,
/// The expected DNS handshakes.
pub(crate) expected_dns_handshakes: VecDeque<QueryId>,
pub(crate) gateways_by_resource: HashMap<ResourceId, GatewayId>,
}
impl RefClient {
@@ -291,9 +301,13 @@ impl RefClient {
};
tracing::Span::current().record("resource", tracing::field::display(resource.id));
let gateway = *self.gateways_by_resource.get(&resource.id).unwrap();
if self.is_connected_to_cidr(resource.id) && self.is_tunnel_ip(src) {
tracing::debug!("Connected to CIDR resource, expecting packet to be routed");
self.expected_icmp_handshakes
.entry(gateway)
.or_default()
.push_back((ResourceDst::Cidr(dst), seq, identifier));
return;
}
@@ -320,6 +334,8 @@ impl RefClient {
tracing::Span::current().record("resource", tracing::field::display(resource));
let gateway = *self.gateways_by_resource.get(&resource).unwrap();
if self
.connected_dns_resources
.contains(&(resource, dst.clone()))
@@ -327,6 +343,8 @@ impl RefClient {
{
tracing::debug!("Connected to DNS resource, expecting packet to be routed");
self.expected_icmp_handshakes
.entry(gateway)
.or_default()
.push_back((ResourceDst::Dns(dst), seq, identifier));
return;
}
@@ -404,11 +422,25 @@ impl RefClient {
/// An ICMP packet is valid if we didn't yet send an ICMP packet with the same seq and identifier.
pub(crate) fn is_valid_icmp_packet(&self, seq: &u16, identifier: &u16) -> bool {
self.expected_icmp_handshakes
.iter()
.all(|(_, existing_seq, existing_identifer)| {
self.expected_icmp_handshakes.values().flatten().all(
|(_, existing_seq, existing_identifer)| {
existing_seq != seq && existing_identifer != identifier
})
},
)
}
pub(crate) fn gateway_by_cidr_resource_ip(&self, dst: IpAddr) -> Option<GatewayId> {
let resource_id = self.cidr_resource_by_ip(dst)?;
let gateway_id = self.gateways_by_resource.get(&resource_id)?;
Some(*gateway_id)
}
pub(crate) fn gateway_by_domain_name(&self, dst: &DomainName) -> Option<GatewayId> {
let resource_id = self.dns_resource_by_domain(dst)?;
let gateway_id = self.gateways_by_resource.get(&resource_id)?;
Some(*gateway_id)
}
pub(crate) fn resolved_v4_domains(&self) -> Vec<DomainName> {
@@ -626,6 +658,7 @@ fn ref_client(
connected_dns_resources: Default::default(),
expected_icmp_handshakes: Default::default(),
expected_dns_handshakes: Default::default(),
gateways_by_resource: Default::default(),
},
)
}

11
rust/connlib/tunnel/src/tests/sim_gateway.rs
View File

@@ -3,7 +3,7 @@ use super::{
sim_net::{any_ip_stack, any_port, host, Host},
};
use crate::{tests::sut::hickory_name_to_domain, GatewayState};
use connlib_shared::{messages::GatewayId, proptest::gateway_id, DomainName};
use connlib_shared::DomainName;
use ip_packet::IpPacket;
use proptest::prelude::*;
use snownet::Transmit;
@@ -15,7 +15,6 @@ use std::{
/// Simulation state for a particular client.
pub(crate) struct SimGateway {
pub(crate) id: GatewayId,
pub(crate) sut: GatewayState,
pub(crate) received_icmp_requests: VecDeque<IpPacket<'static>>,
@@ -24,9 +23,8 @@ pub(crate) struct SimGateway {
}
impl SimGateway {
pub(crate) fn new(id: GatewayId, sut: GatewayState) -> Self {
pub(crate) fn new(sut: GatewayState) -> Self {
Self {
id,
sut,
received_icmp_requests: Default::default(),
buffer: vec![0u8; (1 << 16) - 1],
@@ -88,7 +86,6 @@ impl SimGateway {
/// Reference state for a particular gateway.
#[derive(Debug, Clone)]
pub struct RefGateway {
pub(crate) id: GatewayId,
pub(crate) key: PrivateKey,
}
@@ -97,7 +94,7 @@ impl RefGateway {
///
/// This simulates receiving the `init` message from the portal.
pub(crate) fn init(self) -> SimGateway {
SimGateway::new(self.id, GatewayState::new(self.key))
SimGateway::new(GatewayState::new(self.key))
}
}
@@ -106,5 +103,5 @@ pub(crate) fn ref_gateway_host() -> impl Strategy<Value = Host<RefGateway>> {
}
fn ref_gateway() -> impl Strategy<Value = RefGateway> {
(gateway_id(), private_key()).prop_map(move |(id, key)| RefGateway { id, key })
private_key().prop_map(move |key| RefGateway { key })
}

30
rust/connlib/tunnel/src/tests/sim_portal.rs
View File

@@ -1,22 +1,27 @@
use connlib_shared::messages::{
client::{ResourceDescriptionCidr, ResourceDescriptionDns, SiteId},
ClientId, GatewayId, ResourceId,
GatewayId, ResourceId,
};
use ip_network_table::IpNetworkTable;
use std::collections::{BTreeMap, HashSet};
use std::collections::{BTreeMap, HashMap, HashSet};
/// Stub implementation of the portal.
///
/// Currently, we only simulate a connection between a single client and a single gateway on a single site.
#[derive(Debug, Clone)]
pub(crate) struct SimPortal {
_client: ClientId,
gateway: GatewayId,
gateways_by_site: HashMap<SiteId, GatewayId>, // TODO: Technically, this is wrong because a site has multiple gateways but not the other way round.
}
impl SimPortal {
pub(crate) fn new(_client: ClientId, gateway: GatewayId) -> Self {
Self { _client, gateway }
pub(crate) fn new() -> Self {
Self {
gateways_by_site: HashMap::default(),
}
}
pub(crate) fn register_site(&mut self, site: SiteId, gateway: GatewayId) {
self.gateways_by_site.insert(site, gateway);
}
/// Picks, which gateway and site we should connect to for the given resource.
@@ -38,11 +43,12 @@ impl SimPortal {
.get(&resource)
.and_then(|r| Some(r.sites.first()?.id));
(
self.gateway,
cidr_site
.or(dns_site)
.expect("resource to be a known CIDR or DNS resource"),
)
let site_id = cidr_site
.or(dns_site)
.expect("resource to be a known CIDR or DNS resource");
let gateway_id = self.gateways_by_site.get(&site_id).unwrap();
(*gateway_id, site_id)
}
}

2
rust/connlib/tunnel/src/tests/sim_relay.rs
View File

@@ -20,7 +20,7 @@ pub(crate) struct SimRelay {
pub(crate) fn map_explode<'a>(
relays: impl Iterator<Item = (&'a RelayId, &'a Host<SimRelay>)> + 'a,
username: &'static str,
username: &'a str,
) -> impl Iterator<Item = (RelayId, RelaySocket, String, String, String)> + 'a {
relays.map(move |(id, r)| {
let (socket, username, password, realm) = r.inner().explode(

222
rust/connlib/tunnel/src/tests/sut.rs
View File

@@ -48,7 +48,7 @@ pub(crate) struct TunnelTest {
utc_now: DateTime<Utc>,
pub(crate) client: Host<SimClient>,
pub(crate) gateway: Host<SimGateway>,
pub(crate) gateways: HashMap<GatewayId, Host<SimGateway>>,
relays: HashMap<RelayId, Host<SimRelay>>,
portal: SimPortal,
@@ -76,10 +76,19 @@ impl StateMachineTest for TunnelTest {
let mut client = ref_state
.client
.map(|ref_client, _, _| ref_client.init(), debug_span!("client"));
let mut gateway = ref_state.gateway.map(
|ref_gateway, _, _| ref_gateway.init(),
debug_span!("gateway"),
);
let mut gateways = ref_state
.gateways
.iter()
.map(|(id, gateway)| {
let gateway = gateway.map(
|ref_gateway, _, _| ref_gateway.init(),
debug_span!("gateway", gid = %id),
);
(*id, gateway)
})
.collect::<HashMap<_, _>>();
let relays = ref_state
.relays
@@ -100,7 +109,6 @@ impl StateMachineTest for TunnelTest {
(*id, relay)
})
.collect::<HashMap<_, _>>();
let portal = SimPortal::new(client.inner().id, gateway.inner().id);
// Configure client and gateway with the relays.
client.exec_mut(|c| {
@@ -110,21 +118,23 @@ impl StateMachineTest for TunnelTest {
ref_state.now,
)
});
gateway.exec_mut(|g| {
g.sut.update_relays(
HashSet::default(),
HashSet::from_iter(map_explode(relays.iter(), "gateway")),
ref_state.now,
)
});
for (id, gateway) in &mut gateways {
gateway.exec_mut(|g| {
g.sut.update_relays(
HashSet::default(),
HashSet::from_iter(map_explode(relays.iter(), &format!("gateway_{id}"))),
ref_state.now,
)
});
}
let mut this = Self {
now: ref_state.now,
utc_now: ref_state.utc_now,
network: ref_state.network.clone(),
client,
gateway,
portal,
gateways,
portal: SimPortal::new(),
logger,
relays,
};
@@ -147,14 +157,26 @@ impl StateMachineTest for TunnelTest {
// Act: Apply the transition
match transition {
Transition::AddCidrResource(r) => {
Transition::AddCidrResource { resource, gateway } => {
for site in &resource.sites {
state.portal.register_site(site.id, gateway)
}
state
.client
.exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Cidr(r)]));
.exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Cidr(resource)]));
}
Transition::AddDnsResource {
resource, gateway, ..
} => {
for site in &resource.sites {
state.portal.register_site(site.id, gateway)
}
state
.client
.exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Dns(resource)]))
}
Transition::AddDnsResource { resource, .. } => state
.client
.exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Dns(resource)])),
Transition::RemoveResource(id) => {
state.client.exec_mut(|c| c.sut.remove_resources(&[id]))
}
@@ -267,13 +289,17 @@ impl StateMachineTest for TunnelTest {
) {
let ref_client = ref_state.client.inner();
let sim_client = state.client.inner();
let sim_gateway = state.gateway.inner();
let sim_gateways = state
.gateways
.iter()
.map(|(id, g)| (*id, g.inner()))
.collect();
// Assert our properties: Check that our actual state is equivalent to our expectation (the reference state).
assert_icmp_packets_properties(
ref_client,
sim_client,
sim_gateway,
sim_gateways,
&ref_state.global_dns_records,
);
assert_dns_packets_properties(ref_client, sim_client);
@@ -295,7 +321,7 @@ impl TunnelTest {
///
/// Consequently, this function needs to loop until no host can make progress at which point we consider the [`Transition`] complete.
fn advance(&mut self, ref_state: &ReferenceState, buffered_transmits: &mut BufferedTransmits) {
loop {
'outer: loop {
if let Some(transmit) = buffered_transmits.pop() {
self.dispatch_transmit(transmit, buffered_transmits, &ref_state.global_dns_records);
continue;
@@ -325,24 +351,27 @@ impl TunnelTest {
}
});
if let Some(transmit) = self.gateway.exec_mut(|g| g.sut.poll_transmit()) {
buffered_transmits.push(transmit, &self.gateway);
continue;
}
if let Some(event) = self.gateway.exec_mut(|g| g.sut.poll_event()) {
self.on_gateway_event(self.gateway.inner().id, event);
continue;
for (_, gateway) in self.gateways.iter_mut() {
if let Some(transmit) = gateway.exec_mut(|g| g.sut.poll_transmit()) {
buffered_transmits.push(transmit, gateway);
continue 'outer;
}
}
let mut any_relay_advanced = false;
for (id, gateway) in self.gateways.iter_mut() {
let Some(event) = gateway.exec_mut(|g| g.sut.poll_event()) else {
continue;
};
on_gateway_event(*id, event, &mut self.client, self.now);
continue 'outer;
}
for (_, relay) in self.relays.iter_mut() {
let Some(message) = relay.exec_mut(|r| r.sut.next_command()) else {
continue;
};
any_relay_advanced = true;
match message {
firezone_relay::Command::SendMessage { payload, recipient } => {
let dst = recipient.into_socket();
@@ -369,10 +398,8 @@ impl TunnelTest {
relay.exec_mut(|r| r.allocations.remove(&(family, port)));
}
}
}
if any_relay_advanced {
continue;
continue 'outer;
}
if self.handle_timeout(self.now, self.utc_now) {
@@ -399,16 +426,16 @@ impl TunnelTest {
self.client.exec_mut(|c| c.sut.handle_timeout(now));
};
if self
.gateway
.exec_mut(|g| g.sut.poll_timeout())
.is_some_and(|t| t <= now)
{
any_advanced = true;
for (_, gateway) in self.gateways.iter_mut() {
if gateway
.exec_mut(|g| g.sut.poll_timeout())
.is_some_and(|t| t <= now)
{
any_advanced = true;
self.gateway
.exec_mut(|g| g.sut.handle_timeout(now, utc_now))
};
gateway.exec_mut(|g| g.sut.handle_timeout(now, utc_now))
};
}
for (_, relay) in self.relays.iter_mut() {
if relay
@@ -451,15 +478,16 @@ impl TunnelTest {
self.client
.exec_mut(|c| c.handle_packet(payload, src, dst, self.now));
}
HostId::Gateway(_) => {
let Some(transmit) = self
.gateway
HostId::Gateway(id) => {
let gateway = self.gateways.get_mut(&id).expect("unknown gateway");
let Some(transmit) = gateway
.exec_mut(|g| g.handle_packet(global_dns_records, payload, src, dst, self.now))
else {
return;
};
buffered_transmits.push(transmit, &self.gateway);
buffered_transmits.push(transmit, gateway);
}
HostId::Relay(id) => {
let relay = self.relays.get_mut(&id).expect("unknown relay");
@@ -487,16 +515,30 @@ impl TunnelTest {
global_dns_records: &BTreeMap<DomainName, HashSet<IpAddr>>,
) {
match event {
ClientEvent::AddedIceCandidates { candidates, .. } => self.gateway.exec_mut(|g| {
for candidate in candidates {
g.sut.add_ice_candidate(src, candidate, self.now)
}
}),
ClientEvent::RemovedIceCandidates { candidates, .. } => self.gateway.exec_mut(|g| {
for candidate in candidates {
g.sut.remove_ice_candidate(src, candidate)
}
}),
ClientEvent::AddedIceCandidates {
candidates,
conn_id,
} => {
let gateway = self.gateways.get_mut(&conn_id).expect("unknown gateway");
gateway.exec_mut(|g| {
for candidate in candidates {
g.sut.add_ice_candidate(src, candidate, self.now)
}
})
}
ClientEvent::RemovedIceCandidates {
candidates,
conn_id,
} => {
let gateway = self.gateways.get_mut(&conn_id).expect("unknown gateway");
gateway.exec_mut(|g| {
for candidate in candidates {
g.sut.remove_ice_candidate(src, candidate)
}
})
}
ClientEvent::ConnectionIntent {
resource,
connected_gateway_ids,
@@ -533,8 +575,13 @@ impl TunnelTest {
match request {
Request::NewConnection(new_connection) => {
let answer = self
.gateway
let Some(gateway) = self.gateways.get_mut(&new_connection.gateway_id)
else {
tracing::error!("Unknown gateway");
return;
};
let answer = gateway
.exec_mut(|g| {
g.sut.accept(
self.client.inner().id,
@@ -579,14 +626,19 @@ impl TunnelTest {
},
},
resource_id,
self.gateway.inner().sut.public_key(),
gateway.inner().sut.public_key(),
self.now,
)
})
.unwrap();
}
Request::ReuseConnection(reuse_connection) => {
self.gateway
let gateway = self
.gateways
.get_mut(&reuse_connection.gateway_id)
.expect("unknown gateway");
gateway
.exec_mut(|g| {
g.sut.allow_access(
resource,
@@ -603,6 +655,11 @@ impl TunnelTest {
ClientEvent::SendProxyIps { connections } => {
for reuse_connection in connections {
let gateway = self
.gateways
.get_mut(&reuse_connection.gateway_id)
.expect("unknown gateway");
let resolved_ips = reuse_connection
.payload
.as_ref()
@@ -619,7 +676,7 @@ impl TunnelTest {
reuse_connection.resource_id,
);
self.gateway
gateway
.exec_mut(|g| {
g.sut.allow_access(
resource,
@@ -642,22 +699,6 @@ impl TunnelTest {
}
}
fn on_gateway_event(&mut self, src: GatewayId, event: GatewayEvent) {
match event {
GatewayEvent::AddedIceCandidates { candidates, .. } => self.client.exec_mut(|c| {
for candidate in candidates {
c.sut.add_ice_candidate(src, candidate, self.now)
}
}),
GatewayEvent::RemovedIceCandidates { candidates, .. } => self.client.exec_mut(|c| {
for candidate in candidates {
c.sut.remove_ice_candidate(src, candidate)
}
}),
GatewayEvent::RefreshDns { .. } => todo!(),
}
}
// TODO: Should we vary the following things via proptests?
// - Forwarded DNS query timing out?
// - hickory error?
@@ -694,6 +735,27 @@ impl TunnelTest {
}
}
fn on_gateway_event(
src: GatewayId,
event: GatewayEvent,
client: &mut Host<SimClient>,
now: Instant,
) {
match event {
GatewayEvent::AddedIceCandidates { candidates, .. } => client.exec_mut(|c| {
for candidate in candidates {
c.sut.add_ice_candidate(src, candidate, now)
}
}),
GatewayEvent::RemovedIceCandidates { candidates, .. } => client.exec_mut(|c| {
for candidate in candidates {
c.sut.remove_ice_candidate(src, candidate)
}
}),
GatewayEvent::RefreshDns { .. } => todo!(),
}
}
fn map_client_resource_to_gateway_resource(
client_cidr_resources: &IpNetworkTable<ResourceDescriptionCidr>,
client_dns_resources: &BTreeMap<ResourceId, ResourceDescriptionDns>,

45
rust/connlib/tunnel/src/tests/transition.rs
View File

@@ -5,7 +5,7 @@ use super::{
use connlib_shared::{
messages::{
client::{ResourceDescriptionCidr, ResourceDescriptionDns},
DnsServer, ResourceId,
DnsServer, GatewayId, ResourceId,
},
proptest::*,
DomainName,
@@ -23,7 +23,10 @@ use std::{
#[allow(clippy::large_enum_variant)]
pub(crate) enum Transition {
/// Add a new CIDR resource to the client.
AddCidrResource(ResourceDescriptionCidr),
AddCidrResource {
resource: ResourceDescriptionCidr,
gateway: GatewayId,
},
/// Send an ICMP packet to non-resource IP.
SendICMPPacketToNonResourceIp {
src: IpAddr,
@@ -54,6 +57,7 @@ pub(crate) enum Transition {
resource: ResourceDescriptionDns,
/// The DNS records to add together with the resource.
records: HashMap<DomainName, HashSet<IpAddr>>,
gateway: GatewayId,
},
/// Send a DNS query.
SendDnsQuery {
@@ -178,17 +182,22 @@ where
)
}
pub(crate) fn non_wildcard_dns_resource() -> impl Strategy<Value = Transition> {
(dns_resource(), resolved_ips()).prop_map(|(resource, resolved_ips)| {
pub(crate) fn non_wildcard_dns_resource(
gateway: impl Strategy<Value = GatewayId>,
) -> impl Strategy<Value = Transition> {
(dns_resource(), resolved_ips(), gateway).prop_map(|(resource, resolved_ips, gateway)| {
Transition::AddDnsResource {
records: HashMap::from([(resource.address.parse().unwrap(), resolved_ips)]),
resource,
gateway,
}
})
}
pub(crate) fn star_wildcard_dns_resource() -> impl Strategy<Value = Transition> {
dns_resource().prop_flat_map(move |r| {
pub(crate) fn star_wildcard_dns_resource(
gateway: impl Strategy<Value = GatewayId>,
) -> impl Strategy<Value = Transition> {
(dns_resource(), gateway).prop_flat_map(move |(r, gateway)| {
let wildcard_address = format!("*.{}", r.address);
let records = subdomain_records(r.address, domain_name(1..3));
@@ -197,13 +206,20 @@ pub(crate) fn star_wildcard_dns_resource() -> impl Strategy<Value = Transition>
..r
});
(resource, records)
.prop_map(|(resource, records)| Transition::AddDnsResource { records, resource })
(resource, records, Just(gateway)).prop_map(|(resource, records, gateway)| {
Transition::AddDnsResource {
records,
resource,
gateway,
}
})
})
}
pub(crate) fn question_mark_wildcard_dns_resource() -> impl Strategy<Value = Transition> {
dns_resource().prop_flat_map(move |r| {
pub(crate) fn question_mark_wildcard_dns_resource(
gateway: impl Strategy<Value = GatewayId>,
) -> impl Strategy<Value = Transition> {
(dns_resource(), gateway).prop_flat_map(move |(r, gateway)| {
let wildcard_address = format!("?.{}", r.address);
let records = subdomain_records(r.address, domain_label());
@@ -212,8 +228,13 @@ pub(crate) fn question_mark_wildcard_dns_resource() -> impl Strategy<Value = Tra
..r
});
(resource, records)
.prop_map(|(resource, records)| Transition::AddDnsResource { records, resource })
(resource, records, Just(gateway)).prop_map(|(resource, records, gateway)| {
Transition::AddDnsResource {
records,
resource,
gateway,
}
})
})
}

2
rust/ip-packet/Cargo.toml
View File

@@ -14,7 +14,7 @@ proptest = ["dep:proptest"]
pnet_packet = { version = "0.34" }
hickory-proto = { workspace = true }
thiserror = "1"
proptest = { version = "1.4.0", optional = true }
proptest = { version = "1", optional = true }
tracing = "0.1"
[dev-dependencies]

2
rust/relay/Cargo.toml
View File

@@ -25,7 +25,7 @@ bytes = "1.4.0"
sha2 = "0.10.8"
base64 = "0.22.1"
once_cell = "1.17.1"
proptest = { version = "1.4.0", optional = true }
proptest = { version = "1", optional = true }
derive_more = { version = "0.99.18", features = ["from"] }
uuid = { version = "1.7.0", features = ["v4"] }
phoenix-channel = { path = "../phoenix-channel" }