diff --git a/rust/connlib/snownet/src/allocation.rs b/rust/connlib/snownet/src/allocation.rs index fde3bff74..721435c4a 100644 --- a/rust/connlib/snownet/src/allocation.rs +++ b/rust/connlib/snownet/src/allocation.rs @@ -431,7 +431,7 @@ impl Allocation { let maybe_candidate = message.attributes().find_map(|a| srflx_candidate(local, a)); update_candidate(maybe_candidate, current_srflx_candidate, &mut self.events); - self.log_update(); + self.log_update(now); // Second, check if we have already determined which socket to use for this relay. // We send 2 BINDING requests to start with (one for each IP version) and the first one coming back wins. @@ -485,7 +485,7 @@ impl Allocation { &mut self.events, ); - self.log_update(); + self.log_update(now); while let Some(peer) = self.buffered_channel_bindings.pop() { debug_assert!( @@ -503,7 +503,7 @@ impl Allocation { self.allocation_lifetime = Some((now, lifetime.lifetime())); - self.log_update(); + self.log_update(now); } CHANNEL_BIND => { let Some(channel) = original_request @@ -759,13 +759,13 @@ impl Allocation { self.credentials.is_some() } - fn log_update(&self) { + fn log_update(&self, now: Instant) { tracing::info!( srflx_ip4 = ?self.ip4_srflx_candidate.as_ref().map(|c| c.addr()), srflx_ip6 = ?self.ip6_srflx_candidate.as_ref().map(|c| c.addr()), relay_ip4 = ?self.ip4_allocation.as_ref().map(|c| c.addr()), relay_ip6 = ?self.ip6_allocation.as_ref().map(|c| c.addr()), - lifetime = ?self.allocation_lifetime, + remaining_lifetime = ?self.allocation_lifetime.and_then(|(created_at, d)| d.checked_sub(now.checked_duration_since(created_at)?)), "Updated allocation" ); } diff --git a/rust/connlib/tunnel/proptest-regressions/tests.txt b/rust/connlib/tunnel/proptest-regressions/tests.txt index 603d509ab..7bdc3040b 100644 --- a/rust/connlib/tunnel/proptest-regressions/tests.txt +++ b/rust/connlib/tunnel/proptest-regressions/tests.txt @@ -47,3 +47,4 @@ cc 7f664f0f323f6ee66ce3646f2dcee3060cd44a0ba6cc5dac93b6e05b506a8b77 # shrinks to cc 7f024279e0477056e37717e4b886d5016acc9fab8964611b58f86587767bba2e # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 29920, tv_nsec: 368456206 }, utc_now: 2024-07-08T07:58:51.159038594Z, client: Host { inner: SimNode { id: ClientId(00000000-0000-0000-0000-000000000000), state: (PrivateKey("0000000000000000000000000000000000000000000000000000000000000000"), {}), tunnel_ip4: 100.64.0.1, tunnel_ip6: fd00:2021:1111:: }, ip4: Some(203.0.113.1), ip6: None, old_ips: [], old_ports: {0}, default_port: 1, allocated_ports: {(1, V4)} }, gateway: Host { inner: SimNode { id: GatewayId(00000000-0000-0000-0000-000000000000), state: PrivateKey("0000000000000000000000000000000000000000000000000000000000000001"), tunnel_ip4: 100.64.0.2, tunnel_ip6: fd00:2021:1111::1 }, ip4: Some(203.0.113.2), ip6: None, old_ips: [], old_ports: {0}, default_port: 28202, allocated_ports: {(28202, V4)} }, relay: Host { inner: SimRelay { id: RelayId(447095dd-0f99-00f7-9cdf-14db9370f920), state: 12549090075203197752, ip_stack: Dual { ip4: 203.0.113.3, ip6: 2001:db80::2 }, allocations: {} }, ip4: Some(203.0.113.3), ip6: Some(2001:db80::2), old_ips: [], old_ports: {0}, default_port: 3478, allocated_ports: {(3478, V4), (3478, V6)} }, system_dns_resolvers: [247.236.135.67], upstream_dns_resolvers: [IpPort(IpDnsServer { address: 220.197.151.19:53 }), IpPort(IpDnsServer { address: 89.28.131.88:53 }), IpPort(IpDnsServer { address: 127.0.0.1:53 }), IpPort(IpDnsServer { address: [::ffff:127.0.0.1]:53 }), IpPort(IpDnsServer { address: [d567:8313:60b:e85d:365f:a699:b9d3:e1dd]:53 }), IpPort(IpDnsServer { address: [::ffff:185.167.149.64]:53 })], client_cidr_resources: {}, client_dns_resources: {}, client_dns_records: {}, client_connected_cidr_resources: {}, client_connected_dns_resources: {}, global_dns_records: {Name(crmoi.ecmx.cwku.): {0.0.0.0, ::ffff:48.3.194.225}, Name(hmtwct.hzpvl.ehgvx.): {17.184.192.139, 78.22.114.111, 6d72:a7e2:d73b:4f0b:50e0:ba1:22b8:ac48, 32.78.117.29}, Name(mahy.ekzr.): {c24a:a44c:7cbf:99e:37a3:e6d2:6f61:ab88}, Name(yeth.fbjqx.fnbhvt.): {127.0.0.1}, Name(ouwbw.kznvif.): {46.29.147.154, ::ffff:127.0.0.1, 127.0.0.1, ::ffff:232.178.178.113}, Name(nai.hem.ldep.): {::ffff:127.0.0.1}, Name(sbwn.qkckf.ltura.): {a59f:3fed:f882:3e0d:c967:564f:3f45:68f, 100.103.142.95, 145.154.71.75}, Name(lcqdtu.mot.): {1004:28c1:2c31:9817:cef2:5c3d:9956:d716, ::ffff:50.209.39.213}, Name(ipeyg.hvld.mpsik.): {90.2.106.81}, Name(mlte.snvnft.): {4.133.196.171}, Name(fvdig.sxuuco.): {::ffff:141.228.199.134, 191.77.153.166, 127.0.0.1}, Name(rel.rzveiw.xgd.): {127.0.0.1, 61.187.72.87, 200.81.188.150, 37.142.30.190, ::ffff:118.191.216.138}, Name(kgq.xlj.): {58d1:fd65:d116:b285:d0d8:169b:2249:1edb, 127.0.0.1}}, client_known_host_records: {}, expected_icmp_handshakes: [], expected_dns_handshakes: [] }, [AddDnsResource { resource: ResourceDescriptionDns { id: ResourceId(d8b923a2-6a73-df2d-4baf-557089dcf022), address: "*.tfq.swmt", name: "fkkx", address_description: Some("rzatphisen"), sites: [Site { name: "xewrde", id: SiteId(5ce8469f-4883-b33d-09ff-a468dfe73513) }, Site { name: "gttfyjb", id: SiteId(bc45737f-02af-c129-e90f-52d76e040d63) }] }, records: {Name(kzb.oatiu.tfq.swmt.): {::ffff:132.15.20.125, 165.213.243.89, 2f93:9571:7862:df97:7a8e:79c7:2e87:eb78, ::ffff:237.141.93.128}, Name(jehqf.iicd.tfq.swmt.): {::ffff:127.0.0.1, 126.122.243.61, 8158:f77:e6af:40a3:51ed:b37f:401e:fe16, 29c3:8a5c:43d7:2563:d400:e50a:c999:45f9}} }, SendDnsQuery { domain: Name(jehqf.iicd.tfq.swmt.), r_type: AAAA, query_id: 1402, dns_server: 89.28.131.88:53 }, SendICMPPacketToDnsResource { src: fd00:2021:1111::, dst: Name(jehqf.iicd.tfq.swmt.), seq: 0, identifier: 0 }, RoamClient { ip4: None, ip6: Some(2001:db80::3), port: 1 }, AddCidrResource(ResourceDescriptionCidr { id: ResourceId(00000000-0000-0000-0000-000000000000), address: V6(Ipv6Network { network_address: ::ffff:0.0.0.0, netmask: 127 }), name: "aaaa", address_description: Some("zjmvzpv"), sites: [Site { name: "mvjorgubk", id: SiteId(807103d6-5141-8196-c36c-cd6539fb3210) }] }), SendICMPPacketToDnsResource { src: fd00:2021:1111::, dst: Name(jehqf.iicd.tfq.swmt.), seq: 0, identifier: 0 }], None) cc b51827a79947ba53609819934c11750a46302b58ca9cd993c24a7f1312ef2a0b # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 32296, tv_nsec: 202806172 }, utc_now: 2024-07-08T08:38:26.993388469Z, client: Host { inner: SimNode { id: ClientId(00000000-0000-0000-0000-000000000000), state: (PrivateKey("0000000000000000000000000000000000000000000000000000000000000000"), {}), tunnel_ip4: 100.64.0.1, tunnel_ip6: fd00:2021:1111:: }, ip4: Some(203.0.113.1), ip6: None, old_ips: [], old_ports: {0}, default_port: 1, allocated_ports: {(1, V4)} }, gateway: Host { inner: SimNode { id: GatewayId(00000000-0000-0000-0000-000000000000), state: PrivateKey("0000000000000000000000000000000000000000000000000000000000000001"), tunnel_ip4: 100.64.0.2, tunnel_ip6: fd00:2021:1111::1 }, ip4: Some(203.0.113.2), ip6: None, old_ips: [], old_ports: {0}, default_port: 1, allocated_ports: {(1, V4)} }, relay: Host { inner: SimRelay { id: RelayId(00000000-0000-0000-0000-000000000000), state: 0, ip_stack: Dual { ip4: 203.0.113.3, ip6: 2001:db80:: }, allocations: {} }, ip4: Some(203.0.113.3), ip6: Some(2001:db80::), old_ips: [], old_ports: {0}, default_port: 3478, allocated_ports: {(3478, V6), (3478, V4)} }, system_dns_resolvers: [0.0.0.0], upstream_dns_resolvers: [], client_cidr_resources: {}, client_dns_resources: {}, client_dns_records: {}, client_connected_cidr_resources: {}, client_connected_dns_resources: {}, global_dns_records: {}, client_known_host_records: {}, expected_icmp_handshakes: [], expected_dns_handshakes: [], network: RoutingTable { routes: {(V4(Ipv4Network { network_address: 203.0.113.1, netmask: 32 }), Client(ClientId(00000000-0000-0000-0000-000000000000))), (V4(Ipv4Network { network_address: 203.0.113.2, netmask: 32 }), Gateway(GatewayId(00000000-0000-0000-0000-000000000000))), (V4(Ipv4Network { network_address: 203.0.113.3, netmask: 32 }), Relay(RelayId(00000000-0000-0000-0000-000000000000))), (V6(Ipv6Network { network_address: 2001:db80::, netmask: 128 }), Relay(RelayId(00000000-0000-0000-0000-000000000000)))} } }, [RoamClient { ip4: Some(203.0.113.1), ip6: Some(2001:db80::34), port: 1 }], None) cc 6012a980592466d46f51f36bf511aeca295d174a188cd6ae9e6e799b2eeb4bdd # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 32727, tv_nsec: 104744448 }, utc_now: 2024-07-08T08:45:37.895326745Z, client: Host { inner: SimNode { id: ClientId(00000000-0000-0000-0000-000000000000), state: (PrivateKey("0000000000000000000000000000000000000000000000000000000000000000"), {}), tunnel_ip4: 100.64.0.1, tunnel_ip6: fd00:2021:1111:: }, ip4: Some(203.0.113.1), ip6: None, old_ips: [], old_ports: {0}, default_port: 25710, allocated_ports: {(25710, V4)} }, gateway: Host { inner: SimNode { id: GatewayId(0a90cfcb-9a3f-9f65-601c-8675fc305ce4), state: PrivateKey("3b0a23bbb0d745be5d50cb4c7dcad3a3761a55353133417707f93a6abbfdc255"), tunnel_ip4: 100.64.0.2, tunnel_ip6: fd00:2021:1111::1 }, ip4: None, ip6: Some(2001:db80::24), old_ips: [], old_ports: {0}, default_port: 12684, allocated_ports: {(12684, V6)} }, relay: Host { inner: SimRelay { id: RelayId(7cb9f334-4ccc-019d-85f2-2390848e66c1), state: 5611803068338389074, ip_stack: Dual { ip4: 203.0.113.40, ip6: 2001:db80::58 }, allocations: {} }, ip4: Some(203.0.113.40), ip6: Some(2001:db80::58), old_ips: [], old_ports: {0}, default_port: 3478, allocated_ports: {(3478, V4), (3478, V6)} }, system_dns_resolvers: [127.0.0.1, ::ffff:247.177.145.232, ::ffff:179.131.145.140], upstream_dns_resolvers: [], client_cidr_resources: {}, client_dns_resources: {}, client_dns_records: {}, client_connected_cidr_resources: {}, client_connected_dns_resources: {}, global_dns_records: {}, client_known_host_records: {}, expected_icmp_handshakes: [], expected_dns_handshakes: [], network: RoutingTable { routes: {(V4(Ipv4Network { network_address: 203.0.113.1, netmask: 32 }), Client(ClientId(00000000-0000-0000-0000-000000000000))), (V4(Ipv4Network { network_address: 203.0.113.40, netmask: 32 }), Relay(RelayId(7cb9f334-4ccc-019d-85f2-2390848e66c1))), (V6(Ipv6Network { network_address: 2001:db80::24, netmask: 128 }), Gateway(GatewayId(0a90cfcb-9a3f-9f65-601c-8675fc305ce4))), (V6(Ipv6Network { network_address: 2001:db80::58, netmask: 128 }), Relay(RelayId(7cb9f334-4ccc-019d-85f2-2390848e66c1)))} } }, [AddDnsResource { resource: ResourceDescriptionDns { id: ResourceId(00000000-0000-0000-0000-000000000000), address: "wqgecw.ymss.bdyrft", name: "aaaa", address_description: Some("jydrm"), sites: [Site { name: "xttyddtwg", id: SiteId(deca6d49-b6f5-cff2-8a83-9f3626d89bf6) }] }, records: {Name(wqgecw.ymss.bdyrft.): {229.251.45.251, ::ffff:59.2.60.9}} }, RoamClient { ip4: Some(203.0.113.1), ip6: Some(2001:db80::49), port: 10114 }, SendDnsQuery { domain: Name(wqgecw.ymss.bdyrft.), r_type: A, query_id: 0, dns_server: 127.0.0.1:53 }, RoamClient { ip4: None, ip6: Some(2001:db80::49), port: 64407 }, SendICMPPacketToDnsResource { src: 100.64.0.1, dst: Name(wqgecw.ymss.bdyrft.), seq: 0, identifier: 0 }, SendICMPPacketToDnsResource { src: 100.64.0.1, dst: Name(wqgecw.ymss.bdyrft.), seq: 0, identifier: 0 }], None) +cc 4e94fe9054e7d0a8d3b0cdc9564e562e7a20e33e6969312e8088601600f28101 # shrinks to (initial_state, transitions, seen_counter) = (ReferenceState { now: Instant { tv_sec: 50453, tv_nsec: 990736170 }, utc_now: 2024-07-10T04:15:45.022371741Z, client: Host { inner: RefClient { key: PrivateKey("6308b982f16e56b13315d94d5b028af3aaf089a46fd9364f28bdbad6839c5c5b"), known_hosts: {"sog.iuv": [35.0.92.197, 127.0.0.1], "vziabd.xmwhup.pungj": [ffd6:9c77:debb:b85a:b249:547e:a9a2:cf79, 73.50.5.200], "mbjw.mtq.favlrp": [127.0.0.1, 191.162.14.155], "icjv.adsl.pxd": [::ffff:139.251.78.60, 108.203.183.88, ::ffff:127.0.0.1, 172.85.230.143, a14c:3bc0:da24:1562:8d35:5703:c164:38dd], "spyhqn.jeq.tsrqyz": [159.196.134.93, 127.0.0.1, 93.234.19.166, ::ffff:0.0.0.0], "gxrk.uglvuf": [343d:c04:5b36:3e6b:eeab:ce4a:1f9b:957e, 127.0.0.1, 138.47.94.195], "xlgd.ire": [::ffff:106.75.98.27, 185.25.235.30, 238.179.216.1, ::ffff:86.200.52.17, ::ffff:187.107.84.34], "hpkq.sicp.jxfzq": [231.237.65.176, 127.0.0.1], "zua.irezz": [::ffff:0.0.0.0], "exl.onofs.xwt": [::ffff:172.91.84.123], "pvqerk.wicwl.gycnoo": [::ffff:210.196.90.97, ::ffff:236.5.9.242, 24.90.160.237, 33fb:f5e:78d5:6fcf:2679:804e:4cb4:8685]}, tunnel_ip4: 100.64.0.1, tunnel_ip6: fd00:2021:1111::, system_dns_resolvers: [127.0.0.1, ::ffff:236.238.31.35, 191.46.28.254], upstream_dns_resolvers: [], cidr_resources: {}, dns_resources: {}, dns_records: {}, connected_cidr_resources: {}, connected_dns_resources: {}, expected_icmp_handshakes: [], expected_dns_handshakes: [] }, sim_state: ClientId(d16f3c0b-0ba2-1880-fd7c-fd5d42eb30e0), ip4: None, ip6: Some(2001:db80::63), old_ports: {0}, default_port: 57958, allocated_ports: {(57958, V6)} }, gateway: Host { inner: RefGateway { key: PrivateKey("a2b756447eeed8fef6f097182747a36f8c04e4e16757c660c41330e0638060e5") }, sim_state: GatewayId(f1e32f0c-bea4-415c-f839-e414140b6a7b), ip4: Some(203.0.113.77), ip6: None, old_ports: {0}, default_port: 3275, allocated_ports: {(3275, V4)} }, relays: {RelayId(a1570084-8b27-bfc0-f65a-e6d9828ee8e9): Host { inner: 13930734621424630275, sim_state: (), ip4: Some(203.0.113.9), ip6: Some(2001:db80::55), old_ports: {0}, default_port: 3478, allocated_ports: {(3478, V4), (3478, V6)} }, RelayId(0d3f6d84-4d2c-22cd-7449-f34df43d91cf): Host { inner: 13962067921649706663, sim_state: (), ip4: Some(203.0.113.76), ip6: Some(2001:db80::1), old_ports: {0}, default_port: 3478, allocated_ports: {(3478, V4), (3478, V6)} }}, global_dns_records: {Name(xkaxzj.qcku.wmbrrt.): {::ffff:127.0.0.1, 16.64.128.18}}, network: RoutingTable { routes: {(V4(Ipv4Network { network_address: 203.0.113.9, netmask: 32 }), Relay(RelayId(a1570084-8b27-bfc0-f65a-e6d9828ee8e9))), (V4(Ipv4Network { network_address: 203.0.113.76, netmask: 32 }), Relay(RelayId(0d3f6d84-4d2c-22cd-7449-f34df43d91cf))), (V4(Ipv4Network { network_address: 203.0.113.77, netmask: 32 }), Gateway(GatewayId(f1e32f0c-bea4-415c-f839-e414140b6a7b))), (V6(Ipv6Network { network_address: 2001:db80::1, netmask: 128 }), Relay(RelayId(0d3f6d84-4d2c-22cd-7449-f34df43d91cf))), (V6(Ipv6Network { network_address: 2001:db80::55, netmask: 128 }), Relay(RelayId(a1570084-8b27-bfc0-f65a-e6d9828ee8e9))), (V6(Ipv6Network { network_address: 2001:db80::63, netmask: 128 }), Client(ClientId(d16f3c0b-0ba2-1880-fd7c-fd5d42eb30e0)))} } }, [AddDnsResource { resource: ResourceDescriptionDns { id: ResourceId(7db7a1da-c996-3b3d-bae1-665f356f28f4), address: "?.fnhvin.pei", name: "gmhyshdhlp", address_description: Some("wezrfwmh"), sites: [Site { name: "swkmncjeae", id: SiteId(cadbbea8-4aba-d0a9-e6f4-fc38a36231eb) }] }, records: {Name(fthl.fnhvin.pei.): {175.242.230.69}} }, SendDnsQuery { domain: Name(fthl.fnhvin.pei.), r_type: AAAA, query_id: 27157, dns_server: [::ffff:236.238.31.35]:53 }, SendICMPPacketToDnsResource { src: fd00:2021:1111::, dst: Name(fthl.fnhvin.pei.), seq: 47, identifier: 6323 }, SendICMPPacketToDnsResource { src: fd00:2021:1111::, dst: Name(fthl.fnhvin.pei.), seq: 14926, identifier: 8427 }], None) diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs index b9993ba66..8a6bdcfb0 100644 --- a/rust/connlib/tunnel/src/client.rs +++ b/rust/connlib/tunnel/src/client.rs @@ -317,6 +317,24 @@ impl ClientState { } } + #[cfg(test)] + pub(crate) fn tunnel_ip4(&self) -> Option { + Some(self.interface_config.as_ref()?.ipv4) + } + + #[cfg(test)] + pub(crate) fn tunnel_ip6(&self) -> Option { + Some(self.interface_config.as_ref()?.ipv6) + } + + #[cfg(test)] + pub(crate) fn tunnel_ip_for(&self, dst: IpAddr) -> Option { + Some(match dst { + IpAddr::V4(_) => self.tunnel_ip4()?.into(), + IpAddr::V6(_) => self.tunnel_ip6()?.into(), + }) + } + pub(crate) fn resources(&self) -> Vec { self.resource_ids .values() diff --git a/rust/connlib/tunnel/src/tests.rs b/rust/connlib/tunnel/src/tests.rs index 4a6d4d228..edc70aa39 100644 --- a/rust/connlib/tunnel/src/tests.rs +++ b/rust/connlib/tunnel/src/tests.rs @@ -4,8 +4,9 @@ use proptest::test_runner::Config; mod assertions; mod composite_strategy; mod reference; +mod sim_client; +mod sim_gateway; mod sim_net; -mod sim_node; mod sim_portal; mod sim_relay; mod strategies; diff --git a/rust/connlib/tunnel/src/tests/assertions.rs b/rust/connlib/tunnel/src/tests/assertions.rs index 118685379..20659450a 100644 --- a/rust/connlib/tunnel/src/tests/assertions.rs +++ b/rust/connlib/tunnel/src/tests/assertions.rs @@ -1,4 +1,7 @@ -use super::{reference::ReferenceState, TunnelTest}; +use super::{ + sim_client::{RefClient, SimClient}, + sim_gateway::SimGateway, +}; use crate::tests::reference::ResourceDst; use connlib_shared::DomainName; use ip_packet::IpPacket; @@ -14,10 +17,15 @@ use std::{ /// - For CIDR resources, that is the actual CIDR resource IP. /// - For DNS resources, the IP must match one of the resolved IPs for the domain. /// 3. For DNS resources, the mapping of proxy IP to actual resource IP must be stable. -pub(crate) fn assert_icmp_packets_properties(state: &TunnelTest, ref_state: &ReferenceState) { +pub(crate) fn assert_icmp_packets_properties( + ref_client: &RefClient, + sim_client: &SimClient, + sim_gateway: &SimGateway, + global_dns_records: &BTreeMap>, +) { let unexpected_icmp_replies = find_unexpected_entries( - &ref_state.expected_icmp_handshakes, - &state.client_received_icmp_replies, + &ref_client.expected_icmp_handshakes, + &sim_client.received_icmp_replies, |(_, seq_a, id_a), (seq_b, id_b)| seq_a == seq_b && id_a == id_b, ); assert_eq!( @@ -27,28 +35,28 @@ pub(crate) fn assert_icmp_packets_properties(state: &TunnelTest, ref_state: &Ref ); assert_eq!( - ref_state.expected_icmp_handshakes.len(), - state.gateway_received_icmp_requests.len(), + ref_client.expected_icmp_handshakes.len(), + sim_gateway.received_icmp_requests.len(), "Unexpected ICMP requests on gateway" ); - tracing::info!(target: "assertions", "✅ Performed the expected {} ICMP handshakes", state.gateway_received_icmp_requests.len()); + tracing::info!(target: "assertions", "✅ Performed the expected {} ICMP handshakes", sim_gateway.received_icmp_requests.len()); let mut mapping = HashMap::new(); - for ((resource_dst, seq, identifier), gateway_received_request) in ref_state + for ((resource_dst, seq, identifier), gateway_received_request) in ref_client .expected_icmp_handshakes .iter() - .zip(state.gateway_received_icmp_requests.iter()) + .zip(sim_gateway.received_icmp_requests.iter()) { let _guard = tracing::info_span!(target: "assertions", "icmp", %seq, %identifier).entered(); - let client_sent_request = &state - .client_sent_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 = &state - .client_received_icmp_replies + let client_received_reply = &sim_client + .received_icmp_replies .get(&(*seq, *identifier)) .expect("to have ICMP reply on client"); @@ -56,10 +64,7 @@ pub(crate) fn assert_icmp_packets_properties(state: &TunnelTest, ref_state: &Ref assert_eq!( gateway_received_request.source(), - ref_state - .client - .inner() - .tunnel_ip(gateway_received_request.source()), + ref_client.tunnel_ip_for(gateway_received_request.source()), "ICMP request on gateway to originate from client" ); @@ -70,7 +75,7 @@ pub(crate) fn assert_icmp_packets_properties(state: &TunnelTest, ref_state: &Ref ResourceDst::Dns(domain) => { assert_destination_is_dns_resource( gateway_received_request, - &ref_state.global_dns_records, + global_dns_records, domain, ); @@ -84,18 +89,18 @@ pub(crate) fn assert_icmp_packets_properties(state: &TunnelTest, ref_state: &Ref } } -pub(crate) fn assert_known_hosts_are_valid(state: &TunnelTest, ref_state: &ReferenceState) { - for (record, actual_addrs) in &state.client_dns_records { - if let Some(expected_addrs) = ref_state.client_known_host_records.get(&record.to_string()) { +pub(crate) fn assert_known_hosts_are_valid(ref_client: &RefClient, sim_client: &SimClient) { + for (record, actual_addrs) in &sim_client.dns_records { + if let Some(expected_addrs) = ref_client.known_hosts.get(&record.to_string()) { assert_eq!(actual_addrs, expected_addrs); } } } -pub(crate) fn assert_dns_packets_properties(state: &TunnelTest, ref_state: &ReferenceState) { +pub(crate) fn assert_dns_packets_properties(ref_client: &RefClient, sim_client: &SimClient) { let unexpected_icmp_replies = find_unexpected_entries( - &ref_state.expected_dns_handshakes, - &state.client_received_dns_responses, + &ref_client.expected_dns_handshakes, + &sim_client.received_dns_responses, |id_a, id_b| id_a == id_b, ); @@ -105,15 +110,15 @@ pub(crate) fn assert_dns_packets_properties(state: &TunnelTest, ref_state: &Refe "Unexpected DNS replies on client" ); - for query_id in ref_state.expected_dns_handshakes.iter() { + for query_id in ref_client.expected_dns_handshakes.iter() { let _guard = tracing::info_span!(target: "assertions", "dns", %query_id).entered(); - let client_sent_query = state - .client_sent_dns_queries + let client_sent_query = sim_client + .sent_dns_queries .get(query_id) .expect("to have DNS query on client"); - let client_received_response = state - .client_received_dns_responses + let client_received_response = sim_client + .received_dns_responses .get(query_id) .expect("to have DNS response on client"); diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs index 4e0dd03de..8e6714257 100644 --- a/rust/connlib/tunnel/src/tests/reference.rs +++ b/rust/connlib/tunnel/src/tests/reference.rs @@ -1,25 +1,17 @@ use super::{ - composite_strategy::CompositeStrategy, sim_net::*, sim_node::*, sim_relay::*, strategies::*, - transition::*, IcmpIdentifier, IcmpSeq, QueryId, + composite_strategy::CompositeStrategy, sim_client::*, sim_gateway::*, sim_net::*, sim_relay::*, + strategies::*, transition::*, }; use chrono::{DateTime, Utc}; -use connlib_shared::{ - messages::{ - client::{ResourceDescriptionCidr, ResourceDescriptionDns}, - ClientId, DnsServer, GatewayId, RelayId, ResourceId, - }, - proptest::*, - DomainName, -}; +use connlib_shared::{messages::RelayId, proptest::*, DomainName, StaticSecret}; use hickory_proto::rr::RecordType; -use ip_network_table::IpNetworkTable; -use itertools::Itertools; use prop::collection; use proptest::{prelude::*, sample}; use proptest_state_machine::ReferenceStateMachine; use std::{ - collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}, - net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, + collections::{BTreeMap, HashMap, HashSet}, + fmt, + net::IpAddr, time::{Duration, Instant}, }; @@ -30,46 +22,15 @@ use std::{ pub(crate) struct ReferenceState { pub(crate) now: Instant, pub(crate) utc_now: DateTime, - #[allow(clippy::type_complexity)] - pub(crate) client: Host>)>>, - pub(crate) gateway: Host>, - pub(crate) relays: HashMap>>, - - /// The DNS resolvers configured on the client outside of connlib. - pub(crate) system_dns_resolvers: Vec, - /// The upstream DNS resolvers configured in the portal. - pub(crate) upstream_dns_resolvers: Vec, - - /// The CIDR resources the client is aware of. - pub(crate) client_cidr_resources: IpNetworkTable, - /// The DNS resources the client is aware of. - pub(crate) client_dns_resources: BTreeMap, - - /// The client's DNS records. - /// - /// The IPs assigned to a domain by connlib are an implementation detail that we don't want to model in these tests. - /// Instead, we just remember what _kind_ of records we resolved to be able to sample a matching src IP. - client_dns_records: BTreeMap>, - - /// The CIDR resources the client is connected to. - client_connected_cidr_resources: HashSet, - - /// The DNS resources the client is connected to. - client_connected_dns_resources: HashSet<(ResourceId, DomainName)>, + pub(crate) client: Host, + pub(crate) gateway: Host, + pub(crate) relays: HashMap>, /// All IP addresses a domain resolves to in our test. /// /// This is used to e.g. mock DNS resolution on the gateway. pub(crate) global_dns_records: BTreeMap>, - /// Ips that the client know about without querying the global records - pub(crate) client_known_host_records: HashMap>, - - /// The expected ICMP handshakes. - pub(crate) expected_icmp_handshakes: VecDeque<(ResourceDst, IcmpSeq, IcmpIdentifier)>, - /// The expected DNS handshakes. - pub(crate) expected_dns_handshakes: VecDeque, - pub(crate) network: RoutingTable, } @@ -92,31 +53,17 @@ impl ReferenceStateMachine for ReferenceState { let mut tunnel_ip4s = tunnel_ip4s(); let mut tunnel_ip6s = tunnel_ip6s(); - let client_prototype = sim_node_prototype( - client_id(), - client_state(), - &mut tunnel_ip4s, - &mut tunnel_ip6s, - ); - let gateway_prototype = sim_node_prototype( - gateway_id(), - gateway_state(), - &mut tunnel_ip4s, - &mut tunnel_ip6s, - ); ( - client_prototype, - gateway_prototype, - collection::hash_map(relay_id(), sim_relay_prototype(), 2), - system_dns_servers(), - upstream_dns_servers(), + ref_client_host(&mut tunnel_ip4s, &mut tunnel_ip6s), + ref_gateway_host(), + collection::hash_map(relay_id(), relay_prototype(), 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, system_dns, upstream_dns, global_dns, now, utc_now)| { + |(c, g, relays, global_dns, now, utc_now)| { let mut routing_table = RoutingTable::default(); if !routing_table.add_host(c.inner().id, &c) { @@ -132,78 +79,22 @@ impl ReferenceStateMachine for ReferenceState { }; } - Some(( - c, - g, - relays, - system_dns, - upstream_dns, - global_dns, - now, - utc_now, - routing_table, - )) + Some((c, g, relays, global_dns, now, utc_now, routing_table)) }, ) .prop_filter( "client and gateway priv key must be different", - |(c, g, _, _, _, _, _, _, _)| c.inner().state.0 != g.inner().state, - ) - .prop_filter( - "at least one DNS server needs to be reachable", - |(c, _, _, system_dns, upstream_dns, _, _, _, _)| { - // TODO: PRODUCTION CODE DOES NOT HANDLE THIS! - - if !upstream_dns.is_empty() { - if c.ip4.is_none() && upstream_dns.iter().all(|s| s.ip().is_ipv4()) { - return false; - } - if c.ip6.is_none() && upstream_dns.iter().all(|s| s.ip().is_ipv6()) { - return false; - } - - return true; - } - - if c.ip4.is_none() && system_dns.iter().all(|s| s.is_ipv4()) { - return false; - } - if c.ip6.is_none() && system_dns.iter().all(|s| s.is_ipv6()) { - return false; - } - - true - }, + |(c, g, _, _, _, _, _)| c.inner().key != g.inner().key, ) .prop_map( - |( + |(client, gateway, relays, global_dns_records, now, utc_now, network)| Self { + now, + utc_now, client, gateway, relays, - system_dns_resolvers, - upstream_dns_resolvers, global_dns_records, - now, - utc_now, network, - )| Self { - now, - utc_now, - client: client.clone(), - gateway, - relays, - system_dns_resolvers, - upstream_dns_resolvers, - global_dns_records, - client_known_host_records: client.inner().state.1.clone(), - network, - client_cidr_resources: IpNetworkTable::new(), - client_connected_cidr_resources: Default::default(), - expected_icmp_handshakes: Default::default(), - client_dns_resources: Default::default(), - client_dns_records: Default::default(), - expected_dns_handshakes: Default::default(), - client_connected_dns_resources: Default::default(), }, ) .boxed() @@ -239,35 +130,51 @@ impl ReferenceStateMachine for ReferenceState { question_mark_wildcard_dns_resource(), ], ) - .with_if_not_empty(10, state.ipv4_cidr_resource_dsts(), |ip4_resources| { - icmp_to_cidr_resource( - packet_source_v4(state.client.inner().tunnel_ip4), - sample::select(ip4_resources), - ) - }) - .with_if_not_empty(10, state.ipv6_cidr_resource_dsts(), |ip6_resources| { - icmp_to_cidr_resource( - packet_source_v6(state.client.inner().tunnel_ip6), - sample::select(ip6_resources), - ) - }) - .with_if_not_empty(10, state.resolved_v4_domains(), |dns_v4_domains| { - icmp_to_dns_resource( - packet_source_v4(state.client.inner().tunnel_ip4), - sample::select(dns_v4_domains), - ) - }) - .with_if_not_empty(10, state.resolved_v6_domains(), |dns_v6_domains| { - icmp_to_dns_resource( - packet_source_v6(state.client.inner().tunnel_ip6), - sample::select(dns_v6_domains), - ) - }) + .with_if_not_empty( + 10, + state.client.inner().ipv4_cidr_resource_dsts(), + |ip4_resources| { + icmp_to_cidr_resource( + packet_source_v4(state.client.inner().tunnel_ip4), + sample::select(ip4_resources), + ) + }, + ) + .with_if_not_empty( + 10, + state.client.inner().ipv6_cidr_resource_dsts(), + |ip6_resources| { + icmp_to_cidr_resource( + packet_source_v6(state.client.inner().tunnel_ip6), + sample::select(ip6_resources), + ) + }, + ) + .with_if_not_empty( + 10, + state.client.inner().resolved_v4_domains(), + |dns_v4_domains| { + icmp_to_dns_resource( + packet_source_v4(state.client.inner().tunnel_ip4), + sample::select(dns_v4_domains), + ) + }, + ) + .with_if_not_empty( + 10, + state.client.inner().resolved_v6_domains(), + |dns_v6_domains| { + icmp_to_dns_resource( + packet_source_v6(state.client.inner().tunnel_ip6), + sample::select(dns_v6_domains), + ) + }, + ) .with_if_not_empty( 10, ( - state.all_domains(), - state.v4_dns_servers(), + state.all_domains(state.client.inner()), + state.client.inner().v4_dns_servers(), state.client.ip4, ), |(domains, v4_dns_servers, _)| { @@ -277,8 +184,8 @@ impl ReferenceStateMachine for ReferenceState { .with_if_not_empty( 10, ( - state.all_domains(), - state.v6_dns_servers(), + state.all_domains(state.client.inner()), + state.client.inner().v6_dns_servers(), state.client.ip6, ), |(domains, v6_dns_servers, _)| { @@ -287,7 +194,10 @@ impl ReferenceStateMachine for ReferenceState { ) .with_if_not_empty( 1, - state.resolved_ip4_for_non_resources(), + state + .client + .inner() + .resolved_ip4_for_non_resources(&state.global_dns_records), |resolved_non_resource_ip4s| { ping_random_ip( packet_source_v4(state.client.inner().tunnel_ip4), @@ -297,7 +207,10 @@ impl ReferenceStateMachine for ReferenceState { ) .with_if_not_empty( 1, - state.resolved_ip6_for_non_resources(), + state + .client + .inner() + .resolved_ip6_for_non_resources(&state.global_dns_records), |resolved_non_resource_ip6s| { ping_random_ip( packet_source_v6(state.client.inner().tunnel_ip6), @@ -305,7 +218,7 @@ impl ReferenceStateMachine for ReferenceState { ) }, ) - .with_if_not_empty(1, state.all_resources(), |resources| { + .with_if_not_empty(1, state.client.inner().all_resources(), |resources| { sample::select(resources).prop_map(Transition::RemoveResource) }) .boxed() @@ -317,20 +230,30 @@ impl ReferenceStateMachine for ReferenceState { fn apply(mut state: Self::State, transition: &Self::Transition) -> Self::State { match transition { Transition::AddCidrResource(r) => { - state.client_cidr_resources.insert(r.address, r.clone()); + state + .client + .exec_mut(|client| client.cidr_resources.insert(r.address, r.clone())); } Transition::RemoveResource(id) => { - state.client_cidr_resources.retain(|_, r| &r.id != id); - state.client_connected_cidr_resources.remove(id); - state.client_dns_resources.remove(id); + state + .client + .exec_mut(|client| client.cidr_resources.retain(|_, r| &r.id != id)); + state + .client + .exec_mut(|client| client.connected_cidr_resources.remove(id)); + state + .client + .exec_mut(|client| client.dns_resources.remove(id)); } Transition::AddDnsResource { resource: new_resource, records, } => { - let existing_resource = state - .client_dns_resources - .insert(new_resource.id, new_resource.clone()); + let existing_resource = state.client.exec_mut(|client| { + client + .dns_resources + .insert(new_resource.id, new_resource.clone()) + }); // For the client, there is no difference between a DNS resource and a truly global DNS name. // We store all records in the same map to follow the same model. @@ -339,7 +262,9 @@ impl ReferenceStateMachine for ReferenceState { // If a resource is updated (i.e. same ID but different address) and we are currently connected, we disconnect from it. if let Some(resource) = existing_resource { if new_resource.address != resource.address { - state.client_connected_cidr_resources.remove(&resource.id); + state.client.exec_mut(|client| { + client.connected_cidr_resources.remove(&resource.id) + }); state .global_dns_records @@ -348,9 +273,11 @@ impl ReferenceStateMachine for ReferenceState { // TODO: IN PRODUCTION, WE CANNOT DO THIS. // CHANGING A DNS RESOURCE BREAKS CLIENT UNTIL THEY DECIDE TO RE-QUERY THE RESOURCE. // WE DO THIS HERE TO ENSURE THE TEST DOESN'T RUN INTO THIS. - state - .client_dns_records - .retain(|name, _| !matches_domain(&resource.address, name)); + state.client.exec_mut(|client| { + client + .dns_records + .retain(|name, _| !matches_domain(&resource.address, name)) + }); } } } @@ -360,22 +287,30 @@ impl ReferenceStateMachine for ReferenceState { dns_server, query_id, .. - } => match state.dns_query_via_cidr_resource(dns_server.ip(), domain) { + } => match state + .client + .inner() + .dns_query_via_cidr_resource(dns_server.ip(), domain) + { Some(resource) - if !state.client_connected_cidr_resources.contains(&resource) - && !state - .client_known_host_records - .contains_key(&domain.to_string()) => + if !state.client.inner().is_connected_to_cidr(resource) + && !state.client.inner().is_known_host(&domain.to_string()) => { - state.client_connected_cidr_resources.insert(resource); + state + .client + .exec_mut(|client| client.connected_cidr_resources.insert(resource)); } Some(_) | None => { + state.client.exec_mut(|client| { + client + .dns_records + .entry(domain.clone()) + .or_default() + .insert(*r_type) + }); state - .client_dns_records - .entry(domain.clone()) - .or_default() - .insert(*r_type); - state.expected_dns_handshakes.push_back(*query_id); + .client + .exec_mut(|client| client.expected_dns_handshakes.push_back(*query_id)); } }, Transition::SendICMPPacketToNonResourceIp { .. } => { @@ -388,7 +323,9 @@ impl ReferenceStateMachine for ReferenceState { identifier, .. } => { - state.on_icmp_packet_to_cidr(*src, *dst, *seq, *identifier); + state.client.exec_mut(|client| { + client.on_icmp_packet_to_cidr(*src, *dst, *seq, *identifier) + }); } Transition::SendICMPPacketToDnsResource { src, @@ -396,13 +333,19 @@ impl ReferenceStateMachine for ReferenceState { seq, identifier, .. - } => state.on_icmp_packet_to_dns(*src, dst.clone(), *seq, *identifier), + } => state.client.exec_mut(|client| { + client.on_icmp_packet_to_dns(*src, dst.clone(), *seq, *identifier) + }), Transition::Tick { millis } => state.now += Duration::from_millis(*millis), Transition::UpdateSystemDnsServers { servers } => { - state.system_dns_resolvers.clone_from(servers); + state + .client + .exec_mut(|client| client.system_dns_resolvers.clone_from(servers)); } Transition::UpdateUpstreamDnsServers { servers } => { - state.upstream_dns_resolvers.clone_from(servers); + state + .client + .exec_mut(|client| client.upstream_dns_resolvers.clone_from(servers)); } Transition::RoamClient { ip4, ip6, .. } => { state.network.remove_host(&state.client); @@ -413,8 +356,12 @@ impl ReferenceStateMachine for ReferenceState { .add_host(state.client.inner().id, &state.client)); // When roaming, we are not connected to any resource and wait for the next packet to re-establish a connection. - state.client_connected_cidr_resources.clear(); - state.client_connected_dns_resources.clear(); + state + .client + .exec_mut(|client| client.connected_cidr_resources.clear()); + state + .client + .exec_mut(|client| client.connected_dns_resources.clear()); } }; @@ -457,9 +404,14 @@ impl ReferenceStateMachine for ReferenceState { } // TODO: PRODUCTION CODE DOES NOT HANDLE THIS. - let any_real_ip_overlaps_with_cidr_resource = resolved_ips - .iter() - .any(|resolved_ip| state.cidr_resource_by_ip(*resolved_ip).is_some()); + let any_real_ip_overlaps_with_cidr_resource = + resolved_ips.iter().any(|resolved_ip| { + state + .client + .inner() + .cidr_resource_by_ip(*resolved_ip) + .is_some() + }); if any_real_ip_overlaps_with_cidr_resource { return false; @@ -475,14 +427,20 @@ impl ReferenceStateMachine for ReferenceState { identifier, .. } => { - let is_valid_icmp_packet = state.is_valid_icmp_packet(seq, identifier); - let is_cidr_resource = state.client_cidr_resources.longest_match(*dst).is_some(); + let is_valid_icmp_packet = + state.client.inner().is_valid_icmp_packet(seq, identifier); + let is_cidr_resource = state + .client + .inner() + .cidr_resources + .longest_match(*dst) + .is_some(); is_valid_icmp_packet && !is_cidr_resource } Transition::SendICMPPacketToCidrResource { seq, identifier, .. - } => state.is_valid_icmp_packet(seq, identifier), + } => state.client.inner().is_valid_icmp_packet(seq, identifier), Transition::SendICMPPacketToDnsResource { seq, identifier, @@ -490,9 +448,11 @@ impl ReferenceStateMachine for ReferenceState { src, .. } => { - state.is_valid_icmp_packet(seq, identifier) + state.client.inner().is_valid_icmp_packet(seq, identifier) && state - .client_dns_records + .client + .inner() + .dns_records .get(dst) .is_some_and(|r| match src { IpAddr::V4(_) => r.contains(&RecordType::A), @@ -527,12 +487,13 @@ impl ReferenceStateMachine for ReferenceState { domain, dns_server, .. } => { state.global_dns_records.contains_key(domain) - && state.expected_dns_servers().contains(dns_server) - } - Transition::RemoveResource(id) => { - state.client_cidr_resources.iter().any(|(_, r)| &r.id == id) - || state.client_dns_resources.contains_key(id) + && state + .client + .inner() + .expected_dns_servers() + .contains(dns_server) } + Transition::RemoveResource(id) => state.client.inner().all_resources().contains(id), Transition::RoamClient { ip4, ip6, port } => { // In production, we always rebind to a new port so we never roam to our old existing IP / port combination. @@ -546,262 +507,20 @@ impl ReferenceStateMachine for ReferenceState { } } -/// Pub(crate) functions used across the test suite. -impl ReferenceState { - /// Returns the DNS servers that we expect connlib to use. - /// - /// If there are upstream DNS servers configured in the portal, it should use those. - /// Otherwise it should use whatever was configured on the system prior to connlib starting. - pub(crate) fn expected_dns_servers(&self) -> BTreeSet { - if !self.upstream_dns_resolvers.is_empty() { - return self - .upstream_dns_resolvers - .iter() - .map(|s| s.address()) - .collect(); - } - - self.system_dns_resolvers - .iter() - .map(|ip| SocketAddr::new(*ip, 53)) - .collect() - } -} - /// 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_to_cidr(&mut self, src: IpAddr, dst: IpAddr, seq: u16, identifier: u16) { - tracing::Span::current().record("dst", tracing::field::display(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(dst) else { - tracing::debug!("No resource corresponds to IP"); - return; - }; - tracing::Span::current().record("resource", tracing::field::display(resource.id)); - - if self.client_connected_cidr_resources.contains(&resource.id) - && self.client.inner().is_tunnel_ip(src) - { - tracing::debug!("Connected to CIDR resource, expecting packet to be routed"); - self.expected_icmp_handshakes - .push_back((ResourceDst::Cidr(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); - } - - #[tracing::instrument(level = "debug", skip_all, fields(dst, resource))] - fn on_icmp_packet_to_dns(&mut self, src: IpAddr, dst: DomainName, seq: u16, identifier: u16) { - tracing::Span::current().record("dst", tracing::field::display(&dst)); - - let Some(resource) = self.dns_resource_by_domain(&dst) else { - tracing::debug!("No resource corresponds to IP"); - return; - }; - - tracing::Span::current().record("resource", tracing::field::display(resource)); - - if self - .client_connected_dns_resources - .contains(&(resource, dst.clone())) - && self.client.inner().is_tunnel_ip(src) - { - tracing::debug!("Connected to DNS resource, expecting packet to be routed"); - self.expected_icmp_handshakes - .push_back((ResourceDst::Dns(dst), seq, identifier)); - return; - } - - debug_assert!( - self.client_dns_records.iter().any(|(name, _)| name == &dst), - "Should only sample ICMPs to domains that we resolved" - ); - - tracing::debug!("Not connected to resource, expecting to trigger connection intent"); - self.client_connected_dns_resources.insert((resource, dst)); - } - - fn ipv4_cidr_resource_dsts(&self) -> Vec { - let mut ips = vec![]; - - // This is an imperative loop on purpose because `ip-network` appears to have a bug with its `size_hint` and thus `.extend` does not work reliably? - for (network, _) in self.client_cidr_resources.iter_ipv4() { - if network.netmask() == 31 || network.netmask() == 32 { - ips.push(network.network_address()); - } else { - for ip in network.hosts() { - ips.push(ip) - } - } - } - - ips - } - - fn ipv6_cidr_resource_dsts(&self) -> Vec { - let mut ips = vec![]; - - // This is an imperative loop on purpose because `ip-network` appears to have a bug with its `size_hint` and thus `.extend` does not work reliably? - for (network, _) in self.client_cidr_resources.iter_ipv6() { - if network.netmask() == 127 || network.netmask() == 128 { - ips.push(network.network_address()); - } else { - for ip in network - .subnets_with_prefix(128) - .map(|i| i.network_address()) - { - ips.push(ip) - } - } - } - - ips - } - - fn resolved_v4_domains(&self) -> Vec { - self.resolved_domains() - .filter_map(|(domain, records)| { - records - .iter() - .any(|r| matches!(r, RecordType::A)) - .then_some(domain) - }) - .collect() - } - - fn resolved_v6_domains(&self) -> Vec { - self.resolved_domains() - .filter_map(|(domain, records)| { - records - .iter() - .any(|r| matches!(r, RecordType::AAAA)) - .then_some(domain) - }) - .collect() - } - - fn all_domains(&self) -> Vec { + fn all_domains(&self, client: &RefClient) -> Vec { self.global_dns_records .keys() .cloned() .chain( - self.client_known_host_records + client + .known_hosts .keys() .map(|h| DomainName::vec_from_str(h).unwrap()), ) .collect() } - - fn resolved_domains(&self) -> impl Iterator)> + '_ { - self.client_dns_records - .iter() - .filter(|(domain, _)| self.dns_resource_by_domain(domain).is_some()) - .map(|(domain, ips)| (domain.clone(), ips.clone())) - } - - /// An ICMP packet is valid if we didn't yet send an ICMP packet with the same seq and identifier. - fn is_valid_icmp_packet(&self, seq: &u16, identifier: &u16) -> bool { - self.expected_icmp_handshakes - .iter() - .all(|(_, existing_seq, existing_identifer)| { - existing_seq != seq && existing_identifer != identifier - }) - } - - fn v4_dns_servers(&self) -> Vec { - self.expected_dns_servers() - .into_iter() - .filter_map(|s| match s { - SocketAddr::V4(v4) => Some(v4), - SocketAddr::V6(_) => None, - }) - .collect() - } - - fn v6_dns_servers(&self) -> Vec { - self.expected_dns_servers() - .into_iter() - .filter_map(|s| match s { - SocketAddr::V6(v6) => Some(v6), - SocketAddr::V4(_) => None, - }) - .collect() - } - - fn dns_resource_by_domain(&self, domain: &DomainName) -> Option { - self.client_dns_resources - .values() - .filter(|r| is_subdomain(&domain.to_string(), &r.address)) - .sorted_by_key(|r| r.address.len()) - .rev() - .map(|r| r.id) - .next() - } - - fn cidr_resource_by_ip(&self, ip: IpAddr) -> Option { - self.client_cidr_resources - .longest_match(ip) - .map(|(_, r)| r.id) - } - - fn resolved_ip4_for_non_resources(&self) -> Vec { - self.resolved_ips_for_non_resources() - .filter_map(|ip| match ip { - IpAddr::V4(v4) => Some(v4), - IpAddr::V6(_) => None, - }) - .collect() - } - - fn resolved_ip6_for_non_resources(&self) -> Vec { - self.resolved_ips_for_non_resources() - .filter_map(|ip| match ip { - IpAddr::V6(v6) => Some(v6), - IpAddr::V4(_) => None, - }) - .collect() - } - - fn resolved_ips_for_non_resources(&self) -> impl Iterator + '_ { - self.client_dns_records - .iter() - .filter_map(|(domain, _)| { - self.dns_resource_by_domain(domain) - .is_none() - .then_some(self.global_dns_records.get(domain)) - }) - .flatten() - .flatten() - .copied() - } - - /// Returns the CIDR resource we will forward the DNS query for the given name to. - /// - /// DNS servers may be resources, in which case queries that need to be forwarded actually need to be encapsulated. - fn dns_query_via_cidr_resource( - &self, - dns_server: IpAddr, - domain: &DomainName, - ) -> Option { - // If we are querying a DNS resource, we will issue a connection intent to the DNS resource, not the CIDR resource. - if self.dns_resource_by_domain(domain).is_some() { - return None; - } - - self.cidr_resource_by_ip(dns_server) - } - - fn all_resources(&self) -> Vec { - let cidr_resources = self.client_cidr_resources.iter().map(|(_, r)| r.id); - let dns_resources = self.client_dns_resources.keys().copied(); - - Vec::from_iter(cidr_resources.chain(dns_resources)) - } } fn matches_domain(resource_address: &str, domain: &DomainName) -> bool { @@ -816,21 +535,23 @@ fn matches_domain(resource_address: &str, domain: &DomainName) -> bool { name == resource_address } -fn is_subdomain(name: &str, record: &str) -> bool { - if name == record { - return true; - } - let Some((first, end)) = record.split_once('.') else { - return false; - }; - match first { - "*" => name.ends_with(end) && name.strip_suffix(end).is_some_and(|n| n.ends_with('.')), - "?" => { - name.ends_with(end) - && name - .strip_suffix(end) - .is_some_and(|n| n.ends_with('.') && n.matches('.').count() == 1) - } - _ => false, +pub(crate) fn private_key() -> impl Strategy { + any::<[u8; 32]>().prop_map(PrivateKey) +} + +#[derive(Clone, Copy, PartialEq)] +pub(crate) struct PrivateKey([u8; 32]); + +impl From for StaticSecret { + fn from(key: PrivateKey) -> Self { + StaticSecret::from(key.0) + } +} + +impl fmt::Debug for PrivateKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("PrivateKey") + .field(&hex::encode(self.0)) + .finish() } } diff --git a/rust/connlib/tunnel/src/tests/sim_client.rs b/rust/connlib/tunnel/src/tests/sim_client.rs new file mode 100644 index 000000000..69b7e7608 --- /dev/null +++ b/rust/connlib/tunnel/src/tests/sim_client.rs @@ -0,0 +1,639 @@ +use super::{ + reference::{private_key, PrivateKey, ResourceDst}, + sim_net::{any_ip_stack, any_port, host, Host}, + strategies::{system_dns_servers, upstream_dns_servers}, + sut::domain_to_hickory_name, + IcmpIdentifier, IcmpSeq, QueryId, +}; +use crate::{tests::sut::hickory_name_to_domain, ClientState}; +use bimap::BiMap; +use connlib_shared::{ + messages::{ + client::{ResourceDescriptionCidr, ResourceDescriptionDns}, + ClientId, DnsServer, Interface, ResourceId, + }, + proptest::{client_id, domain_name}, + DomainName, +}; +use hickory_proto::{ + op::MessageType, + rr::{rdata, RData, RecordType}, + serialize::binary::BinDecodable as _, +}; +use ip_network_table::IpNetworkTable; +use ip_packet::{IpPacket, MutableIpPacket, Packet as _}; +use itertools::Itertools as _; +use prop::collection; +use proptest::prelude::*; +use snownet::Transmit; +use std::{ + collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}, + net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, + time::Instant, +}; + +/// Simulation state for a particular client. +pub(crate) struct SimClient { + pub(crate) id: ClientId, + + pub(crate) sut: ClientState, + + /// The DNS records created on the client as a result of received DNS responses. + /// + /// This contains results from both, queries to DNS resources and non-resources. + pub(crate) dns_records: HashMap>, + + /// Bi-directional mapping between connlib's sentinel DNS IPs and the effective DNS servers. + pub(crate) dns_by_sentinel: BiMap, + + pub(crate) sent_dns_queries: HashMap>, + pub(crate) received_dns_responses: HashMap>, + + pub(crate) sent_icmp_requests: HashMap<(u16, u16), IpPacket<'static>>, + pub(crate) received_icmp_replies: HashMap<(u16, u16), IpPacket<'static>>, + + buffer: Vec, +} + +impl SimClient { + pub(crate) fn new(id: ClientId, sut: ClientState) -> Self { + Self { + id, + sut, + dns_records: Default::default(), + dns_by_sentinel: Default::default(), + sent_dns_queries: Default::default(), + received_dns_responses: Default::default(), + sent_icmp_requests: Default::default(), + received_icmp_replies: Default::default(), + buffer: vec![0u8; (1 << 16) - 1], + } + } + + /// Returns the _effective_ DNS servers that connlib is using. + pub(crate) fn effective_dns_servers(&self) -> BTreeSet { + self.dns_by_sentinel.right_values().copied().collect() + } + + pub(crate) fn send_dns_query_for( + &mut self, + domain: DomainName, + r_type: RecordType, + query_id: u16, + dns_server: SocketAddr, + now: Instant, + ) -> Option> { + let dns_server = *self + .dns_by_sentinel + .get_by_right(&dns_server) + .expect("to have a sentinel DNS server for the sampled one"); + + let name = domain_to_hickory_name(domain); + + let src = self + .sut + .tunnel_ip_for(dns_server) + .expect("tunnel should be initialised"); + + let packet = ip_packet::make::dns_query( + name, + r_type, + SocketAddr::new(src, 9999), // An application would pick a random source port that is free. + SocketAddr::new(dns_server, 53), + query_id, + ); + + self.encapsulate(packet, now) + } + + pub(crate) fn encapsulate( + &mut self, + packet: MutableIpPacket<'static>, + now: Instant, + ) -> Option> { + { + let packet = packet.to_owned().into_immutable(); + + if let Some(icmp) = packet.as_icmp() { + let echo_request = icmp.as_echo_request().expect("to be echo request"); + + self.sent_icmp_requests + .insert((echo_request.sequence(), echo_request.identifier()), packet); + } + } + + { + let packet = packet.to_owned().into_immutable(); + + if let Some(udp) = packet.as_udp() { + if let Ok(message) = hickory_proto::op::Message::from_bytes(udp.payload()) { + debug_assert_eq!( + message.message_type(), + MessageType::Query, + "every DNS message sent from the client should be a DNS query" + ); + + self.sent_dns_queries.insert(message.id(), packet); + } + } + } + + Some(self.sut.encapsulate(packet, now)?.into_owned()) + } + + pub(crate) fn handle_packet( + &mut self, + payload: &[u8], + src: SocketAddr, + dst: SocketAddr, + now: Instant, + ) { + let Some(packet) = self + .sut + .decapsulate(dst, src, payload, now, &mut self.buffer) + else { + return; + }; + let packet = packet.to_owned(); + + self.on_received_packet(packet); + } + + /// Process an IP packet received on the client. + pub(crate) fn on_received_packet(&mut self, packet: IpPacket<'_>) { + if let Some(icmp) = packet.as_icmp() { + let echo_reply = icmp.as_echo_reply().expect("to be echo reply"); + + self.received_icmp_replies.insert( + (echo_reply.sequence(), echo_reply.identifier()), + packet.to_owned(), + ); + + return; + }; + + if let Some(udp) = packet.as_udp() { + if udp.get_source() == 53 { + let mut message = hickory_proto::op::Message::from_bytes(udp.payload()) + .expect("ip packets on port 53 to be DNS packets"); + + self.received_dns_responses + .insert(message.id(), packet.to_owned()); + + for record in message.take_answers().into_iter() { + let domain = hickory_name_to_domain(record.name().clone()); + + let ip = match record.data() { + Some(RData::A(rdata::A(ip4))) => IpAddr::from(*ip4), + Some(RData::AAAA(rdata::AAAA(ip6))) => IpAddr::from(*ip6), + unhandled => { + panic!("Unexpected record data: {unhandled:?}") + } + }; + + self.dns_records.entry(domain).or_default().push(ip); + } + + // Ensure all IPs are always sorted. + for ips in self.dns_records.values_mut() { + ips.sort() + } + + return; + } + } + + unimplemented!("Unhandled packet") + } +} + +/// Reference state for a particular client. +#[derive(Debug, Clone)] +pub struct RefClient { + pub(crate) id: ClientId, + pub(crate) key: PrivateKey, + pub(crate) known_hosts: HashMap>, + pub(crate) tunnel_ip4: Ipv4Addr, + pub(crate) tunnel_ip6: Ipv6Addr, + + /// The DNS resolvers configured on the client outside of connlib. + pub(crate) system_dns_resolvers: Vec, + /// The upstream DNS resolvers configured in the portal. + pub(crate) upstream_dns_resolvers: Vec, + + /// The CIDR resources the client is aware of. + pub(crate) cidr_resources: IpNetworkTable, + /// The DNS resources the client is aware of. + pub(crate) dns_resources: BTreeMap, + + /// The client's DNS records. + /// + /// The IPs assigned to a domain by connlib are an implementation detail that we don't want to model in these tests. + /// Instead, we just remember what _kind_ of records we resolved to be able to sample a matching src IP. + pub(crate) dns_records: BTreeMap>, + + /// The CIDR resources the client is connected to. + pub(crate) connected_cidr_resources: HashSet, + + /// The DNS resources the client is connected to. + pub(crate) connected_dns_resources: HashSet<(ResourceId, DomainName)>, + + /// The expected ICMP handshakes. + pub(crate) expected_icmp_handshakes: VecDeque<(ResourceDst, IcmpSeq, IcmpIdentifier)>, + /// The expected DNS handshakes. + pub(crate) expected_dns_handshakes: VecDeque, +} + +impl RefClient { + /// Initialize the [`ClientState`]. + /// + /// This simulates receiving the `init` message from the portal. + pub(crate) fn init(self) -> SimClient { + let mut client_state = ClientState::new(self.key, self.known_hosts); + let _ = client_state.update_interface_config(Interface { + ipv4: self.tunnel_ip4, + ipv6: self.tunnel_ip6, + upstream_dns: self.upstream_dns_resolvers, + }); + let _ = client_state.update_system_resolvers(self.system_dns_resolvers); + + SimClient::new(self.id, client_state) + } + + pub(crate) fn is_tunnel_ip(&self, ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip4) => self.tunnel_ip4 == ip4, + IpAddr::V6(ip6) => self.tunnel_ip6 == ip6, + } + } + + pub(crate) fn tunnel_ip_for(&self, dst: IpAddr) -> IpAddr { + match dst { + IpAddr::V4(_) => self.tunnel_ip4.into(), + IpAddr::V6(_) => self.tunnel_ip6.into(), + } + } + + #[tracing::instrument(level = "debug", skip_all, fields(dst, resource))] + pub(crate) fn on_icmp_packet_to_cidr( + &mut self, + src: IpAddr, + dst: IpAddr, + seq: u16, + identifier: u16, + ) { + tracing::Span::current().record("dst", tracing::field::display(dst)); + + // Second, if we are not yet connected, check if we have a resource for this IP. + let Some((_, resource)) = self.cidr_resources.longest_match(dst) else { + tracing::debug!("No resource corresponds to IP"); + return; + }; + tracing::Span::current().record("resource", tracing::field::display(resource.id)); + + 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 + .push_back((ResourceDst::Cidr(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.connected_cidr_resources.insert(resource.id); + } + + #[tracing::instrument(level = "debug", skip_all, fields(dst, resource))] + pub(crate) fn on_icmp_packet_to_dns( + &mut self, + src: IpAddr, + dst: DomainName, + seq: u16, + identifier: u16, + ) { + tracing::Span::current().record("dst", tracing::field::display(&dst)); + + let Some(resource) = self.dns_resource_by_domain(&dst) else { + tracing::debug!("No resource corresponds to IP"); + return; + }; + + tracing::Span::current().record("resource", tracing::field::display(resource)); + + if self + .connected_dns_resources + .contains(&(resource, dst.clone())) + && self.is_tunnel_ip(src) + { + tracing::debug!("Connected to DNS resource, expecting packet to be routed"); + self.expected_icmp_handshakes + .push_back((ResourceDst::Dns(dst), seq, identifier)); + return; + } + + debug_assert!( + self.dns_records.iter().any(|(name, _)| name == &dst), + "Should only sample ICMPs to domains that we resolved" + ); + + tracing::debug!("Not connected to resource, expecting to trigger connection intent"); + self.connected_dns_resources.insert((resource, dst)); + } + + pub(crate) fn ipv4_cidr_resource_dsts(&self) -> Vec { + let mut ips = vec![]; + + // This is an imperative loop on purpose because `ip-network` appears to have a bug with its `size_hint` and thus `.extend` does not work reliably? + for (network, _) in self.cidr_resources.iter_ipv4() { + if network.netmask() == 31 || network.netmask() == 32 { + ips.push(network.network_address()); + } else { + for ip in network.hosts() { + ips.push(ip) + } + } + } + + ips + } + + pub(crate) fn ipv6_cidr_resource_dsts(&self) -> Vec { + let mut ips = vec![]; + + // This is an imperative loop on purpose because `ip-network` appears to have a bug with its `size_hint` and thus `.extend` does not work reliably? + for (network, _) in self.cidr_resources.iter_ipv6() { + if network.netmask() == 127 || network.netmask() == 128 { + ips.push(network.network_address()); + } else { + for ip in network + .subnets_with_prefix(128) + .map(|i| i.network_address()) + { + ips.push(ip) + } + } + } + + ips + } + + pub(crate) fn is_connected_to_cidr(&self, id: ResourceId) -> bool { + self.connected_cidr_resources.contains(&id) + } + + pub(crate) fn is_known_host(&self, name: &str) -> bool { + self.known_hosts.contains_key(name) + } + + fn dns_resource_by_domain(&self, domain: &DomainName) -> Option { + self.dns_resources + .values() + .filter(|r| is_subdomain(&domain.to_string(), &r.address)) + .sorted_by_key(|r| r.address.len()) + .rev() + .map(|r| r.id) + .next() + } + + fn resolved_domains(&self) -> impl Iterator)> + '_ { + self.dns_records + .iter() + .filter(|(domain, _)| self.dns_resource_by_domain(domain).is_some()) + .map(|(domain, ips)| (domain.clone(), ips.clone())) + } + + /// 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)| { + existing_seq != seq && existing_identifer != identifier + }) + } + + pub(crate) fn resolved_v4_domains(&self) -> Vec { + self.resolved_domains() + .filter_map(|(domain, records)| { + records + .iter() + .any(|r| matches!(r, RecordType::A)) + .then_some(domain) + }) + .collect() + } + + pub(crate) fn resolved_v6_domains(&self) -> Vec { + self.resolved_domains() + .filter_map(|(domain, records)| { + records + .iter() + .any(|r| matches!(r, RecordType::AAAA)) + .then_some(domain) + }) + .collect() + } + + /// Returns the DNS servers that we expect connlib to use. + /// + /// If there are upstream DNS servers configured in the portal, it should use those. + /// Otherwise it should use whatever was configured on the system prior to connlib starting. + pub(crate) fn expected_dns_servers(&self) -> BTreeSet { + if !self.upstream_dns_resolvers.is_empty() { + return self + .upstream_dns_resolvers + .iter() + .map(|s| s.address()) + .collect(); + } + + self.system_dns_resolvers + .iter() + .map(|ip| SocketAddr::new(*ip, 53)) + .collect() + } + + pub(crate) fn v4_dns_servers(&self) -> Vec { + self.expected_dns_servers() + .into_iter() + .filter_map(|s| match s { + SocketAddr::V4(v4) => Some(v4), + SocketAddr::V6(_) => None, + }) + .collect() + } + + pub(crate) fn v6_dns_servers(&self) -> Vec { + self.expected_dns_servers() + .into_iter() + .filter_map(|s| match s { + SocketAddr::V6(v6) => Some(v6), + SocketAddr::V4(_) => None, + }) + .collect() + } + + pub(crate) fn cidr_resource_by_ip(&self, ip: IpAddr) -> Option { + self.cidr_resources.longest_match(ip).map(|(_, r)| r.id) + } + + pub(crate) fn resolved_ip4_for_non_resources( + &self, + global_dns_records: &BTreeMap>, + ) -> Vec { + self.resolved_ips_for_non_resources(global_dns_records) + .filter_map(|ip| match ip { + IpAddr::V4(v4) => Some(v4), + IpAddr::V6(_) => None, + }) + .collect() + } + + pub(crate) fn resolved_ip6_for_non_resources( + &self, + global_dns_records: &BTreeMap>, + ) -> Vec { + self.resolved_ips_for_non_resources(global_dns_records) + .filter_map(|ip| match ip { + IpAddr::V6(v6) => Some(v6), + IpAddr::V4(_) => None, + }) + .collect() + } + + fn resolved_ips_for_non_resources<'a>( + &'a self, + global_dns_records: &'a BTreeMap>, + ) -> impl Iterator + 'a { + self.dns_records + .iter() + .filter_map(|(domain, _)| { + self.dns_resource_by_domain(domain) + .is_none() + .then_some(global_dns_records.get(domain)) + }) + .flatten() + .flatten() + .copied() + } + + /// Returns the CIDR resource we will forward the DNS query for the given name to. + /// + /// DNS servers may be resources, in which case queries that need to be forwarded actually need to be encapsulated. + pub(crate) fn dns_query_via_cidr_resource( + &self, + dns_server: IpAddr, + domain: &DomainName, + ) -> Option { + // If we are querying a DNS resource, we will issue a connection intent to the DNS resource, not the CIDR resource. + if self.dns_resource_by_domain(domain).is_some() { + return None; + } + + self.cidr_resource_by_ip(dns_server) + } + + pub(crate) fn all_resources(&self) -> Vec { + let cidr_resources = self.cidr_resources.iter().map(|(_, r)| r.id); + let dns_resources = self.dns_resources.keys().copied(); + + Vec::from_iter(cidr_resources.chain(dns_resources)) + } +} + +fn is_subdomain(name: &str, record: &str) -> bool { + if name == record { + return true; + } + let Some((first, end)) = record.split_once('.') else { + return false; + }; + match first { + "*" => name.ends_with(end) && name.strip_suffix(end).is_some_and(|n| n.ends_with('.')), + "?" => { + name.ends_with(end) + && name + .strip_suffix(end) + .is_some_and(|n| n.ends_with('.') && n.matches('.').count() == 1) + } + _ => false, + } +} + +pub(crate) fn ref_client_host( + tunnel_ip4s: &mut impl Iterator, + tunnel_ip6s: &mut impl Iterator, +) -> impl Strategy> { + host( + any_ip_stack(), + any_port(), + ref_client(tunnel_ip4s, tunnel_ip6s), + ) + .prop_filter("at least one DNS server needs to be reachable", |host| { + // TODO: PRODUCTION CODE DOES NOT HANDLE THIS! + + let upstream_dns_resolvers = &host.inner().upstream_dns_resolvers; + let system_dns = &host.inner().system_dns_resolvers; + + if !upstream_dns_resolvers.is_empty() { + if host.ip4.is_none() && upstream_dns_resolvers.iter().all(|s| s.ip().is_ipv4()) { + return false; + } + if host.ip6.is_none() && upstream_dns_resolvers.iter().all(|s| s.ip().is_ipv6()) { + return false; + } + + return true; + } + + if host.ip4.is_none() && system_dns.iter().all(|s| s.is_ipv4()) { + return false; + } + if host.ip6.is_none() && system_dns.iter().all(|s| s.is_ipv6()) { + return false; + } + + true + }) +} + +fn ref_client( + tunnel_ip4s: &mut impl Iterator, + tunnel_ip6s: &mut impl Iterator, +) -> impl Strategy { + let tunnel_ip4 = tunnel_ip4s.next().unwrap(); + let tunnel_ip6 = tunnel_ip6s.next().unwrap(); + + ( + client_id(), + private_key(), + known_hosts(), + system_dns_servers(), + upstream_dns_servers(), + ) + .prop_map( + move |(id, key, known_hosts, system_dns_resolvers, upstream_dns_resolvers)| RefClient { + id, + key, + known_hosts, + tunnel_ip4, + tunnel_ip6, + system_dns_resolvers, + upstream_dns_resolvers, + cidr_resources: IpNetworkTable::new(), + dns_resources: Default::default(), + dns_records: Default::default(), + connected_cidr_resources: Default::default(), + connected_dns_resources: Default::default(), + expected_icmp_handshakes: Default::default(), + expected_dns_handshakes: Default::default(), + }, + ) +} + +fn known_hosts() -> impl Strategy>> { + collection::hash_map( + domain_name(2..4).prop_map(|d| d.parse().unwrap()), + collection::vec(any::(), 1..6), + 0..15, + ) +} diff --git a/rust/connlib/tunnel/src/tests/sim_gateway.rs b/rust/connlib/tunnel/src/tests/sim_gateway.rs new file mode 100644 index 000000000..a5be49721 --- /dev/null +++ b/rust/connlib/tunnel/src/tests/sim_gateway.rs @@ -0,0 +1,110 @@ +use super::{ + reference::{private_key, PrivateKey}, + 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 ip_packet::IpPacket; +use proptest::prelude::*; +use snownet::Transmit; +use std::{ + collections::{BTreeMap, HashSet, VecDeque}, + net::{IpAddr, SocketAddr}, + time::Instant, +}; + +/// Simulation state for a particular client. +pub(crate) struct SimGateway { + pub(crate) id: GatewayId, + pub(crate) sut: GatewayState, + + pub(crate) received_icmp_requests: VecDeque>, + + buffer: Vec, +} + +impl SimGateway { + pub(crate) fn new(id: GatewayId, sut: GatewayState) -> Self { + Self { + id, + sut, + received_icmp_requests: Default::default(), + buffer: vec![0u8; (1 << 16) - 1], + } + } + + pub(crate) fn handle_packet( + &mut self, + global_dns_records: &BTreeMap>, + payload: &[u8], + src: SocketAddr, + dst: SocketAddr, + now: Instant, + ) -> Option> { + let packet = self + .sut + .decapsulate(dst, src, payload, now, &mut self.buffer)? + .to_owned(); + + self.on_received_packet(global_dns_records, packet, now) + } + + /// Process an IP packet received on the gateway. + fn on_received_packet( + &mut self, + global_dns_records: &BTreeMap>, + packet: IpPacket<'_>, + now: Instant, + ) -> Option> { + let packet = packet.to_owned(); + + if packet.as_icmp().is_some() { + self.received_icmp_requests.push_back(packet.clone()); + + let echo_response = ip_packet::make::icmp_response_packet(packet); + let transmit = self.sut.encapsulate(echo_response, now)?.into_owned(); + + return Some(transmit); + } + + if packet.as_udp().is_some() { + let response = ip_packet::make::dns_ok_response(packet, |name| { + global_dns_records + .get(&hickory_name_to_domain(name.clone())) + .cloned() + .into_iter() + .flatten() + }); + + let transmit = self.sut.encapsulate(response, now)?.into_owned(); + + return Some(transmit); + } + + panic!("Unhandled packet") + } +} + +/// Reference state for a particular gateway. +#[derive(Debug, Clone)] +pub struct RefGateway { + pub(crate) id: GatewayId, + pub(crate) key: PrivateKey, +} + +impl RefGateway { + /// Initialize the [`GatewayState`]. + /// + /// This simulates receiving the `init` message from the portal. + pub(crate) fn init(self) -> SimGateway { + SimGateway::new(self.id, GatewayState::new(self.key)) + } +} + +pub(crate) fn ref_gateway_host() -> impl Strategy> { + host(any_ip_stack(), any_port(), ref_gateway()) +} + +fn ref_gateway() -> impl Strategy { + (gateway_id(), private_key()).prop_map(move |(id, key)| RefGateway { id, key }) +} diff --git a/rust/connlib/tunnel/src/tests/sim_net.rs b/rust/connlib/tunnel/src/tests/sim_net.rs index b861a55fc..afe0d495a 100644 --- a/rust/connlib/tunnel/src/tests/sim_net.rs +++ b/rust/connlib/tunnel/src/tests/sim_net.rs @@ -1,16 +1,15 @@ -use super::{sim_node::SimNode, sim_relay::SimRelay}; -use crate::{ClientState, GatewayState}; use connlib_shared::messages::{ClientId, GatewayId, RelayId}; -use firezone_relay::AddressFamily; -use ip_network::IpNetwork; +use firezone_relay::{AddressFamily, IpStack}; +use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; use ip_network_table::IpNetworkTable; -use ip_packet::MutableIpPacket; -use rand::rngs::StdRng; -use snownet::Transmit; +use itertools::Itertools as _; +use prop::sample; +use proptest::prelude::*; use std::{ collections::HashSet, + fmt, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, - time::Instant, + num::NonZeroU16, }; use tracing::Span; @@ -98,39 +97,17 @@ impl Host { self.allocate_port(port, AddressFamily::V6); } } - - /// Sets the `src` of the given [`Transmit`] in case it is missing. - /// - /// The `src` of a [`Transmit`] is empty if we want to send if via the default interface. - /// In production, the kernel does this for us. - /// In this test, we need to always set a `src` so that the remote peer knows where the packet is coming from. - fn set_transmit_src(&self, transmit: Transmit<'static>) -> Option> { - if transmit.src.is_some() { - return Some(transmit); - } - - let Some(src) = self.sending_socket_for(transmit.dst.ip()) else { - tracing::debug!(dst = %transmit.dst, "No socket"); - - return None; - }; - - Some(Transmit { - src: Some(src), - ..transmit - }) - } } impl Host where T: Clone, { - pub(crate) fn map( + pub(crate) fn map( &self, - f: impl FnOnce(T, Option, Option) -> S, + f: impl FnOnce(T, Option, Option) -> U, span: Span, - ) -> Host { + ) -> Host { Host { inner: f(self.inner.clone(), self.ip4, self.ip6), ip4: self.ip4, @@ -143,78 +120,6 @@ where } } -#[allow(private_bounds)] -impl Host -where - T: PollTransmit, -{ - pub(crate) fn poll_transmit(&mut self) -> Option> { - let _guard = self.span.enter(); - let transmit = self.span.in_scope(|| self.inner.poll_transmit())?; - - self.set_transmit_src(transmit) - } -} - -#[allow(private_bounds)] -impl Host -where - T: Encapsulate, -{ - pub(crate) fn encapsulate( - &mut self, - packet: MutableIpPacket<'_>, - now: Instant, - ) -> Option> { - let _guard = self.span.enter(); - - let transmit = self - .span - .in_scope(|| self.inner.encapsulate(packet, now))? - .into_owned(); - - self.set_transmit_src(transmit) - } -} - -trait Encapsulate { - fn encapsulate(&mut self, packet: MutableIpPacket<'_>, now: Instant) -> Option>; -} - -impl Encapsulate for SimNode { - fn encapsulate(&mut self, packet: MutableIpPacket<'_>, now: Instant) -> Option> { - self.state.encapsulate(packet, now) - } -} - -impl Encapsulate for SimNode { - fn encapsulate(&mut self, packet: MutableIpPacket<'_>, now: Instant) -> Option> { - self.state.encapsulate(packet, now) - } -} - -trait PollTransmit { - fn poll_transmit(&mut self) -> Option>; -} - -impl PollTransmit for SimNode { - fn poll_transmit(&mut self) -> Option> { - self.state.poll_transmit() - } -} - -impl PollTransmit for SimNode { - fn poll_transmit(&mut self) -> Option> { - self.state.poll_transmit() - } -} - -impl PollTransmit for SimRelay> { - fn poll_transmit(&mut self) -> Option> { - None - } -} - #[derive(Debug, Clone)] pub(crate) struct RoutingTable { routes: IpNetworkTable, @@ -323,3 +228,62 @@ impl From for HostId { Self::Client(v) } } + +pub(crate) fn host( + socket_ips: impl Strategy, + default_port: impl Strategy, + state: impl Strategy, +) -> impl Strategy> +where + T: fmt::Debug, +{ + (state, socket_ips, default_port).prop_map(move |(state, ip_stack, port)| { + let mut host = Host::new(state); + host.update_interface(ip_stack.as_v4().copied(), ip_stack.as_v6().copied(), port); + + host + }) +} + +pub(crate) fn any_port() -> impl Strategy { + any::().prop_map(|v| v.into()) +} + +pub(crate) fn any_ip_stack() -> impl Strategy { + prop_oneof![ + host_ip4s().prop_map(IpStack::Ip4), + host_ip6s().prop_map(IpStack::Ip6), + dual_ip_stack() + ] +} + +pub(crate) fn dual_ip_stack() -> impl Strategy { + (host_ip4s(), host_ip6s()).prop_map(|(ip4, ip6)| IpStack::Dual { ip4, ip6 }) +} + +/// A [`Strategy`] of [`Ipv4Addr`]s used for routing packets between hosts within our test. +/// +/// This uses the `TEST-NET-3` (`203.0.113.0/24`) address space reserved for documentation and examples in [RFC5737](https://datatracker.ietf.org/doc/html/rfc5737). +pub(crate) fn host_ip4s() -> impl Strategy { + let ips = Ipv4Network::new(Ipv4Addr::new(203, 0, 113, 0), 24) + .unwrap() + .hosts() + .take(100) + .collect_vec(); + + sample::select(ips) +} + +/// A [`Strategy`] of [`Ipv6Addr`]s used for routing packets between hosts within our test. +/// +/// This uses the `2001:DB8::/32` address space reserved for documentation and examples in [RFC3849](https://datatracker.ietf.org/doc/html/rfc3849). +pub(crate) fn host_ip6s() -> impl Strategy { + let ips = Ipv6Network::new(Ipv6Addr::new(0x2001, 0xDB80, 0, 0, 0, 0, 0, 0), 32) + .unwrap() + .subnets_with_prefix(128) + .map(|n| n.network_address()) + .take(100) + .collect_vec(); + + sample::select(ips) +} diff --git a/rust/connlib/tunnel/src/tests/sim_node.rs b/rust/connlib/tunnel/src/tests/sim_node.rs deleted file mode 100644 index 4e0111aef..000000000 --- a/rust/connlib/tunnel/src/tests/sim_node.rs +++ /dev/null @@ -1,213 +0,0 @@ -use super::{ - sim_net::Host, - sim_relay::SimRelay, - strategies::{host_ip4s, host_ip6s}, -}; -use crate::{ClientState, GatewayState}; -use connlib_shared::{ - messages::{ClientId, DnsServer, GatewayId, Interface, RelayId}, - proptest::domain_name, - StaticSecret, -}; -use firezone_relay::IpStack; -use ip_network::{Ipv4Network, Ipv6Network}; -use proptest::{collection, prelude::*}; -use rand::rngs::StdRng; -use std::{ - collections::{HashMap, HashSet}, - fmt, - net::{IpAddr, Ipv4Addr, Ipv6Addr}, - time::Instant, -}; - -#[derive(Clone, Debug)] -pub(crate) struct SimNode { - pub(crate) id: ID, - pub(crate) state: S, - - pub(crate) tunnel_ip4: Ipv4Addr, - pub(crate) tunnel_ip6: Ipv6Addr, -} - -impl SimNode { - pub(crate) fn new(id: ID, state: S, tunnel_ip4: Ipv4Addr, tunnel_ip6: Ipv6Addr) -> Self { - Self { - id, - state, - tunnel_ip4, - tunnel_ip6, - } - } -} - -impl SimNode -where - ID: Copy, - S: Clone, -{ - pub(crate) fn map(&self, f: impl FnOnce(S) -> T) -> SimNode { - SimNode { - id: self.id, - state: f(self.state.clone()), - tunnel_ip4: self.tunnel_ip4, - tunnel_ip6: self.tunnel_ip6, - } - } -} - -impl SimNode { - pub(crate) fn init_relays<'a>( - &mut self, - relays: impl Iterator< - Item = ( - &'a RelayId, - &'a Host>>, - ), - >, - now: Instant, - ) { - self.state.update_relays( - HashSet::default(), - HashSet::from_iter(relays.map(|(id, r)| { - let (socket, username, password, realm) = r.inner().explode("client"); - - (*id, socket, username, password, realm) - })), - now, - ) - } - - pub(crate) fn update_upstream_dns(&mut self, upstream_dns_resolvers: Vec) { - let _ = self.state.update_interface_config(Interface { - ipv4: self.tunnel_ip4, - ipv6: self.tunnel_ip6, - upstream_dns: upstream_dns_resolvers, - }); - } -} - -impl SimNode { - pub(crate) fn init_relays<'a>( - &mut self, - relays: impl Iterator< - Item = ( - &'a RelayId, - &'a Host>>, - ), - >, - now: Instant, - ) { - self.state.update_relays( - HashSet::default(), - HashSet::from_iter(relays.map(|(id, r)| { - let (socket, username, password, realm) = r.inner().explode("gateway"); - - (*id, socket, username, password, realm) - })), - now, - ) - } -} - -impl SimNode { - pub(crate) fn tunnel_ip(&self, dst: impl Into) -> IpAddr { - match dst.into() { - IpAddr::V4(_) => IpAddr::from(self.tunnel_ip4), - IpAddr::V6(_) => IpAddr::from(self.tunnel_ip6), - } - } - - pub(crate) fn is_tunnel_ip(&self, ip: IpAddr) -> bool { - self.tunnel_ip(ip) == ip - } -} - -#[derive(Clone, Copy, PartialEq)] -pub(crate) struct PrivateKey([u8; 32]); - -impl From for StaticSecret { - fn from(key: PrivateKey) -> Self { - StaticSecret::from(key.0) - } -} - -impl fmt::Debug for PrivateKey { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("PrivateKey") - .field(&hex::encode(self.0)) - .finish() - } -} - -pub(crate) fn sim_node_prototype( - id: impl Strategy, - state: impl Strategy, - tunnel_ip4s: &mut impl Iterator, - tunnel_ip6s: &mut impl Iterator, -) -> impl Strategy>> -where - ID: fmt::Debug, - S: fmt::Debug, -{ - let socket_ips = prop_oneof![ - host_ip4s().prop_map(IpStack::Ip4), - host_ip6s().prop_map(IpStack::Ip6), - (host_ip4s(), host_ip6s()).prop_map(|(ip4, ip6)| IpStack::Dual { ip4, ip6 }) - ]; - - let tunnel_ip4 = tunnel_ip4s.next().unwrap(); - let tunnel_ip6 = tunnel_ip6s.next().unwrap(); - - ( - id, - state, - socket_ips, - any::().prop_filter("port must not be 0", |p| *p != 0), - ) - .prop_map(move |(id, state, ip_stack, port)| { - let mut host = Host::new(SimNode::new(id, state, tunnel_ip4, tunnel_ip6)); - host.update_interface(ip_stack.as_v4().copied(), ip_stack.as_v6().copied(), port); - - host - }) -} - -/// An [`Iterator`] over the possible IPv4 addresses of a tunnel interface. -/// -/// We use the CG-NAT range for IPv4. -/// See . -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()) -} - -fn private_key() -> impl Strategy { - any::<[u8; 32]>().prop_map(PrivateKey) -} - -pub(crate) fn gateway_state() -> impl Strategy { - private_key() -} - -pub(crate) fn client_state() -> impl Strategy>)> { - (private_key(), known_hosts()) -} - -pub(crate) fn known_hosts() -> impl Strategy>> { - collection::hash_map( - domain_name(2..4).prop_map(|d| d.parse().unwrap()), - collection::vec(any::(), 1..6), - 0..15, - ) -} diff --git a/rust/connlib/tunnel/src/tests/sim_relay.rs b/rust/connlib/tunnel/src/tests/sim_relay.rs index 75d77207d..8193c4a42 100644 --- a/rust/connlib/tunnel/src/tests/sim_relay.rs +++ b/rust/connlib/tunnel/src/tests/sim_relay.rs @@ -1,64 +1,54 @@ -use super::sim_net::Host; -use super::strategies::{host_ip4s, host_ip6s}; +use super::sim_net::{dual_ip_stack, host, Host}; +use connlib_shared::messages::RelayId; use firezone_relay::{AddressFamily, AllocationPort, ClientSocket, IpStack, PeerSocket}; use proptest::prelude::*; use rand::rngs::StdRng; +use secrecy::SecretString; use snownet::{RelaySocket, Transmit}; use std::{ borrow::Cow, collections::HashSet, - net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6}, + net::{SocketAddr, SocketAddrV4, SocketAddrV6}, time::{Duration, Instant, SystemTime}, }; -#[derive(Clone, derivative::Derivative)] -#[derivative(Debug)] -pub(crate) struct SimRelay { - pub(crate) state: S, - - ip_stack: firezone_relay::IpStack, +pub(crate) struct SimRelay { + pub(crate) sut: firezone_relay::Server, pub(crate) allocations: HashSet<(AddressFamily, AllocationPort)>, - - #[derivative(Debug = "ignore")] buffer: Vec, } -impl SimRelay { - pub(crate) fn new(state: S, ip_stack: firezone_relay::IpStack) -> Self { +pub(crate) fn map_explode<'a>( + relays: impl Iterator)> + 'a, + username: &'static str, +) -> impl Iterator + 'a { + relays.map(move |(id, r)| { + let (socket, username, password, realm) = r.inner().explode( + username, + r.inner().sut.auth_secret(), + r.inner().sut.public_address(), + ); + + (*id, socket, username, password, realm) + }) +} + +impl SimRelay { + pub(crate) fn new(sut: firezone_relay::Server) -> Self { Self { - state, - ip_stack, + sut, allocations: Default::default(), buffer: vec![0u8; (1 << 16) - 1], } } - pub(crate) fn ip4(&self) -> Option { - self.ip_stack.as_v4().copied().map(|i| i.into()) - } - - pub(crate) fn ip6(&self) -> Option { - self.ip_stack.as_v6().copied().map(|i| i.into()) - } -} - -impl SimRelay -where - S: Copy, -{ - pub(crate) fn map(&self, f: impl FnOnce(S) -> T) -> SimRelay { - SimRelay { - state: f(self.state), - allocations: self.allocations.clone(), - buffer: self.buffer.clone(), - ip_stack: self.ip_stack, - } - } -} - -impl SimRelay> { - pub(crate) fn explode(&self, username: &str) -> (RelaySocket, String, String, String) { - let relay_socket = match self.ip_stack { + fn explode( + &self, + username: &str, + auth_secret: &SecretString, + public_address: IpStack, + ) -> (RelaySocket, String, String, String) { + let relay_socket = match public_address { firezone_relay::IpStack::Ip4(ip4) => RelaySocket::V4(SocketAddrV4::new(ip4, 3478)), firezone_relay::IpStack::Ip6(ip6) => { RelaySocket::V6(SocketAddrV6::new(ip6, 3478, 0, 0)) @@ -69,15 +59,19 @@ impl SimRelay> { }, }; - let (username, password) = self.make_credentials(username); + let (username, password) = self.make_credentials(username, auth_secret); (relay_socket, username, password, "firezone".to_owned()) } - fn matching_listen_socket(&self, other: SocketAddr) -> Option { + fn matching_listen_socket( + &self, + other: SocketAddr, + public_address: IpStack, + ) -> Option { match other { - SocketAddr::V4(_) => Some(SocketAddr::new((*self.ip_stack.as_v4()?).into(), 3478)), - SocketAddr::V6(_) => Some(SocketAddr::new((*self.ip_stack.as_v6()?).into(), 3478)), + SocketAddr::V4(_) => Some(SocketAddr::new((*public_address.as_v4()?).into(), 3478)), + SocketAddr::V6(_) => Some(SocketAddr::new((*public_address.as_v6()?).into(), 3478)), } } @@ -88,7 +82,10 @@ impl SimRelay> { dst: SocketAddr, now: Instant, ) -> Option> { - if self.matching_listen_socket(dst).is_some_and(|s| s == dst) { + if self + .matching_listen_socket(dst, self.sut.public_address()) + .is_some_and(|s| s == dst) + { return self.handle_client_input(payload, ClientSocket::new(sender), now); } @@ -105,7 +102,7 @@ impl SimRelay> { client: ClientSocket, now: Instant, ) -> Option> { - let (port, peer) = self.state.handle_client_input(payload, client, now)?; + let (port, peer) = self.sut.handle_client_input(payload, client, now)?; let payload = &payload[4..]; @@ -120,7 +117,9 @@ impl SimRelay> { "IPv4 allocation to be present if we want to send to an IPv4 socket" ); - self.ip4().expect("listen on IPv4 if we have an allocation") + self.sut + .public_ip4() + .expect("listen on IPv4 if we have an allocation") } SocketAddr::V6(_) => { assert!( @@ -128,7 +127,9 @@ impl SimRelay> { "IPv6 allocation to be present if we want to send to an IPv6 socket" ); - self.ip6().expect("listen on IPv6 if we have an allocation") + self.sut + .public_ip6() + .expect("listen on IPv6 if we have an allocation") } }; @@ -148,7 +149,7 @@ impl SimRelay> { peer: PeerSocket, port: AllocationPort, ) -> Option> { - let (client, channel) = self.state.handle_peer_traffic(payload, peer, port)?; + let (client, channel) = self.sut.handle_peer_traffic(payload, peer, port)?; let full_length = firezone_relay::ChannelData::encode_header_to_slice( channel, @@ -158,7 +159,9 @@ impl SimRelay> { self.buffer[4..full_length].copy_from_slice(payload); let receiving_socket = client.into_socket(); - let sending_socket = self.matching_listen_socket(receiving_socket).unwrap(); + let sending_socket = self + .matching_listen_socket(receiving_socket, self.sut.public_address()) + .unwrap(); Some(Transmit { src: Some(sending_socket), @@ -167,7 +170,7 @@ impl SimRelay> { }) } - fn make_credentials(&self, username: &str) -> (String, String) { + fn make_credentials(&self, username: &str, auth_secret: &SecretString) -> (String, String) { let expiry = SystemTime::now() + Duration::from_secs(60); let secs = expiry @@ -175,21 +178,16 @@ impl SimRelay> { .expect("expiry must be later than UNIX_EPOCH") .as_secs(); - let password = - firezone_relay::auth::generate_password(self.state.auth_secret(), expiry, username); + let password = firezone_relay::auth::generate_password(auth_secret, expiry, username); (format!("{secs}:{username}"), password) } } -pub(crate) fn sim_relay_prototype() -> impl Strategy>> { - // For this test, our relays always run in dual-stack mode to ensure connectivity! - let socket_ips = (host_ip4s(), host_ip6s()).prop_map(|(ip4, ip6)| IpStack::Dual { ip4, ip6 }); - - (any::(), socket_ips).prop_map(move |(seed, ip_stack)| { - let mut host = Host::new(SimRelay::new(seed, ip_stack)); - host.update_interface(ip_stack.as_v4().copied(), ip_stack.as_v6().copied(), 3478); - - host - }) +pub(crate) fn relay_prototype() -> impl Strategy> { + host( + dual_ip_stack(), // For this test, our relays always run in dual-stack mode to ensure connectivity! + Just(3478), + any::(), + ) } diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs index 046234fea..0ecc228a5 100644 --- a/rust/connlib/tunnel/src/tests/strategies.rs +++ b/rust/connlib/tunnel/src/tests/strategies.rs @@ -1,7 +1,6 @@ use connlib_shared::{messages::DnsServer, proptest::domain_name, DomainName}; use ip_network::{Ipv4Network, Ipv6Network}; -use itertools::Itertools as _; -use proptest::{collection, prelude::*, sample}; +use proptest::{collection, prelude::*}; use std::{ collections::{BTreeMap, HashMap, HashSet}, net::{IpAddr, Ipv4Addr, Ipv6Addr}, @@ -51,33 +50,6 @@ pub(crate) fn upstream_dns_servers() -> impl Strategy> { ] } -/// A [`Strategy`] of [`Ipv4Addr`]s used for routing packets between hosts within our test. -/// -/// This uses the `TEST-NET-3` (`203.0.113.0/24`) address space reserved for documentation and examples in [RFC5737](https://datatracker.ietf.org/doc/html/rfc5737). -pub(crate) fn host_ip4s() -> impl Strategy { - let ips = Ipv4Network::new(Ipv4Addr::new(203, 0, 113, 0), 24) - .unwrap() - .hosts() - .take(100) - .collect_vec(); - - sample::select(ips) -} - -/// A [`Strategy`] of [`Ipv6Addr`]s used for routing packets between hosts within our test. -/// -/// This uses the `2001:DB8::/32` address space reserved for documentation and examples in [RFC3849](https://datatracker.ietf.org/doc/html/rfc3849). -pub(crate) fn host_ip6s() -> impl Strategy { - let ips = Ipv6Network::new(Ipv6Addr::new(0x2001, 0xDB80, 0, 0, 0, 0, 0, 0), 32) - .unwrap() - .subnets_with_prefix(128) - .map(|n| n.network_address()) - .take(100) - .collect_vec(); - - sample::select(ips) -} - pub(crate) fn system_dns_servers() -> impl Strategy> { collection::vec(any::(), 1..4) // Always need at least 1 system DNS server. TODO: Should we test what happens if we don't? } @@ -103,3 +75,23 @@ pub(crate) fn packet_source_v6(client: Ipv6Addr) -> impl Strategy any::() ] } + +/// An [`Iterator`] over the possible IPv4 addresses of a tunnel interface. +/// +/// We use the CG-NAT range for IPv4. +/// See . +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()) +} diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs index f2c9ee7bd..95b262c68 100644 --- a/rust/connlib/tunnel/src/tests/sut.rs +++ b/rust/connlib/tunnel/src/tests/sut.rs @@ -1,15 +1,15 @@ use super::reference::ReferenceState; +use super::sim_client::SimClient; +use super::sim_gateway::SimGateway; use super::sim_net::{Host, HostId, RoutingTable}; -use super::sim_node::SimNode; use super::sim_portal::SimPortal; use super::sim_relay::SimRelay; -use super::QueryId; use crate::tests::assertions::*; +use crate::tests::sim_relay::map_explode; use crate::tests::transition::Transition; -use crate::{dns::DnsQuery, ClientEvent, ClientState, GatewayEvent, GatewayState, Request}; -use bimap::BiMap; +use crate::{dns::DnsQuery, ClientEvent, GatewayEvent, Request}; use chrono::{DateTime, Utc}; -use connlib_shared::messages::RelayId; +use connlib_shared::messages::{Interface, RelayId}; use connlib_shared::{ messages::{ client::{ResourceDescription, ResourceDescriptionCidr, ResourceDescriptionDns}, @@ -19,21 +19,19 @@ use connlib_shared::{ }; use firezone_relay::IpStack; use hickory_proto::{ - op::{MessageType, Query}, - rr::{rdata, RData, Record, RecordType}, - serialize::binary::BinDecodable as _, + op::Query, + rr::{RData, Record, RecordType}, }; use hickory_resolver::lookup::Lookup; use ip_network_table::IpNetworkTable; -use ip_packet::{IpPacket, MutableIpPacket, Packet as _}; use proptest_state_machine::{ReferenceStateMachine, StateMachineTest}; -use rand::{rngs::StdRng, SeedableRng as _}; +use rand::SeedableRng as _; use secrecy::ExposeSecret as _; use snownet::Transmit; -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; use std::{ collections::{HashMap, HashSet, VecDeque}, - net::{IpAddr, SocketAddr}, + net::IpAddr, str::FromStr as _, sync::Arc, time::{Duration, Instant}, @@ -49,26 +47,11 @@ pub(crate) struct TunnelTest { now: Instant, utc_now: DateTime, - client: Host>, - gateway: Host>, - relays: HashMap>>>, + pub(crate) client: Host, + pub(crate) gateway: Host, + relays: HashMap>, portal: SimPortal, - /// The DNS records created on the client as a result of received DNS responses. - /// - /// This contains results from both, queries to DNS resources and non-resources. - pub(crate) client_dns_records: HashMap>, - - /// Bi-directional mapping between connlib's sentinel DNS IPs and the effective DNS servers. - client_dns_by_sentinel: BiMap, - - pub(crate) client_sent_dns_queries: HashMap>, - pub(crate) client_received_dns_responses: HashMap>, - - pub(crate) client_sent_icmp_requests: HashMap<(u16, u16), IpPacket<'static>>, - pub(crate) client_received_icmp_replies: HashMap<(u16, u16), IpPacket<'static>>, - pub(crate) gateway_received_icmp_requests: VecDeque>, - network: RoutingTable, #[allow(dead_code)] @@ -90,12 +73,11 @@ impl StateMachineTest for TunnelTest { .set_default(); // Construct client, gateway and relay from the initial state. - let mut client = ref_state.client.map( - |sim_node, _, _| sim_node.map(|(k, h)| ClientState::new(k, h)), - debug_span!("client"), - ); + let mut client = ref_state + .client + .map(|ref_client, _, _| ref_client.init(), debug_span!("client")); let mut gateway = ref_state.gateway.map( - |sim_node, _, _| sim_node.map(GatewayState::new), + |ref_gateway, _, _| ref_gateway.init(), debug_span!("gateway"), ); @@ -104,16 +86,14 @@ impl StateMachineTest for TunnelTest { .iter() .map(|(id, relay)| { let relay = relay.map( - |relay, ip4, ip6| { - relay.map(|seed| { - firezone_relay::Server::new( - IpStack::from((ip4, ip6)), - rand::rngs::StdRng::seed_from_u64(seed), - 3478, - 49152, - 65535, - ) - }) + |seed, ip4, ip6| { + SimRelay::new(firezone_relay::Server::new( + IpStack::from((ip4, ip6)), + rand::rngs::StdRng::seed_from_u64(seed), + 3478, + 49152, + 65535, + )) }, debug_span!("relay", rid = %id), ); @@ -124,13 +104,19 @@ impl StateMachineTest for TunnelTest { let portal = SimPortal::new(client.inner().id, gateway.inner().id); // Configure client and gateway with the relays. - client.exec_mut(|c| c.init_relays(relays.iter(), ref_state.now)); - gateway.exec_mut(|g| g.init_relays(relays.iter(), ref_state.now)); - - client.exec_mut(|c| c.update_upstream_dns(ref_state.upstream_dns_resolvers.clone())); client.exec_mut(|c| { - c.state - .update_system_resolvers(ref_state.system_dns_resolvers.clone()) + c.sut.update_relays( + HashSet::default(), + HashSet::from_iter(map_explode(relays.iter(), "client")), + ref_state.now, + ) + }); + gateway.exec_mut(|g| { + g.sut.update_relays( + HashSet::default(), + HashSet::from_iter(map_explode(relays.iter(), "gateway")), + ref_state.now, + ) }); let mut this = Self { @@ -142,16 +128,9 @@ impl StateMachineTest for TunnelTest { portal, logger, relays, - client_dns_records: Default::default(), - client_dns_by_sentinel: Default::default(), - client_sent_icmp_requests: Default::default(), - client_received_icmp_replies: Default::default(), - gateway_received_icmp_requests: Default::default(), - client_received_dns_responses: Default::default(), - client_sent_dns_queries: Default::default(), }; - let mut buffered_transmits = VecDeque::new(); + let mut buffered_transmits = BufferedTransmits::default(); this.advance(ref_state, &mut buffered_transmits); // Perform initial setup before we apply the first transition. debug_assert!(buffered_transmits.is_empty()); @@ -165,20 +144,20 @@ impl StateMachineTest for TunnelTest { ref_state: &::State, transition: ::Transition, ) -> Self::SystemUnderTest { - let mut buffered_transmits = VecDeque::new(); + let mut buffered_transmits = BufferedTransmits::default(); // Act: Apply the transition match transition { Transition::AddCidrResource(r) => { state .client - .exec_mut(|c| c.state.add_resources(&[ResourceDescription::Cidr(r)])); + .exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Cidr(r)])); } Transition::AddDnsResource { resource, .. } => state .client - .exec_mut(|c| c.state.add_resources(&[ResourceDescription::Dns(resource)])), + .exec_mut(|c| c.sut.add_resources(&[ResourceDescription::Dns(resource)])), Transition::RemoveResource(id) => { - state.client.exec_mut(|c| c.state.remove_resources(&[id])) + state.client.exec_mut(|c| c.sut.remove_resources(&[id])) } Transition::SendICMPPacketToNonResourceIp { src, @@ -194,7 +173,11 @@ impl StateMachineTest for TunnelTest { } => { let packet = ip_packet::make::icmp_request_packet(src, dst, seq, identifier); - buffered_transmits.extend(state.send_ip_packet_client_to_gateway(packet)); + let transmit = state + .client + .exec_mut(|sim| sim.encapsulate(packet, state.now)); + + buffered_transmits.push(transmit, &state.client); } Transition::SendICMPPacketToDnsResource { src, @@ -203,21 +186,26 @@ impl StateMachineTest for TunnelTest { identifier, resolved_ip, } => { - let available_ips = - state - .client_dns_records - .get(&dst) - .unwrap() - .iter() - .filter(|ip| match ip { - IpAddr::V4(_) => src.is_ipv4(), - IpAddr::V6(_) => src.is_ipv6(), - }); + let available_ips = state + .client + .inner() + .dns_records + .get(&dst) + .unwrap() + .iter() + .filter(|ip| match ip { + IpAddr::V4(_) => src.is_ipv4(), + IpAddr::V6(_) => src.is_ipv6(), + }); let dst = *resolved_ip.select(available_ips); let packet = ip_packet::make::icmp_request_packet(src, dst, seq, identifier); - buffered_transmits.extend(state.send_ip_packet_client_to_gateway(packet)); + let transmit = state + .client + .exec_mut(|sim| Some(sim.encapsulate(packet, state.now)?.into_owned())); + + buffered_transmits.push(transmit, &state.client); } Transition::SendDnsQuery { domain, @@ -225,9 +213,11 @@ impl StateMachineTest for TunnelTest { query_id, dns_server, } => { - let transmit = state.send_dns_query_for(domain, r_type, query_id, dns_server); + let transmit = state.client.exec_mut(|sim| { + sim.send_dns_query_for(domain, r_type, query_id, dns_server, state.now) + }); - buffered_transmits.extend(transmit) + buffered_transmits.push(transmit, &state.client); } Transition::Tick { millis } => { state.now += Duration::from_millis(millis); @@ -235,10 +225,16 @@ impl StateMachineTest for TunnelTest { Transition::UpdateSystemDnsServers { servers } => { state .client - .exec_mut(|c| c.state.update_system_resolvers(servers)); + .exec_mut(|c| c.sut.update_system_resolvers(servers)); } Transition::UpdateUpstreamDnsServers { servers } => { - state.client.exec_mut(|c| c.update_upstream_dns(servers)); + state.client.exec_mut(|c| { + c.sut.update_interface_config(Interface { + ipv4: c.sut.tunnel_ip4().unwrap(), + ipv6: c.sut.tunnel_ip6().unwrap(), + upstream_dns: servers, + }) + }); } Transition::RoamClient { ip4, ip6, port } => { state.network.remove_host(&state.client); @@ -248,10 +244,14 @@ impl StateMachineTest for TunnelTest { .add_host(state.client.inner().id, &state.client)); state.client.exec_mut(|c| { - c.state.reset(); + c.sut.reset(); // In prod, we reconnect to the portal and receive a new `init` message. - c.init_relays(state.relays.iter(), ref_state.now); + c.sut.update_relays( + HashSet::default(), + HashSet::from_iter(map_explode(state.relays.iter(), "client")), + ref_state.now, + ) }); } }; @@ -266,13 +266,22 @@ impl StateMachineTest for TunnelTest { state: &Self::SystemUnderTest, ref_state: &::State, ) { + let ref_client = ref_state.client.inner(); + let sim_client = state.client.inner(); + let sim_gateway = state.gateway.inner(); + // Assert our properties: Check that our actual state is equivalent to our expectation (the reference state). - assert_icmp_packets_properties(state, ref_state); - assert_dns_packets_properties(state, ref_state); - assert_known_hosts_are_valid(state, ref_state); + assert_icmp_packets_properties( + ref_client, + sim_client, + sim_gateway, + &ref_state.global_dns_records, + ); + assert_dns_packets_properties(ref_client, sim_client); + assert_known_hosts_are_valid(ref_client, sim_client); assert_eq!( - state.effective_dns_servers(), - ref_state.expected_dns_servers(), + sim_client.effective_dns_servers(), + ref_client.expected_dns_servers(), "Effective DNS servers should match either system or upstream DNS" ); } @@ -286,48 +295,42 @@ impl TunnelTest { /// Dispatching a [`Transmit`] (read: packet) to a host can trigger more packets, i.e. receiving a STUN request may trigger a STUN response. /// /// 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 VecDeque>, - ) { + fn advance(&mut self, ref_state: &ReferenceState, buffered_transmits: &mut BufferedTransmits) { loop { - if let Some(transmit) = buffered_transmits.pop_front() { + if let Some(transmit) = buffered_transmits.pop() { self.dispatch_transmit(transmit, buffered_transmits, &ref_state.global_dns_records); continue; } - if let Some(transmit) = self.client.poll_transmit() { - buffered_transmits.push_back(transmit); + if let Some(transmit) = self.client.exec_mut(|sim| sim.sut.poll_transmit()) { + buffered_transmits.push(transmit, &self.client); continue; } - if let Some(event) = self.client.exec_mut(|c| c.state.poll_event()) { + if let Some(event) = self.client.exec_mut(|c| c.sut.poll_event()) { self.on_client_event( self.client.inner().id, event, - &ref_state.client_cidr_resources, - &ref_state.client_dns_resources, + &ref_state.client.inner().cidr_resources, + &ref_state.client.inner().dns_resources, &ref_state.global_dns_records, ); continue; } - if let Some(query) = self - .client - .exec_mut(|client| client.state.poll_dns_queries()) - { + if let Some(query) = self.client.exec_mut(|client| client.sut.poll_dns_queries()) { self.on_forwarded_dns_query(query, ref_state); continue; } - if let Some(packet) = self.client.exec_mut(|client| client.state.poll_packets()) { - self.on_client_received_packet(packet); - continue; - } + self.client.exec_mut(|sim| { + while let Some(packet) = sim.sut.poll_packets() { + sim.on_received_packet(packet) + } + }); - if let Some(transmit) = self.gateway.poll_transmit() { - buffered_transmits.push_back(transmit); + 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(|gateway| gateway.state.poll_event()) { + if let Some(event) = self.gateway.exec_mut(|g| g.sut.poll_event()) { self.on_gateway_event(self.gateway.inner().id, event); continue; } @@ -335,7 +338,7 @@ impl TunnelTest { let mut any_relay_advanced = false; for (_, relay) in self.relays.iter_mut() { - let Some(message) = relay.exec_mut(|relay| relay.state.next_command()) else { + let Some(message) = relay.exec_mut(|r| r.sut.next_command()) else { continue; }; @@ -348,11 +351,14 @@ impl TunnelTest { .sending_socket_for(dst.ip()) .expect("relay to never emit packets without a matching socket"); - buffered_transmits.push_back(Transmit { - src: Some(src), - dst, - payload: payload.into(), - }); + buffered_transmits.push( + Transmit { + src: Some(src), + dst, + payload: payload.into(), + }, + relay, + ); } firezone_relay::Command::CreateAllocation { port, family } => { @@ -378,14 +384,6 @@ impl TunnelTest { } } - /// Returns the _effective_ DNS servers that connlib is using. - fn effective_dns_servers(&self) -> BTreeSet { - self.client_dns_by_sentinel - .right_values() - .copied() - .collect() - } - /// Forwards time to the given instant iff the corresponding host would like that (i.e. returns a timestamp <= from `poll_timeout`). /// /// Tying the forwarding of time to the result of `poll_timeout` gives us better coverage because in production, we suspend until the value of `poll_timeout`. @@ -394,81 +392,39 @@ impl TunnelTest { if self .client - .exec_mut(|client| client.state.poll_timeout()) + .exec_mut(|c| c.sut.poll_timeout()) .is_some_and(|t| t <= now) { any_advanced = true; - self.client - .exec_mut(|client| client.state.handle_timeout(now)); + self.client.exec_mut(|c| c.sut.handle_timeout(now)); }; if self .gateway - .exec_mut(|gateway| gateway.state.poll_timeout()) + .exec_mut(|g| g.sut.poll_timeout()) .is_some_and(|t| t <= now) { any_advanced = true; self.gateway - .exec_mut(|gateway| gateway.state.handle_timeout(now, utc_now)) + .exec_mut(|g| g.sut.handle_timeout(now, utc_now)) }; for (_, relay) in self.relays.iter_mut() { if relay - .exec_mut(|relay| relay.state.poll_timeout()) + .exec_mut(|r| r.sut.poll_timeout()) .is_some_and(|t| t <= now) { any_advanced = true; - relay.exec_mut(|relay| relay.state.handle_timeout(now)) + relay.exec_mut(|r| r.sut.handle_timeout(now)) }; } any_advanced } - fn send_ip_packet_client_to_gateway( - &mut self, - packet: MutableIpPacket<'_>, - ) -> Option> { - { - let packet = packet.to_owned().into_immutable(); - - if let Some(icmp) = packet.as_icmp() { - let echo_request = icmp.as_echo_request().expect("to be echo request"); - - self.client_sent_icmp_requests - .insert((echo_request.sequence(), echo_request.identifier()), packet); - } - } - - { - let packet = packet.to_owned().into_immutable(); - - if let Some(udp) = packet.as_udp() { - if let Ok(message) = hickory_proto::op::Message::from_bytes(udp.payload()) { - debug_assert_eq!( - message.message_type(), - MessageType::Query, - "every DNS message sent from the client should be a DNS query" - ); - - self.client_sent_dns_queries.insert(message.id(), packet); - } - } - } - - self.client.encapsulate(packet, self.now) - } - - fn send_ip_packet_gateway_to_client( - &mut self, - packet: MutableIpPacket<'_>, - ) -> Option> { - self.gateway.encapsulate(packet, self.now) - } - /// Dispatches a [`Transmit`] to the correct host. /// /// This function is basically the "network layer" of our tests. @@ -478,7 +434,7 @@ impl TunnelTest { fn dispatch_transmit( &mut self, transmit: Transmit, - buffered_transmits: &mut VecDeque>, + buffered_transmits: &mut BufferedTransmits, global_dns_records: &BTreeMap>, ) { let src = transmit @@ -491,28 +447,20 @@ impl TunnelTest { panic!("Unhandled packet: {src} -> {dst}") }; - let mut buf = [0u8; 1000]; - match host { HostId::Client(_) => { - let Some(packet) = self - .client - .exec_mut(|c| c.state.decapsulate(dst, src, payload, self.now, &mut buf)) - else { - return; - }; - - self.on_client_received_packet(packet); + self.client + .exec_mut(|c| c.handle_packet(payload, src, dst, self.now)); } HostId::Gateway(_) => { - let Some(packet) = self + let Some(transmit) = self .gateway - .exec_mut(|g| g.state.decapsulate(dst, src, payload, self.now, &mut buf)) + .exec_mut(|g| g.handle_packet(global_dns_records, payload, src, dst, self.now)) else { return; }; - self.on_gateway_received_packet(packet, global_dns_records, buffered_transmits); + buffered_transmits.push(transmit, &self.gateway); } HostId::Relay(id) => { let relay = self.relays.get_mut(&id).expect("unknown relay"); @@ -523,7 +471,7 @@ impl TunnelTest { return; }; - buffered_transmits.push_back(transmit); + buffered_transmits.push(transmit, relay); } HostId::Stale => { tracing::debug!(%dst, "Dropping packet because host roamed away or is offline"); @@ -540,20 +488,16 @@ impl TunnelTest { global_dns_records: &BTreeMap>, ) { match event { - ClientEvent::AddedIceCandidates { candidates, .. } => { - self.gateway.exec_mut(|gateway| { - for candidate in candidates { - gateway.state.add_ice_candidate(src, candidate, self.now) - } - }) - } - ClientEvent::RemovedIceCandidates { candidates, .. } => { - self.gateway.exec_mut(|gateway| { - for candidate in candidates { - gateway.state.remove_ice_candidate(src, candidate) - } - }) - } + 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::ConnectionIntent { resource, connected_gateway_ids, @@ -567,7 +511,7 @@ impl TunnelTest { let request = self .client - .exec_mut(|c| c.state.create_or_reuse_connection(resource, gateway, site)) + .exec_mut(|c| c.sut.create_or_reuse_connection(resource, gateway, site)) .unwrap() .unwrap(); @@ -592,9 +536,9 @@ impl TunnelTest { Request::NewConnection(new_connection) => { let answer = self .gateway - .exec_mut(|gateway| { - gateway.state.accept( - self.client.exec_mut(|c| c.id), + .exec_mut(|g| { + g.sut.accept( + self.client.inner().id, snownet::Offer { session_key: new_connection .client_preshared_key @@ -612,9 +556,9 @@ impl TunnelTest { .password, }, }, - self.client.exec_mut(|c| c.state.public_key()), - self.client.exec_mut(|c| c.tunnel_ip4), - self.client.exec_mut(|c| c.tunnel_ip6), + self.client.inner().sut.public_key(), + self.client.inner().sut.tunnel_ip4().unwrap(), + self.client.inner().sut.tunnel_ip6().unwrap(), new_connection .client_payload .domain @@ -627,8 +571,8 @@ impl TunnelTest { .unwrap(); self.client - .exec_mut(|client| { - client.state.accept_answer( + .exec_mut(|c| { + c.sut.accept_answer( snownet::Answer { credentials: snownet::Credentials { username: answer.username, @@ -636,7 +580,7 @@ impl TunnelTest { }, }, resource_id, - self.gateway.exec_mut(|g| g.state.public_key()), + self.gateway.inner().sut.public_key(), self.now, ) }) @@ -644,10 +588,10 @@ impl TunnelTest { } Request::ReuseConnection(reuse_connection) => { self.gateway - .exec_mut(|gateway| { - gateway.state.allow_access( + .exec_mut(|g| { + g.sut.allow_access( resource, - self.client.exec_mut(|c| c.id), + self.client.inner().id, None, reuse_connection.payload.map(|r| (r.name, r.proxy_ips)), self.now, @@ -677,10 +621,10 @@ impl TunnelTest { ); self.gateway - .exec_mut(|gateway| { - gateway.state.allow_access( + .exec_mut(|g| { + g.sut.allow_access( resource, - self.client.exec_mut(|c| c.id), + self.client.inner().id, None, reuse_connection.payload.map(|r| (r.name, r.proxy_ips)), self.now, @@ -693,142 +637,28 @@ impl TunnelTest { tracing::warn!("Unimplemented"); } ClientEvent::DnsServersChanged { dns_by_sentinel } => { - self.client_dns_by_sentinel = dns_by_sentinel; + self.client + .exec_mut(|c| c.dns_by_sentinel = dns_by_sentinel); } } } fn on_gateway_event(&mut self, src: GatewayId, event: GatewayEvent) { match event { - GatewayEvent::AddedIceCandidates { candidates, .. } => self.client.exec_mut(|client| { + GatewayEvent::AddedIceCandidates { candidates, .. } => self.client.exec_mut(|c| { for candidate in candidates { - client.state.add_ice_candidate(src, candidate, self.now) + 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::RemovedIceCandidates { candidates, .. } => { - self.client.exec_mut(|client| { - for candidate in candidates { - client.state.remove_ice_candidate(src, candidate) - } - }) - } GatewayEvent::RefreshDns { .. } => todo!(), } } - /// Process an IP packet received on the client. - fn on_client_received_packet(&mut self, packet: IpPacket<'_>) { - if let Some(icmp) = packet.as_icmp() { - let echo_reply = icmp.as_echo_reply().expect("to be echo reply"); - - self.client_received_icmp_replies.insert( - (echo_reply.sequence(), echo_reply.identifier()), - packet.to_owned(), - ); - - return; - }; - - if let Some(udp) = packet.as_udp() { - if udp.get_source() == 53 { - let mut message = hickory_proto::op::Message::from_bytes(udp.payload()) - .expect("ip packets on port 53 to be DNS packets"); - - self.client_received_dns_responses - .insert(message.id(), packet.to_owned()); - - for record in message.take_answers().into_iter() { - let domain = hickory_name_to_domain(record.name().clone()); - - let ip = match record.data() { - Some(RData::A(rdata::A(ip4))) => IpAddr::from(*ip4), - Some(RData::AAAA(rdata::AAAA(ip6))) => IpAddr::from(*ip6), - unhandled => { - panic!("Unexpected record data: {unhandled:?}") - } - }; - - self.client_dns_records.entry(domain).or_default().push(ip); - } - - // Ensure all IPs are always sorted. - for ips in self.client_dns_records.values_mut() { - ips.sort() - } - - return; - } - } - - unimplemented!("Unhandled packet") - } - - /// Process an IP packet received on the gateway. - fn on_gateway_received_packet( - &mut self, - packet: IpPacket<'_>, - global_dns_records: &BTreeMap>, - buffered_transmits: &mut VecDeque>, - ) { - let packet = packet.to_owned(); - - if packet.as_icmp().is_some() { - self.gateway_received_icmp_requests - .push_back(packet.clone()); - - let echo_response = ip_packet::make::icmp_response_packet(packet); - let maybe_transmit = self.send_ip_packet_gateway_to_client(echo_response); - - buffered_transmits.extend(maybe_transmit); - - return; - } - - if packet.as_udp().is_some() { - let response = ip_packet::make::dns_ok_response(packet, |name| { - global_dns_records - .get(&hickory_name_to_domain(name.clone())) - .cloned() - .into_iter() - .flatten() - }); - - let maybe_transmit = self.send_ip_packet_gateway_to_client(response); - buffered_transmits.extend(maybe_transmit); - - return; - } - - panic!("Unhandled packet") - } - - fn send_dns_query_for( - &mut self, - domain: DomainName, - r_type: RecordType, - query_id: u16, - dns_server: SocketAddr, - ) -> Option> { - let dns_server = *self - .client_dns_by_sentinel - .get_by_right(&dns_server) - .expect("to have a sentinel DNS server for the sampled one"); - - let name = domain_to_hickory_name(domain); - - let src = self.client.exec_mut(|c| c.tunnel_ip(dns_server)); - - let packet = ip_packet::make::dns_query( - name, - r_type, - SocketAddr::new(src, 9999), // An application would pick a random source port that is free. - SocketAddr::new(dns_server, 53), - query_id, - ); - - self.send_ip_packet_client_to_gateway(packet) - } - // TODO: Should we vary the following things via proptests? // - Forwarded DNS query timing out? // - hickory error? @@ -854,7 +684,7 @@ impl TunnelTest { .collect::>(); self.client.exec_mut(|c| { - c.state.on_dns_result( + c.sut.on_dns_result( query, Ok(Ok(Ok(Lookup::new_with_max_ttl( Query::query(name, requested_type), @@ -896,7 +726,7 @@ fn map_client_resource_to_gateway_resource( .expect("resource to be a known CIDR or DNS resource") } -fn hickory_name_to_domain(mut name: hickory_proto::rr::Name) -> DomainName { +pub(crate) fn hickory_name_to_domain(mut name: hickory_proto::rr::Name) -> DomainName { name.set_fqdn(false); // Hack to work around hickory always parsing as FQ let name = name.to_string(); @@ -906,7 +736,7 @@ fn hickory_name_to_domain(mut name: hickory_proto::rr::Name) -> DomainName { domain } -fn domain_to_hickory_name(domain: DomainName) -> hickory_proto::rr::Name { +pub(crate) fn domain_to_hickory_name(domain: DomainName) -> hickory_proto::rr::Name { let domain = domain.to_string(); let name = hickory_proto::rr::Name::from_str(&domain).unwrap(); @@ -914,3 +744,44 @@ fn domain_to_hickory_name(domain: DomainName) -> hickory_proto::rr::Name { name } + +#[derive(Debug, Default)] +struct BufferedTransmits { + inner: VecDeque>, +} + +impl BufferedTransmits { + fn push(&mut self, transmit: impl Into>>, sending_host: &Host) { + let Some(transmit) = transmit.into() else { + return; + }; + + if transmit.src.is_some() { + self.inner.push_back(transmit); + return; + } + + // The `src` of a [`Transmit`] is empty if we want to send if via the default interface. + // In production, the kernel does this for us. + // In this test, we need to always set a `src` so that the remote peer knows where the packet is coming from. + + let Some(src) = sending_host.sending_socket_for(transmit.dst.ip()) else { + tracing::debug!(dst = %transmit.dst, "No socket"); + + return; + }; + + self.inner.push_back(Transmit { + src: Some(src), + ..transmit + }); + } + + fn pop(&mut self) -> Option> { + self.inner.pop_front() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} diff --git a/rust/connlib/tunnel/src/tests/transition.rs b/rust/connlib/tunnel/src/tests/transition.rs index d47517399..b532d8363 100644 --- a/rust/connlib/tunnel/src/tests/transition.rs +++ b/rust/connlib/tunnel/src/tests/transition.rs @@ -1,4 +1,7 @@ -use super::strategies::*; +use super::{ + sim_net::{any_ip_stack, any_port}, + strategies::*, +}; use connlib_shared::{ messages::{ client::{ResourceDescriptionCidr, ResourceDescriptionDns}, @@ -7,7 +10,6 @@ use connlib_shared::{ proptest::*, DomainName, }; -use firezone_relay::IpStack; use hickory_proto::rr::RecordType; use proptest::{prelude::*, sample}; use std::{ @@ -216,19 +218,9 @@ pub(crate) fn question_mark_wildcard_dns_resource() -> impl Strategy impl Strategy { - let ip_stack = prop_oneof![ - host_ip4s().prop_map(IpStack::Ip4), - host_ip6s().prop_map(IpStack::Ip6), - (host_ip4s(), host_ip6s()).prop_map(|(ip4, ip6)| IpStack::Dual { ip4, ip6 }) - ]; - - ( - ip_stack, - any::().prop_filter("port must not be 0", |p| *p != 0), - ) - .prop_map(move |(ip_stack, port)| Transition::RoamClient { - ip4: ip_stack.as_v4().copied(), - ip6: ip_stack.as_v6().copied(), - port, - }) + (any_ip_stack(), any_port()).prop_map(move |(ip_stack, port)| Transition::RoamClient { + ip4: ip_stack.as_v4().copied(), + ip6: ip_stack.as_v6().copied(), + port, + }) } diff --git a/rust/relay/src/server.rs b/rust/relay/src/server.rs index 1f4bcc541..9d7f03a7f 100644 --- a/rust/relay/src/server.rs +++ b/rust/relay/src/server.rs @@ -205,6 +205,18 @@ where &self.auth_secret } + pub fn public_address(&self) -> IpStack { + self.public_address + } + + pub fn public_ip4(&self) -> Option { + Some(IpAddr::V4(*self.public_address.as_v4()?)) + } + + pub fn public_ip6(&self) -> Option { + Some(IpAddr::V6(*self.public_address.as_v6()?)) + } + pub fn listen_port(&self) -> u16 { self.listen_port }