refactor(connlib): reuse ResourceDst for assertion (#5206)

Currently, we field `expected_icmp_handshakes` tracks the resource
destination as an `IpAddr` with a separate `ResourceKind` enum to
correctly interpret the address.

When generating the transition, we already use a `ResourceDst` enum
differentiate between the two different kinds of resources.

We can reuse that same enum to make the assertion clearer.
This commit is contained in:
Thomas Eizinger
2024-06-04 14:34:54 +10:00
committed by GitHub
parent 40cba442a4
commit 708375bbb4
2 changed files with 66 additions and 54 deletions

View File

@@ -24,3 +24,5 @@ cc b9ec9994018db89c645fc97c5211f5c96211280f8d542719a4f42f3366c4efeb # shrinks to
cc c8bf9f43adf08de7d824b6a7f28395c44ce68bb6630929721a50d35fd26d4b0c
cc 4d3b3ccb53bd19c67f6bfc5b92ddecba5ced019a135551b6c311bd8d7eca6d87 # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 53399, tv_nsec: 781785315 }, utc_now: 2024-05-29T07:03:25.677460793Z, client: SimNode { id: ClientId(b44bfb87-8cd2-ed87-3ef5-5cd59bf26896), state: PrivateKey("b9e465a6d969ea1865616eaef184a0cc39beb74e8c735ad07869ffb6f5418607"), ip4_socket: None, ip6_socket: Some([::ffff:127.0.0.1]:62725), tunnel_ip4: 100.98.132.149, tunnel_ip6: ::ffff:39.172.151.89 }, gateway: SimNode { id: GatewayId(ded1f941-ee2e-89ae-310a-a80112b7b703), state: PrivateKey("04c49163fd9a275ff0d0c4b8e2799285b95b45b4a02eaefb4dcae5336ddd5314"), ip4_socket: Some(148.232.205.50:25094), ip6_socket: None, tunnel_ip4: 100.101.30.207, tunnel_ip6: 2c32:f92a:5f9c:f5a7:9f6c:3c01:6d5c:d5dc }, relay: SimRelay { id: RelayId(1a497298-250c-6eb4-5bc9-f2554aba9b0a), ip_stack: Dual { ip4: 127.0.0.1, ip6: 388:303:eaa1:5cf2:e5b6:efc8:e6c:16c8 }, allocations: {} }, system_dns_resolvers: [::ffff:155.173.240.25, 112.55.144.77, 0.0.0.0], upstream_dns_resolvers: [], client_cidr_resources: {}, client_dns_resources: {}, resolved_domain_names: {}, client_dns_cache: {}, connected_cidr_resources: {}, expected_icmp_handshakes: [] }, [AddDnsResource(ResourceDescriptionDns { id: ResourceId(00000000-0000-0000-0000-000000000000), address: "aaa.aaa", name: "aaaa", address_description: "aaac", sites: [Site { name: "aacu", id: SiteId(8d98ea90-34fc-d8a7-d9ff-e78c0a83ac8a) }] }), SendQueryToDnsResource { r_idx: Index(10952768576437935965), r_type: AAAA, query_id: 12136, dns_server_idx: Index(16815185780125962912), resolved_ips: [::ffff:161.201.149.31, c734:3f2d:6208:de4e:c2ff:4edc:a82f:7e52, ::ffff:84.4.110.208] }, SendICMPPacketToResolvedAddress { a_idx: Index(6259982162279900885), seq: 9698, identifier: 59707 }], None)
cc 74d80501f87f21717ef6716bb411d2c86a038c2afd14a55ae0836101c9934a84 # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 7768, tv_nsec: 250403424 }, utc_now: 2024-05-30T02:01:37.147716325Z, client: SimNode { id: ClientId(9d83c082-f0c5-a80c-7030-a63095328b1e), state: PrivateKey("b9c59d3cebc50944388b411d646d51cc73ca5ec98ddc88819bdff27fb02296a2"), ip4_socket: Some(60.84.174.227:48562), ip6_socket: None, tunnel_ip4: 100.90.50.225, tunnel_ip6: 900:70d0:2dcc:d8ba:4e3c:3c16:6bb1:8850 }, gateway: SimNode { id: GatewayId(849688fc-8b3b-2952-423c-6fe647524a5f), state: PrivateKey("fdd5811cefe48f6102da7faea8678ee88405681464dc8c724b296d9e5ee0b0a7"), ip4_socket: Some(127.0.0.1:8110), ip6_socket: None, tunnel_ip4: 100.71.235.46, tunnel_ip6: ::ffff:40.42.59.67 }, relay: SimRelay { id: RelayId(77f0ba0d-260c-6bcd-6393-5e8312d9979e), ip_stack: Dual { ip4: 194.126.216.130, ip6: 8d85:fe6a:c7bf:522d:eeaa:2196:2fa:34fa }, allocations: {} }, system_dns_resolvers: [195.217.80.80], upstream_dns_resolvers: [IpPort(IpDnsServer { address: [::ffff:127.0.0.1]:53 }), IpPort(IpDnsServer { address: 90.60.79.207:53 })], client_cidr_resources: {}, client_dns_resources: {}, client_dns_records: {}, client_connected_cidr_resources: {}, resolved_domain_names: {}, expected_icmp_handshakes: [] }, [AddDnsResource { resource: ResourceDescriptionDns { id: ResourceId(00000000-0000-0000-0000-000000000000), address: "aaa.aaa", name: "aaaa", address_description: "aagx", sites: [Site { name: "aaaa", id: SiteId(3b5953ef-c3d4-080e-0f0f-cd9687c9c47d) }] }, resolved_ips: {0.0.0.0, ed15:424:2fc2:720b:50b3:cec8:d98d:aca6, 127.0.0.1} }, AddCidrResource(ResourceDescriptionCidr { id: ResourceId(8be4162b-381a-1c16-ad75-b3046d883a67), address: V4(Ipv4Network { network_address: 127.0.0.1, netmask: 32 }), name: "qmrom", address_description: "sdeslw", sites: [Site { name: "mljenoza", id: SiteId(d346d198-98bb-e503-391b-75961cc35c46) }] }), AddCidrResource(ResourceDescriptionCidr { id: ResourceId(a51a69e5-ee7d-2346-b5ac-c567b47bd414), address: V4(Ipv4Network { network_address: 0.0.0.0, netmask: 26 }), name: "zdauam", address_description: "ktnfyyqh", sites: [Site { name: "ndzgrwjpcc", id: SiteId(f0af320e-3ab4-7c01-d098-e59349e30169) }, Site { name: "ovplh", id: SiteId(2cf8d363-9246-b03b-07dc-2faa66d02aef) }] }), SendICMPPacketToIp4Resource { r_idx: Index(11883287362675254181), seq: 57419, identifier: 34979 }, SendQueryToDnsResource { r_idx: Index(9880014133669589572), r_type: AAAA, query_id: 51384, dns_server_idx: Index(8501323410582601301) }, SendICMPPacketToIp4Resource { r_idx: Index(11771623222579691892), seq: 14899, identifier: 41403 }], None)
cc 06903deacc497922cc9aedd36057972b0359557cf8059d64a157cec6c955a5cc
cc 7e67359c1d9c75586e7fe0cc51e2eee24eba2aca62636195f23b706198599b83 # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 10684, tv_nsec: 44735276 }, utc_now: 2024-06-04T02:25:25.942486401Z, client: SimNode { id: ClientId(d6d58c72-329f-3b21-e04c-f0dac70f9624), state: PrivateKey("a30c602b873ed38203a46651dc6b1a99727747c1b0e6c7bbef3771ecb362e39c"), ip4_socket: None, ip6_socket: Some([abe5:6b84:872f:9555:fd9d:1903:5405:2080]:50619), tunnel_ip4: 100.79.77.32, tunnel_ip6: fd00:2021:1111::16:d758 }, gateway: SimNode { id: GatewayId(feba32ae-7694-e0e8-dff1-f8d97c8d84e3), state: PrivateKey("b3b5e02616d05cbb9571f9f9141d111018f71afa7752b03235022a2c9bf4a9c6"), ip4_socket: Some(127.0.0.1:18210), ip6_socket: None, tunnel_ip4: 100.83.224.20, tunnel_ip6: fd00:2021:1111::10:364b }, relay: SimRelay { id: RelayId(83a23b67-df4d-f152-dd86-0a1b751fb20c), ip_stack: Dual { ip4: 222.66.120.215, ip6: ::ffff:120.157.205.82 }, allocations: {} }, system_dns_resolvers: [::ffff:91.131.146.146, 239.247.149.60, 6261:7566:b7cf:1456:c05c:87e5:cece:3a80], upstream_dns_resolvers: [IpPort(IpDnsServer { address: 0.0.0.0:53 }), IpPort(IpDnsServer { address: [::ffff:65.59.137.206]:53 })], client_cidr_resources: {}, client_dns_resources: {}, client_dns_records: {}, client_connected_cidr_resources: {}, global_dns_records: {}, expected_icmp_handshakes: [] }, [AddDnsResource { resource: ResourceDescriptionDns { id: ResourceId(1cf14706-3f70-8050-e975-cec7dd733c76), address: "mjeen.gpb", name: "chkv", address_description: "fjglgqy", sites: [Site { name: "uxkzphj", id: SiteId(042c51f7-5588-3920-b9dc-62cf9d98b575) }] }, resolved_ips: {::ffff:76.150.216.79, ::ffff:127.0.0.1} }, UpdateUpstreamDnsServers { servers: [IpPort(IpDnsServer { address: 0.0.0.0:53 }), IpPort(IpDnsServer { address: [884f:28d5:cf64:1b47:3349:5b74:4561:ed6b]:53 }), IpPort(IpDnsServer { address: [5cc2:ad51:e599:7253:32e5:824c:249b:fec7]:53 })] }, SendQueryToDnsResource { r_idx: Index(17860973355611427927), r_type: AAAA, query_id: 15986, dns_server_idx: Index(8754080677899521430) }, SendICMPPacketToResource { idx: Index(11870326947086465412), seq: 48005, identifier: 17804, src: None }], None)

View File

@@ -127,13 +127,7 @@ struct ReferenceState {
global_dns_records: HashMap<DomainName, HashSet<IpAddr>>,
/// The expected ICMP handshakes.
expected_icmp_handshakes: VecDeque<(IpAddr, IcmpSeq, IcmpIdentifier, ResourceKind)>,
}
#[derive(Debug, Clone, Copy)]
enum ResourceKind {
Cidr,
Dns,
expected_icmp_handshakes: VecDeque<(ResourceDst, IcmpSeq, IcmpIdentifier)>,
}
type QueryId = u16;
@@ -598,9 +592,10 @@ impl ReferenceStateMachine for ReferenceState {
state
.client_dns_records
.entry(domain)
.entry(domain.clone())
.or_default()
.extend(ips_resolved_by_query);
state.client_dns_records.entry(domain).or_default().sort();
}
Transition::SendICMPPacketToNonResourceIp { .. } => {
// Packets to non-resources are dropped, no state change required.
@@ -614,7 +609,6 @@ impl ReferenceStateMachine for ReferenceState {
let dst = state
.sample_resource_dst(idx)
.expect("Transition to only be sampled if we have at least one resource");
let dst = dst.into_reference_packet_dst();
state.on_icmp_packet(*src, dst, *seq, *identifier);
}
@@ -1227,6 +1221,10 @@ impl TunnelTest {
.entry(requested_domain.clone())
.or_default()
.extend(proxy_ips);
self.client_dns_records
.entry(requested_domain)
.or_default()
.sort();
return;
}
@@ -1239,45 +1237,46 @@ impl TunnelTest {
/// Several helper functions to make the reference state more readable.
impl ReferenceState {
#[tracing::instrument(level = "debug", skip_all, fields(dst, resource))]
fn on_icmp_packet(
&mut self,
src: Option<IpAddr>,
dst: impl Into<IpAddr>,
seq: u16,
identifier: u16,
) {
let dst = dst.into();
fn on_icmp_packet(&mut self, src: Option<IpAddr>, dst: ResourceDst, seq: u16, identifier: u16) {
let packet_originates_from_client =
src.is_none() || src.is_some_and(|src| src == self.client.tunnel_ip(src)); // Packets from IPs other than the client's are dropped by the gateway. We still create connections for them though.
tracing::Span::current().record("dst", tracing::field::display(dst));
if self.client_dns_records.values().flatten().contains(&dst)
&& packet_originates_from_client
{
tracing::debug!("Connected to DNS resource, expecting packet to be routed");
self.expected_icmp_handshakes
.push_back((dst, seq, identifier, ResourceKind::Dns));
return;
match &dst {
ResourceDst::Cidr(ip_dst) => {
tracing::Span::current().record("dst", tracing::field::display(ip_dst));
// Second, if we are not yet connected, check if we have a resource for this IP.
let Some((_, resource)) = self.client_cidr_resources.longest_match(*ip_dst) else {
tracing::debug!("No resource corresponds to IP");
return;
};
if self.client_connected_cidr_resources.contains(&resource.id)
&& packet_originates_from_client
{
tracing::debug!("Connected to CIDR resource, expecting packet to be routed");
self.expected_icmp_handshakes
.push_back((dst, seq, identifier));
return;
}
// If we have a resource, the first packet will initiate a connection to the gateway.
tracing::debug!(
"Not connected to resource, expecting to trigger connection intent"
);
self.client_connected_cidr_resources.insert(resource.id);
}
ResourceDst::Dns(domain, _) => {
tracing::Span::current().record("dst", tracing::field::display(domain));
if self.client_dns_records.contains_key(domain) && packet_originates_from_client {
tracing::debug!("Connected to DNS resource, expecting packet to be routed");
self.expected_icmp_handshakes
.push_back((dst, seq, identifier));
return;
}
}
}
// Second, if we are not yet connected, check if we have a resource for this IP.
let Some((_, resource)) = self.client_cidr_resources.longest_match(dst) else {
tracing::debug!("No resource corresponds to IP");
return;
};
if self.client_connected_cidr_resources.contains(&resource.id)
&& packet_originates_from_client
{
tracing::debug!("Connected to CIDR resource, expecting packet to be routed");
self.expected_icmp_handshakes
.push_back((dst, seq, identifier, ResourceKind::Cidr));
return;
}
// If we have a resource, the first packet will initiate a connection to the gateway.
tracing::debug!("Not connected to resource, expecting to trigger connection intent");
self.client_connected_cidr_resources.insert(resource.id);
}
fn sample_resource_dst(&self, idx: &sample::Index) -> Option<ResourceDst> {
@@ -1380,7 +1379,7 @@ impl ReferenceState {
let not_existing_icmp =
self.expected_icmp_handshakes
.iter()
.all(|(_, existing_seq, existing_identifer, _)| {
.all(|(_, existing_seq, existing_identifer)| {
existing_seq != seq && existing_identifer != identifier
});
@@ -2034,12 +2033,12 @@ fn assert_icmp_packets_properties(state: &mut TunnelTest, ref_state: &ReferenceS
let unexpected_icmp_replies = find_unexpected_entries(
&ref_state.expected_icmp_handshakes,
&state.client_received_icmp_replies,
|(_, seq_a, id_a, _), (seq_b, id_b)| seq_a == seq_b && id_a == id_b,
|(_, seq_a, id_a), (seq_b, id_b)| seq_a == seq_b && id_a == id_b,
);
let unexpected_icmp_requests = find_unexpected_entries(
&ref_state.expected_icmp_handshakes,
&state.gateway_received_icmp_requests,
|(_, seq_a, id_a, _), (seq_b, id_b)| seq_a == seq_b && id_a == id_b,
|(_, seq_a, id_a), (seq_b, id_b)| seq_a == seq_b && id_a == id_b,
);
assert_eq!(
@@ -2053,7 +2052,7 @@ fn assert_icmp_packets_properties(state: &mut TunnelTest, ref_state: &ReferenceS
"Unexpected ICMP requests on gateway"
);
for (resource_dst, seq, identifier, kind) in ref_state.expected_icmp_handshakes.iter() {
for (resource_dst, seq, identifier) in ref_state.expected_icmp_handshakes.iter() {
let client_sent_request = &state
.client_sent_icmp_requests
.get(&(*seq, *identifier))
@@ -2085,8 +2084,8 @@ fn assert_icmp_packets_properties(state: &mut TunnelTest, ref_state: &ReferenceS
"ICMP request source == ICMP reply destination"
);
match kind {
ResourceKind::Cidr => {
match resource_dst {
ResourceDst::Cidr(resource_dst) => {
// For CIDR resources, the expected dst is always known.
assert_eq!(
@@ -2095,11 +2094,22 @@ fn assert_icmp_packets_properties(state: &mut TunnelTest, ref_state: &ReferenceS
"ICMP request on gateway to target correct CIDR resource"
);
}
ResourceKind::Dns => {
ResourceDst::Dns(domain, _) => {
// For DNS resources, we need to look up, which proxy IP we used.
// We consider it an implementation detail, how connlib assigns those IPs.
// We do want to assert that the mapping is stable.
let actual_destination = gateway_received_request.destination();
let possible_resource_ips = ref_state
.global_dns_records
.get(domain)
.expect("ICMP packet for DNS resource to target known domain");
assert!(
possible_resource_ips.contains(&actual_destination),
"ICMP request on gateway to target a known resource IP"
);
match state
.client_proxy_ip_mapping
.entry(client_sent_request.destination())
@@ -2107,13 +2117,13 @@ fn assert_icmp_packets_properties(state: &mut TunnelTest, ref_state: &ReferenceS
Entry::Vacant(v) => {
// We have to gradually discover connlib's mapping ...
// For the first packet, we just save the IP that we ended up talking to.
v.insert(gateway_received_request.source());
v.insert(gateway_received_request.destination());
}
Entry::Occupied(o) => {
assert_eq!(
gateway_received_request.source(),
gateway_received_request.destination(),
*o.get(),
"ICMP request on gateway to target correct DNS resource"
"ICMP request on client to target correct same IP of DNS resource"
);
}
}