fix(transport): stop listener accept loops when the transport layer is cancelled

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<std::io::Result<(TcpStream, SocketAddr)>> {
+    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 {

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,67 @@ 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(())
+}
+
+#[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(())
+}

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;
 

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;

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;