fix(server): detect dead clients and reclaim their sockets (#471)

* fix(server): detect dead clients and tear down their data channel pool

The server never noticed a disconnected client. The control channel task
only wrote heartbeats, and writes keep succeeding on a half-closed socket,
so the ControlChannelHandle stayed in the map, its shutdown_tx never
dropped, and the data channel pool's listening sockets leaked.

Read the control channel concurrently with writes (the client sends nothing
after the handshake, so any completed read means it disconnected), and
remove the handle from the map when the control channel ends so shutdown_tx
drops and the pool is torn down.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(server): harden control channel cleanup

Avoid retaining the control-channel map from cleanup tasks and treat unexpected client data as a protocol violation. Cover half-closed channels, map teardown, listener release, and heartbeat-disabled client shutdown.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Stephen Tierney <sjtrny@gmail.com>

Author: 3 people

src/server.rs

@@ -16,7 +16,7 @@ use backoff::ExponentialBackoff;
 use rand::RngCore;
 use std::collections::HashMap;
 use std::path::Path;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::time::Duration;
 use tokio::io::{self, copy_bidirectional, AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream, UdpSocket};
@@ -355,8 +355,13 @@ async fn do_control_channel_handshake<T: 'static + Transport>(
         conn.flush().await?;
 
         info!(service = %service_config.name, "Control channel established");
-        let handle =
-            ControlChannelHandle::new(conn, service_config, server_config.heartbeat_interval);
+        let handle = ControlChannelHandle::new(
+            conn,
+            service_config,
+            server_config.heartbeat_interval,
+            Arc::downgrade(&control_channels),
+            session_key,
+        );
 
         // Insert the new handle
         let _ = h.insert(service_digest, session_key, handle);
@@ -410,6 +415,8 @@ where
         conn: T::Stream,
         service: ServerServiceConfig,
         heartbeat_interval: u64,
+        control_channels: Weak<RwLock<ControlChannelMap<T>>>,
+        session_key: Nonce,
     ) -> ControlChannelHandle<T> {
         // Create a shutdown channel
         let (shutdown_tx, shutdown_rx) = broadcast::channel::<bool>(1);
@@ -505,12 +512,17 @@ where
             heartbeat_interval,
         };
 
-        // Run the control channel
+        // On exit, drop the handle so `shutdown_tx` drops and the data channel pool
+        // closes; otherwise its sockets leak. session_key is unique, so this never
+        // removes a reconnected session.
         tokio::spawn(
             async move {
                 if let Err(err) = ch.run().await {
                     error!("{:#}", err);
                 }
+                if let Some(control_channels) = control_channels.upgrade() {
+                    control_channels.write().await.remove2(&session_key);
+                }
             }
             .instrument(Span::current()),
         );
@@ -532,25 +544,48 @@ struct ControlChannel<T: Transport> {
 }
 
 impl<T: Transport> ControlChannel<T> {
-    async fn write_and_flush(&mut self, data: &[u8]) -> Result<()> {
-        write_and_flush(&mut self.conn, data)
-            .await
-            .with_context(|| "Failed to write control cmds")?;
-        Ok(())
-    }
     // Run a control channel
     #[instrument(skip_all)]
-    async fn run(mut self) -> Result<()> {
+    async fn run(self) -> Result<()> {
         let create_ch_cmd = bincode::serialize(&ControlChannelCmd::CreateDataChannel).unwrap();
         let heartbeat = bincode::serialize(&ControlChannelCmd::HeartBeat).unwrap();
 
+        // Split so we can read (to detect a dead client) and write concurrently.
+        let ControlChannel {
+            conn,
+            mut shutdown_rx,
+            mut data_ch_req_rx,
+            heartbeat_interval,
+        } = self;
+        let (mut rd, mut wr) = tokio::io::split(conn);
+        let mut probe = [0u8; 1];
+
         // Wait for data channel requests and the shutdown signal
         loop {
             tokio::select! {
-                val = self.data_ch_req_rx.recv() => {
+                // The client sends nothing after the handshake, so any completed
+                // read means it's gone. Heartbeat writes alone never notice a
+                // half-closed client, which leaks its sockets.
+                res = rd.read(&mut probe) => {
+                    match res {
+                        Ok(0) => {
+                            debug!("Control channel closed by the client");
+                            break;
+                        }
+                        Ok(bytes_read) => {
+                            warn!(bytes_read, "Unexpected data on control channel");
+                            break;
+                        }
+                        Err(e) => {
+                            error!("Control channel read error: {:#}", e);
+                            break;
+                        }
+                    }
+                },
+                val = data_ch_req_rx.recv() => {
                     match val {
                         Some(_) => {
-                            if let Err(e) = self.write_and_flush(&create_ch_cmd).await {
+                            if let Err(e) = write_and_flush(&mut wr, &create_ch_cmd).await {
                                 error!("{:#}", e);
                                 break;
                             }
@@ -560,14 +595,14 @@ impl<T: Transport> ControlChannel<T> {
                         }
                     }
                 },
-                _ = time::sleep(Duration::from_secs(self.heartbeat_interval)), if self.heartbeat_interval != 0 => {
-                            if let Err(e) = self.write_and_flush(&heartbeat).await {
+                _ = time::sleep(Duration::from_secs(heartbeat_interval)), if heartbeat_interval != 0 => {
+                            if let Err(e) = write_and_flush(&mut wr, &heartbeat).await {
                                 error!("{:#}", e);
                                 break;
                             }
                 }
                 // Wait for the shutdown signal
-                _ = self.shutdown_rx.recv() => {
+                _ = shutdown_rx.recv() => {
                     break;
                 }
             }
@@ -888,3 +923,159 @@ async fn run_udp_connection_pool<T: Transport>(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::ErrorKind;
+    use std::net::SocketAddr;
+
+    const CLEANUP_TIMEOUT: Duration = Duration::from_secs(5);
+
+    fn unused_tcp_addr() -> Result<SocketAddr> {
+        let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
+        Ok(listener.local_addr()?)
+    }
+
+    async fn connected_tcp_pair() -> Result<(TcpStream, TcpStream)> {
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let addr = listener.local_addr()?;
+        let (client, (server, _)) = tokio::try_join!(TcpStream::connect(addr), listener.accept())?;
+        Ok((server, client))
+    }
+
+    async fn wait_until_listener_is_bound(addr: SocketAddr) -> Result<()> {
+        time::timeout(CLEANUP_TIMEOUT, async {
+            loop {
+                match TcpListener::bind(addr).await {
+                    Ok(listener) => {
+                        drop(listener);
+                        time::sleep(Duration::from_millis(10)).await;
+                    }
+                    Err(error) if error.kind() == ErrorKind::AddrInUse => return Ok(()),
+                    Err(error) => return Err(error.into()),
+                }
+            }
+        })
+        .await
+        .context("service listener was not created before the deadline")?
+    }
+
+    async fn wait_until_listener_is_released(addr: SocketAddr) -> Result<TcpListener> {
+        time::timeout(CLEANUP_TIMEOUT, async {
+            loop {
+                match TcpListener::bind(addr).await {
+                    Ok(listener) => return Ok(listener),
+                    Err(error) if error.kind() == ErrorKind::AddrInUse => {
+                        time::sleep(Duration::from_millis(10)).await;
+                    }
+                    Err(error) => return Err(error.into()),
+                }
+            }
+        })
+        .await
+        .context("service listener was not released before the deadline")?
+    }
+
+    async fn wait_until_channel_is_removed(
+        control_channels: &RwLock<ControlChannelMap<TcpTransport>>,
+        service_digest: ServiceDigest,
+        session_key: Nonce,
+    ) -> Result<()> {
+        time::timeout(CLEANUP_TIMEOUT, async {
+            loop {
+                let removed = {
+                    let channels = control_channels.read().await;
+                    channels.get1(&service_digest).is_none()
+                        && channels.get2(&session_key).is_none()
+                };
+                if removed {
+                    return;
+                }
+                time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .context("control channel was not removed before the deadline")?;
+        Ok(())
+    }
+
+    fn tcp_service(bind_addr: SocketAddr) -> ServerServiceConfig {
+        ServerServiceConfig {
+            service_type: ServiceType::Tcp,
+            name: "cleanup-test".to_owned(),
+            bind_addr: bind_addr.to_string(),
+            ..Default::default()
+        }
+    }
+
+    #[tokio::test]
+    async fn half_closed_control_channel_releases_map_entry_and_listener_without_heartbeat(
+    ) -> Result<()> {
+        let service_addr = unused_tcp_addr()?;
+        let (server_conn, mut client_conn) = connected_tcp_pair().await?;
+        let service_digest = [1_u8; HASH_WIDTH_IN_BYTES];
+        let session_key = [2_u8; HASH_WIDTH_IN_BYTES];
+        let control_channels = Arc::new(RwLock::new(ControlChannelMap::new()));
+
+        {
+            // Match the handshake's locking order so cleanup cannot run before insertion.
+            let mut channels = control_channels.write().await;
+            let handle = ControlChannelHandle::<TcpTransport>::new(
+                server_conn,
+                tcp_service(service_addr),
+                0,
+                Arc::downgrade(&control_channels),
+                session_key,
+            );
+            if channels
+                .insert(service_digest, session_key, handle)
+                .is_err()
+            {
+                bail!("failed to insert test control channel");
+            }
+        }
+
+        wait_until_listener_is_bound(service_addr).await?;
+        client_conn.shutdown().await?;
+
+        wait_until_channel_is_removed(&control_channels, service_digest, session_key).await?;
+        let released_listener = wait_until_listener_is_released(service_addr).await?;
+        drop(released_listener);
+        drop(client_conn);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn dropping_channel_map_releases_listener_without_retaining_cycle() -> Result<()> {
+        let service_addr = unused_tcp_addr()?;
+        let (server_conn, _client_conn) = connected_tcp_pair().await?;
+        let service_digest = [3_u8; HASH_WIDTH_IN_BYTES];
+        let session_key = [4_u8; HASH_WIDTH_IN_BYTES];
+        let control_channels = Arc::new(RwLock::new(ControlChannelMap::new()));
+
+        {
+            let mut channels = control_channels.write().await;
+            let handle = ControlChannelHandle::<TcpTransport>::new(
+                server_conn,
+                tcp_service(service_addr),
+                0,
+                Arc::downgrade(&control_channels),
+                session_key,
+            );
+            if channels
+                .insert(service_digest, session_key, handle)
+                .is_err()
+            {
+                bail!("failed to insert test control channel");
+            }
+        }
+
+        wait_until_listener_is_bound(service_addr).await?;
+        drop(control_channels);
+
+        let released_listener = wait_until_listener_is_released(service_addr).await?;
+        drop(released_listener);
+        Ok(())
+    }
+}