diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs index 6cebdade9..abdc6f07b 100644 --- a/rust/connlib/tunnel/src/tests/reference.rs +++ b/rust/connlib/tunnel/src/tests/reference.rs @@ -40,6 +40,7 @@ pub(crate) struct ReferenceState { /// This is used to e.g. mock DNS resolution on the gateway. pub(crate) global_dns_records: BTreeMap>, + pub(crate) drop_direct_client_traffic: bool, pub(crate) network: RoutingTable, } @@ -67,12 +68,21 @@ impl ReferenceStateMachine for ReferenceState { gateways_and_portal(), collection::hash_map(relay_id(), relay_prototype(), 1..=2), global_dns_records(), // Start out with a set of global DNS records so we have something to resolve outside of DNS resources. + any::(), Just(Instant::now()), Just(Utc::now()), ) .prop_filter_map( "network IPs must be unique", - |(c, (gateways, portal, records), relays, mut global_dns, now, utc_now)| { + |( + c, + (gateways, portal, records), + relays, + mut global_dns, + drop_direct_client_traffic, + now, + utc_now, + )| { let mut routing_table = RoutingTable::default(); if !routing_table.add_host(c.inner().id, &c) { @@ -99,6 +109,7 @@ impl ReferenceStateMachine for ReferenceState { relays, portal, global_dns, + drop_direct_client_traffic, now, utc_now, routing_table, @@ -107,7 +118,7 @@ impl ReferenceStateMachine for ReferenceState { ) .prop_filter( "private keys must be unique", - |(c, gateways, _, _, _, _, _, _)| { + |(c, gateways, _, _, _, _, _, _, _)| { let different_keys = gateways .iter() .map(|(_, g)| g.inner().key) @@ -118,7 +129,17 @@ impl ReferenceStateMachine for ReferenceState { }, ) .prop_map( - |(client, gateways, relays, portal, global_dns_records, now, utc_now, network)| { + |( + client, + gateways, + relays, + portal, + global_dns_records, + drop_direct_client_traffic, + now, + utc_now, + network, + )| { Self { now, utc_now, @@ -128,6 +149,7 @@ impl ReferenceStateMachine for ReferenceState { portal, global_dns_records, network, + drop_direct_client_traffic, } }, ) diff --git a/rust/connlib/tunnel/src/tests/sim_gateway.rs b/rust/connlib/tunnel/src/tests/sim_gateway.rs index c16cc3691..24573278e 100644 --- a/rust/connlib/tunnel/src/tests/sim_gateway.rs +++ b/rust/connlib/tunnel/src/tests/sim_gateway.rs @@ -1,6 +1,6 @@ use super::{ reference::{private_key, PrivateKey}, - sim_net::{any_ip_stack, any_port, host, Host}, + sim_net::{any_port, dual_ip_stack, host, Host}, }; use crate::{tests::sut::hickory_name_to_domain, GatewayState}; use connlib_shared::DomainName; @@ -99,7 +99,7 @@ impl RefGateway { } pub(crate) fn ref_gateway_host() -> impl Strategy> { - host(any_ip_stack(), any_port(), ref_gateway()) + host(dual_ip_stack(), any_port(), ref_gateway()) } fn ref_gateway() -> impl Strategy { diff --git a/rust/connlib/tunnel/src/tests/sim_net.rs b/rust/connlib/tunnel/src/tests/sim_net.rs index 5f651d6db..002b8d25a 100644 --- a/rust/connlib/tunnel/src/tests/sim_net.rs +++ b/rust/connlib/tunnel/src/tests/sim_net.rs @@ -99,10 +99,10 @@ impl Host { } } - pub(crate) fn can_route_to(&self, network: IpNetwork) -> bool { - match network { - IpNetwork::V4(_) => self.ip4.is_some(), - IpNetwork::V6(_) => self.ip6.is_some(), + pub(crate) fn is_sender(&self, src: IpAddr) -> bool { + match src { + IpAddr::V4(src) => self.ip4.is_some_and(|v4| v4 == src), + IpAddr::V6(src) => self.ip6.is_some_and(|v6| v6 == src), } } } diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs index 4d871c812..f20e0ef33 100644 --- a/rust/connlib/tunnel/src/tests/strategies.rs +++ b/rust/connlib/tunnel/src/tests/strategies.rs @@ -189,21 +189,6 @@ pub(crate) fn gateways_and_portal() -> impl Strategy< (Just(gateways), Just(portal), dns_resource_records) }, ) - .prop_filter( - "gateway must be able to access their assigned CIDR resources", - |(gateways, portal, _)| { - portal.cidr_resources().all(|(rid, r)| { - portal - .gateway_for_resource(*rid) - .and_then(|g| gateways.get(g)) - .is_some_and(|g| { - // TODO: PRODUCTION CODE DOES NOT HANDLE THIS! - - g.can_route_to(r.address) - }) - }) - }, - ) } fn any_site(sites: HashSet) -> impl Strategy { diff --git a/rust/connlib/tunnel/src/tests/stub_portal.rs b/rust/connlib/tunnel/src/tests/stub_portal.rs index 456cf6f41..51e14398b 100644 --- a/rust/connlib/tunnel/src/tests/stub_portal.rs +++ b/rust/connlib/tunnel/src/tests/stub_portal.rs @@ -78,12 +78,6 @@ impl StubPortal { .collect() } - pub(crate) fn cidr_resources( - &self, - ) -> impl Iterator + '_ { - self.cidr_resources.iter() - } - /// Picks, which gateway and site we should connect to for the given resource. pub(crate) fn handle_connection_intent( &self, diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs index 4fada946d..cb91dd162 100644 --- a/rust/connlib/tunnel/src/tests/sut.rs +++ b/rust/connlib/tunnel/src/tests/sut.rs @@ -50,6 +50,7 @@ pub(crate) struct TunnelTest { pub(crate) gateways: HashMap>, relays: HashMap>, + drop_direct_client_traffic: bool, network: RoutingTable, #[allow(dead_code)] @@ -130,6 +131,7 @@ impl StateMachineTest for TunnelTest { now: ref_state.now, utc_now: ref_state.utc_now, network: ref_state.network.clone(), + drop_direct_client_traffic: ref_state.drop_direct_client_traffic, client, gateways, logger, @@ -504,6 +506,12 @@ impl TunnelTest { .exec_mut(|c| c.handle_packet(payload, src, dst, self.now)); } HostId::Gateway(id) => { + if self.drop_direct_client_traffic && self.client.is_sender(src.ip()) { + tracing::debug!("Dropping direct traffic from client -> gateway"); + + return; + } + let gateway = self.gateways.get_mut(&id).expect("unknown gateway"); let Some(transmit) = gateway