diff --git a/rust/connlib/snownet/src/allocation.rs b/rust/connlib/snownet/src/allocation.rs index 3cec115c2..a76acfe4e 100644 --- a/rust/connlib/snownet/src/allocation.rs +++ b/rust/connlib/snownet/src/allocation.rs @@ -174,6 +174,14 @@ impl Socket { } } +#[derive(Debug, thiserror::Error, PartialEq)] +pub enum FreeReason { + #[error("authentication error")] + AuthenticationError, + #[error("no response received. Is STUN blocked?")] + NoResponseReceived, +} + impl Allocation { pub fn new( server: RelaySocket, @@ -715,9 +723,32 @@ impl Allocation { /// Whether this [`Allocation`] can be freed. /// - /// This is tied to having our credentials cleared (i.e due to an authentication error) and having emitted all events. - pub fn can_be_freed(&self) -> bool { - self.credentials.is_none() && self.events.is_empty() + /// This is tied to having our credentials cleared (i.e due to an authentication error) and having emitted all events or not having received a single response. + pub fn can_be_freed(&self) -> Option { + let pending_work = !self.events.is_empty() + || !self.buffered_transmits.is_empty() + || !self.sent_requests.is_empty(); + + let no_responses = !self.received_any_response(); + let auth_failure = !self.has_credentials(); + + if !pending_work && no_responses { + return Some(FreeReason::NoResponseReceived); + } + + if !pending_work && auth_failure { + return Some(FreeReason::AuthenticationError); + } + + None + } + + pub fn received_any_response(&self) -> bool { + self.active_socket.is_some() + } + + pub fn has_credentials(&self) -> bool { + self.credentials.is_some() } fn log_update(&self) { @@ -2155,7 +2186,10 @@ mod tests { CandidateEvent::Invalid(Candidate::relayed(RELAY_ADDR_IP6, Protocol::Udp).unwrap()), ] ); - assert!(allocation.can_be_freed()); + assert_eq!( + allocation.can_be_freed(), + Some(FreeReason::AuthenticationError) + ); } #[test] @@ -2183,7 +2217,7 @@ mod tests { fn allocation_is_not_freed_on_startup() { let allocation = Allocation::for_test_ip4(Instant::now()); - assert!(!allocation.can_be_freed()); + assert_eq!(allocation.can_be_freed(), None); } #[test] @@ -2325,6 +2359,29 @@ mod tests { assert_eq!(allocation.poll_transmit(), None); } + #[test] + fn allocation_can_be_freed_after_all_requests_time_out() { + let mut allocation = Allocation::for_test_dual(Instant::now()); + + loop { + let Some(timeout) = allocation.poll_timeout() else { + break; + }; + + allocation.handle_timeout(timeout); + + // We expect two transmits. + // The order is not deterministic because internally it is a `HashMap`. + let _ = allocation.poll_transmit().unwrap(); + let _ = allocation.poll_transmit().unwrap(); + } + + assert_eq!( + allocation.can_be_freed(), + Some(FreeReason::NoResponseReceived) + ); + } + fn ch(peer: SocketAddr, now: Instant) -> Channel { Channel { peer, diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs index 21b6601c8..1ef3dd532 100644 --- a/rust/connlib/snownet/src/node.rs +++ b/rust/connlib/snownet/src/node.rs @@ -451,15 +451,15 @@ where self.next_rate_limiter_reset = Some(now + Duration::from_secs(1)); } - self.allocations.retain(|id, allocation| { - if allocation.can_be_freed() { - tracing::info!(%id, "Freeing memory of allocation"); + self.allocations + .retain(|id, allocation| match allocation.can_be_freed() { + Some(e) => { + tracing::info!(%id, "Disconnecting from relay; {e}"); - return false; - } - - true - }); + false + } + None => true, + }); self.connections.remove_failed(&mut self.pending_events); }