From a4d4b4fbfcbe2754d17f44a29d5aaec03003369f Mon Sep 17 00:00:00 2001 From: Gabi Date: Fri, 26 Apr 2024 18:05:38 -0300 Subject: [PATCH] chore(connlib): make peer pure by taking utc time from parameters (#4773) This came up while working on #2030 and thinking about testing `Peer`. Not entirely convinced of taking both `Instant` and `DateTime` but unless we convert the expiration to an instant, which would bring a bunch of new problems, I don't see another way to do this. --- rust/connlib/tunnel/src/gateway.rs | 6 ++++-- rust/connlib/tunnel/src/lib.rs | 3 ++- rust/connlib/tunnel/src/peer.rs | 4 ++-- rust/connlib/tunnel/src/tests.rs | 19 +++++++++++++++---- 4 files changed, 23 insertions(+), 9 deletions(-) diff --git a/rust/connlib/tunnel/src/gateway.rs b/rust/connlib/tunnel/src/gateway.rs index e14b135ad..ce70936bb 100644 --- a/rust/connlib/tunnel/src/gateway.rs +++ b/rust/connlib/tunnel/src/gateway.rs @@ -266,12 +266,14 @@ impl GatewayState { earliest(self.next_expiry_resources_check, self.node.poll_timeout()) } - pub fn handle_timeout(&mut self, now: Instant) { + pub fn handle_timeout(&mut self, now: Instant, utc_now: DateTime) { self.node.handle_timeout(now); match self.next_expiry_resources_check { Some(next_expiry_resources_check) if now >= next_expiry_resources_check => { - self.peers.iter_mut().for_each(|p| p.expire_resources()); + self.peers + .iter_mut() + .for_each(|p| p.expire_resources(utc_now)); self.peers.retain(|_, p| !p.is_emptied()); self.next_expiry_resources_check = Some(now + EXPIRE_RESOURCES_INTERVAL); diff --git a/rust/connlib/tunnel/src/lib.rs b/rust/connlib/tunnel/src/lib.rs index 55f7cfd9a..608e5ad2d 100644 --- a/rust/connlib/tunnel/src/lib.rs +++ b/rust/connlib/tunnel/src/lib.rs @@ -4,6 +4,7 @@ //! [Tunnel] is the main entry-point for this crate. use boringtun::x25519::StaticSecret; +use chrono::Utc; use connlib_shared::{ messages::{ClientId, GatewayId, Relay, RelayId, ResourceId, ReuseConnection}, Callbacks, Result, @@ -201,7 +202,7 @@ where self.device_read_buf.as_mut(), )? { Poll::Ready(io::Input::Timeout(timeout)) => { - self.role_state.handle_timeout(timeout); + self.role_state.handle_timeout(timeout, Utc::now()); continue; } Poll::Ready(io::Input::Device(packet)) => { diff --git a/rust/connlib/tunnel/src/peer.rs b/rust/connlib/tunnel/src/peer.rs index f06cfc907..0e4f2b651 100644 --- a/rust/connlib/tunnel/src/peer.rs +++ b/rust/connlib/tunnel/src/peer.rs @@ -101,9 +101,9 @@ impl ClientOnGateway { self.resources.is_empty() } - pub(crate) fn expire_resources(&mut self) { + pub(crate) fn expire_resources(&mut self, now: DateTime) { self.resources - .retain(|_, (_, e)| !e.is_some_and(|e| e <= Utc::now())); + .retain(|_, (_, e)| !e.is_some_and(|e| e <= now)); } pub(crate) fn remove_resource(&mut self, resource: &ResourceId) { diff --git a/rust/connlib/tunnel/src/tests.rs b/rust/connlib/tunnel/src/tests.rs index b829810cf..3720189c3 100644 --- a/rust/connlib/tunnel/src/tests.rs +++ b/rust/connlib/tunnel/src/tests.rs @@ -1,4 +1,5 @@ use crate::{ClientEvent, ClientState, GatewayState}; +use chrono::{DateTime, Utc}; use connlib_shared::{ messages::{ResourceDescription, ResourceDescriptionCidr, ResourceId}, proptest::cidr_resource, @@ -38,6 +39,7 @@ proptest_state_machine::prop_state_machine! { /// [`proptest`] manipulates this using [`Transition`]s and we assert it against [`ReferenceState`]. struct TunnelTest { now: Instant, + utc_now: DateTime, client: ClientState, gateway: GatewayState, @@ -52,6 +54,7 @@ struct TunnelTest { #[derive(Clone, Debug)] struct ReferenceState { now: Instant, + utc_now: DateTime, client_priv_key: [u8; 32], gateway_priv_key: [u8; 32], @@ -73,6 +76,7 @@ impl StateMachineTest for TunnelTest { ) -> Self::SystemUnderTest { Self { now: ref_state.now, + utc_now: ref_state.utc_now, client: ClientState::new(StaticSecret::from(ref_state.client_priv_key)), gateway: GatewayState::new(StaticSecret::from(ref_state.gateway_priv_key)), logger: tracing_subscriber::fmt() @@ -108,8 +112,9 @@ impl StateMachineTest for TunnelTest { } Transition::Tick { millis } => { state.now += Duration::from_millis(millis); + state.utc_now += Duration::from_millis(millis); state.client.handle_timeout(state.now); - state.gateway.handle_timeout(state.now); + state.gateway.handle_timeout(state.now, state.utc_now); } }; @@ -257,13 +262,19 @@ impl ReferenceStateMachine for ReferenceState { type Transition = Transition; fn init_state() -> proptest::prelude::BoxedStrategy { - (any::<[u8; 32]>(), any::<[u8; 32]>(), Just(Instant::now())) + ( + any::<[u8; 32]>(), + any::<[u8; 32]>(), + Just(Instant::now()), + Just(Utc::now()), + ) .prop_filter( "client and gateway priv key must be different", - |(c, g, _)| c != g, + |(c, g, _, _)| c != g, ) - .prop_map(|(client_priv_key, gateway_priv_key, now)| Self { + .prop_map(|(client_priv_key, gateway_priv_key, now, utc_now)| Self { now, + utc_now, client_priv_key, gateway_priv_key, client_resources: IpNetworkTable::new(),