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..eaba05bf 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,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(()) +} 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; 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;