From 58f53df0c28d584a2f7377dfd29ac47d30cc35ba Mon Sep 17 00:00:00 2001 From: Wael Jami Date: Tue, 13 Oct 2026 09:20:15 +0200 Subject: [PATCH 1/2] fix(transport): stop listener accept loops when the transport layer is cancelled The TCP, TLS and WebSocket listener accept loops ran in detached tasks that never observed the transport layer's cancel token, so a cancelled transport layer kept its listening sockets bound and kept TransportLayerInner (and its connection map) alive. Race accept() against the token, as the UDP serve loop already does, and exit the loop on cancel. --- src/transport/tcp_listener.rs | 40 ++++++++++++++---- src/transport/tests/test_listener_api.rs | 53 +++++++++++++++++++++++- src/transport/tls.rs | 21 ++++++---- src/transport/websocket.rs | 21 ++++++---- 4 files changed, 112 insertions(+), 23 deletions(-) diff --git a/src/transport/tcp_listener.rs b/src/transport/tcp_listener.rs index 0aec80af..1894d756 100644 --- a/src/transport/tcp_listener.rs +++ b/src/transport/tcp_listener.rs @@ -1,3 +1,4 @@ +use crate::platform::CancellationToken; use crate::transport::tcp::TcpConnection; use crate::transport::transport_layer::TransportLayerInnerRef; use crate::transport::SipAddr; @@ -6,7 +7,7 @@ use crate::Result; use socket2::{Domain, Protocol, Socket, Type}; use std::fmt; use std::{net::SocketAddr, sync::Arc}; -use tokio::net::TcpListener; +use tokio::net::{TcpListener, TcpStream}; use tracing::{debug, warn}; pub struct TcpListenerConnectionInner { pub local_addr: SipAddr, @@ -56,13 +57,19 @@ impl TcpListenerConnection { let listener_local_addr = self.get_addr().clone(); crate::platform::spawn(async move { loop { - let (stream, remote_addr) = match listener.accept().await { - Ok((stream, remote_addr)) => (stream, remote_addr), - Err(e) => { - warn!(error = ?e, "Failed to accept connection"); - continue; - } - }; + let (stream, remote_addr) = + match accept_or_cancelled(&listener, &transport_layer_inner.cancel_token).await + { + Some(Ok((stream, remote_addr))) => (stream, remote_addr), + Some(Err(e)) => { + warn!(error = ?e, "Failed to accept connection"); + continue; + } + None => { + debug!(local = %listener_local_addr, "TCP listener cancelled"); + break; + } + }; if !transport_layer_inner.is_whitelisted(remote_addr.ip()).await { debug!(remote = %remote_addr, "tcp connection rejected by whitelist"); continue; @@ -88,6 +95,23 @@ impl TcpListenerConnection { } } +/// Accepts the next connection, or returns `None` once `cancel_token` fires. +/// +/// Listener accept loops run in detached tasks; racing the token lets them exit +/// and release the listening socket (and the transport layer they hold) when the +/// transport layer is cancelled. +pub(crate) async fn accept_or_cancelled( + listener: &TcpListener, + cancel_token: &CancellationToken, +) -> Option> { + let mut cancel_f = core::pin::pin!(cancel_token.cancelled()); + let mut accept_f = core::pin::pin!(listener.accept()); + match crate::platform::select::select2(&mut cancel_f, &mut accept_f).await { + crate::platform::select::Either::A(()) => None, + crate::platform::select::Either::B(accepted) => Some(accepted), + } +} + impl TcpListenerConnection { pub fn get_addr(&self) -> &SipAddr { if let Some(external) = &self.inner.external { diff --git a/src/transport/tests/test_listener_api.rs b/src/transport/tests/test_listener_api.rs index 318606dc..74a67811 100644 --- a/src/transport/tests/test_listener_api.rs +++ b/src/transport/tests/test_listener_api.rs @@ -1,7 +1,11 @@ use crate::{ - transport::{TcpListenerConnection, WebSocketListenerConnection}, + platform::CancellationToken, + transport::{ + SipConnection, TcpListenerConnection, TransportLayer, WebSocketListenerConnection, + }, Result, }; +use std::{net::SocketAddr, time::Duration}; /// Test new TcpListenerConnection API #[tokio::test] @@ -59,3 +63,50 @@ async fn test_websocket_listener_connection_api() -> Result<()> { Ok(()) } + +/// Reserves a free loopback port (listeners bind lazily in `serve_listens`). +fn free_port() -> SocketAddr { + let probe = std::net::TcpListener::bind("127.0.0.1:0").expect("bind probe"); + probe.local_addr().expect("probe addr") +} + +/// The accept loop must stop when the transport layer is cancelled, closing the +/// listening socket even while the `TransportLayer` itself is still alive. +async fn assert_listener_closed_on_cancel(listener: SipConnection, addr: SocketAddr) { + let cancel = CancellationToken::new(); + let transport_layer = TransportLayer::new(cancel.clone()); + transport_layer.add_transport(listener); + transport_layer + .serve_listens() + .await + .expect("serve listeners"); + tokio::net::TcpStream::connect(addr) + .await + .expect("listener accepts before cancel"); + + cancel.cancel(); + for _ in 0..50 { + if tokio::net::TcpStream::connect(addr).await.is_err() { + return; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("listener on {addr} still accepting after cancel"); +} + +#[tokio::test] +async fn test_tcp_listener_closed_on_cancel() -> Result<()> { + let addr = free_port(); + let listener = TcpListenerConnection::new(addr, None).await?; + assert_listener_closed_on_cancel(listener.into(), addr).await; + Ok(()) +} + +#[cfg(feature = "websocket")] +#[tokio::test] +async fn test_websocket_listener_closed_on_cancel() -> Result<()> { + let addr = free_port(); + let listener = WebSocketListenerConnection::new(addr, None, false).await?; + assert_listener_closed_on_cancel(listener.into(), addr).await; + Ok(()) +} diff --git a/src/transport/tls.rs b/src/transport/tls.rs index ff294b91..63d0d7cc 100644 --- a/src/transport/tls.rs +++ b/src/transport/tls.rs @@ -6,6 +6,7 @@ use super::{ }; use crate::platform::CancellationToken; use crate::sip::SipMessage; +use crate::transport::tcp_listener::accept_or_cancelled; use crate::{error::Error, transport::transport_layer::TransportLayerInnerRef, Result}; use core::future::Future; use core::pin::Pin; @@ -321,13 +322,19 @@ impl TlsListenerConnection { crate::platform::spawn(async move { loop { - let (stream, remote_addr) = match listener.accept().await { - Ok((stream, remote_addr)) => (stream, remote_addr), - Err(e) => { - warn!(error = ?e, "Failed to accept TLS connection"); - continue; - } - }; + let (stream, remote_addr) = + match accept_or_cancelled(&listener, &transport_layer_inner.cancel_token).await + { + Some(Ok((stream, remote_addr))) => (stream, remote_addr), + Some(Err(e)) => { + warn!(error = ?e, "Failed to accept TLS connection"); + continue; + } + None => { + debug!(local = %listener_local_addr, "TLS listener cancelled"); + break; + } + }; if !transport_layer_inner.is_whitelisted(remote_addr.ip()).await { debug!(remote = %remote_addr, "tls connection rejected by whitelist"); continue; diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs index 250ae5c5..3453474a 100644 --- a/src/transport/websocket.rs +++ b/src/transport/websocket.rs @@ -1,5 +1,6 @@ use crate::platform::CancellationToken; use crate::sip::SipMessage; +use crate::transport::tcp_listener::accept_or_cancelled; use crate::{ transport::{ connection::{TransportSender, KEEPALIVE_REQUEST, KEEPALIVE_RESPONSE}, @@ -101,13 +102,19 @@ impl WebSocketListenerConnection { debug!(local = %self.inner.local_addr, "Starting WebSocket listener"); crate::platform::spawn(async move { loop { - let (stream, remote_addr) = match listener.accept().await { - Ok((stream, remote_addr)) => (stream, remote_addr), - Err(e) => { - warn!(error = ?e, "Failed to accept WebSocket connection"); - continue; - } - }; + let (stream, remote_addr) = + match accept_or_cancelled(&listener, &transport_layer_inner.cancel_token).await + { + Some(Ok((stream, remote_addr))) => (stream, remote_addr), + Some(Err(e)) => { + warn!(error = ?e, "Failed to accept WebSocket connection"); + continue; + } + None => { + debug!(local = %listener_local_addr, "WebSocket listener cancelled"); + break; + } + }; if !transport_layer_inner.is_whitelisted(remote_addr.ip()).await { debug!(remote = %remote_addr, "websocket connection rejected by whitelist"); continue; From 61a6762155b55e5b2a2549e616a6a1cfa2060968 Mon Sep 17 00:00:00 2001 From: Wael Jami Date: Wed, 14 Oct 2026 10:00:33 +0200 Subject: [PATCH 2/2] test(transport): cover TLS listener shutdown on cancel Test suggested by @tgeorge06 in review. --- src/transport/tests/test_listener_api.rs | 17 +++++++++++++++++ src/transport/tests/test_tls_reload.rs | 4 +++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/transport/tests/test_listener_api.rs b/src/transport/tests/test_listener_api.rs index 74a67811..eaba05bf 100644 --- a/src/transport/tests/test_listener_api.rs +++ b/src/transport/tests/test_listener_api.rs @@ -110,3 +110,20 @@ async fn test_websocket_listener_closed_on_cancel() -> Result<()> { assert_listener_closed_on_cancel(listener.into(), addr).await; Ok(()) } + +#[cfg(feature = "rustls")] +#[tokio::test] +async fn test_tls_listener_closed_on_cancel() -> Result<()> { + use crate::transport::{TlsConfig, TlsListenerConnection}; + + let (cert, key) = super::test_tls_reload::generate_test_cert("localhost")?; + let config = TlsConfig { + cert: Some(cert.into_bytes()), + key: Some(key.into_bytes()), + ..Default::default() + }; + let addr = free_port(); + let listener = TlsListenerConnection::new(addr, None, config).await?; + assert_listener_closed_on_cancel(listener.into(), addr).await; + Ok(()) +} diff --git a/src/transport/tests/test_tls_reload.rs b/src/transport/tests/test_tls_reload.rs index bf6fac64..7375fe39 100644 --- a/src/transport/tests/test_tls_reload.rs +++ b/src/transport/tests/test_tls_reload.rs @@ -3,7 +3,9 @@ use std::sync::Arc; /// Generate a self-signed certificate and key for testing #[cfg(feature = "rustls")] -fn generate_test_cert(subject_cn: &str) -> std::result::Result<(String, String), std::io::Error> { +pub(super) fn generate_test_cert( + subject_cn: &str, +) -> std::result::Result<(String, String), std::io::Error> { use std::process::Command; use tempfile::TempDir;