drop clients on drop packets

Author: lherman-cs

Cargo.lock

@@ -38,6 +38,7 @@ event-listener = "5.4.1"
 parking_lot = "0.12.5"
 dashmap = "6.1.0"
 async-channel = "2.5.0"
+tokio-util = "0.7.17"
 
 [dev-dependencies]
 criterion = { version = "0.8.1", features = ["async", "async_tokio"] }

pulsebeam-runtime/Cargo.toml

@@ -130,6 +130,12 @@ pub enum UnifiedSocketReader {
 }
 
 impl UnifiedSocketReader {
+    pub fn close_peer(&mut self, peer_addr: &SocketAddr) {
+        if let Self::Tcp(inner) = self {
+            inner.close_peer(peer_addr);
+        }
+    }
+
     pub fn local_addr(&self) -> SocketAddr {
         match self {
             Self::Udp(inner) => inner.local_addr(),

pulsebeam-runtime/src/net/mod.rs

@@ -1,43 +1,82 @@
-use super::{BATCH_SIZE, CHUNK_SIZE, RecvPacketBatch, SendPacketBatch};
+use super::{BATCH_SIZE, RecvPacketBatch, SendPacketBatch};
 use crate::net::Transport;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use dashmap::DashMap;
 use std::sync::Arc;
-use std::{io, net::SocketAddr};
+use std::{
+    io,
+    net::{IpAddr, SocketAddr},
+    time::Duration,
+};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
+use tokio_util::sync::CancellationToken;
+
+const MAX_FRAME_SIZE: usize = 1500;
+const MAX_CONNECTIONS: usize = 10_000;
+const MAX_CONNS_PER_IP: usize = 20;
+const READ_TIMEOUT: Duration = Duration::from_secs(30);
+
+/// Shared state for connection management
+type ConnMap = Arc<DashMap<SocketAddr, (async_channel::Sender<Bytes>, CancellationToken)>>;
 
-/// Creates a TCP Reader and Writer pair.
 pub async fn bind(
     addr: SocketAddr,
     external_addr: Option<SocketAddr>,
 ) -> io::Result<(TcpTransportReader, TcpTransportWriter)> {
     let listener = TcpListener::bind(addr).await?;
     let local_addr = external_addr.unwrap_or(listener.local_addr()?);
 
-    let (packet_tx, packet_rx) = async_channel::bounded(BATCH_SIZE * 64);
-    let conns = Arc::new(DashMap::new());
+    let (packet_tx, packet_rx) = async_channel::bounded(BATCH_SIZE * 128);
+    let conns: ConnMap = Arc::new(DashMap::new());
+    let ip_counts = Arc::new(DashMap::<IpAddr, usize>::new());
+
     let readable_notifier = Arc::new(tokio::sync::Notify::new());
     let writable_notifier = Arc::new(tokio::sync::Notify::new());
+    let conn_semaphore = Arc::new(tokio::sync::Semaphore::new(MAX_CONNECTIONS));
 
     let conns_clone = conns.clone();
     let r_notify = readable_notifier.clone();
     let w_notify = writable_notifier.clone();
+    let ip_counts_clone = ip_counts.clone();
+    let semaphore_clone = conn_semaphore.clone();
 
-    // Passive Listener Task (RFC 6544)
+    // Passive Listener Task
     tokio::spawn(async move {
         while let Ok((stream, peer_addr)) = listener.accept().await {
+            let ip = peer_addr.ip();
+
+            // 1. Global Connection Limit
+            let permit = match semaphore_clone.clone().try_acquire_owned() {
+                Ok(p) => p,
+                Err(_) => {
+                    tracing::warn!(%peer_addr, "Rejecting: Global limit reached");
+                    continue;
+                }
+            };
+
+            // 2. Per-IP Limit
+            {
+                let mut count = ip_counts_clone.entry(ip).or_insert(0);
+                if *count >= MAX_CONNS_PER_IP {
+                    tracing::warn!(%peer_addr, "Rejecting: Per-IP limit reached");
+                    continue;
+                }
+                *count += 1;
+            }
+
             let _ = stream.set_nodelay(true);
-            let _ = stream.set_linger(None);
 
             handle_new_connection(
                 stream,
                 peer_addr,
                 local_addr,
                 packet_tx.clone(),
                 conns_clone.clone(),
+                ip_counts_clone.clone(),
                 r_notify.clone(),
                 w_notify.clone(),
+                permit,
             );
             w_notify.notify_waiters();
         }
@@ -47,6 +86,7 @@ pub async fn bind(
         local_addr,
         packet_rx,
         readable_notifier,
+        conns: conns.clone(), // Reader now has access to conns
     };
 
     let writer = TcpTransportWriter {
@@ -62,14 +102,25 @@ pub struct TcpTransportReader {
     packet_rx: async_channel::Receiver<RecvPacketBatch>,
     readable_notifier: Arc<tokio::sync::Notify>,
     local_addr: SocketAddr,
+    /// Shared connection map to allow forced disconnects
+    conns: ConnMap,
 }
 
 impl TcpTransportReader {
+    /// Closes the connection for a specific peer.
+    /// Useful when high-level demux or auth fails.
+    pub fn close_peer(&self, peer_addr: &SocketAddr) {
+        if let Some((_, (tx, cancel))) = self.conns.remove(peer_addr) {
+            cancel.cancel(); // Stops the background reader task
+            tx.close(); // Stops the background writer task
+            tracing::info!(%peer_addr, "Reader forced connection close");
+        }
+    }
+
     pub fn local_addr(&self) -> SocketAddr {
         self.local_addr
     }
 
-    /// Waits until the socket is readable.
     pub async fn readable(&self) -> io::Result<()> {
         loop {
             if !self.packet_rx.is_empty() {
@@ -83,8 +134,6 @@ impl TcpTransportReader {
         }
     }
 
-    /// Pulls packets from the internal channel into the provided vector.
-    /// Takes `&mut self` as requested for the Reader API.
     #[inline]
     pub fn try_recv_batch(&mut self, out: &mut Vec<RecvPacketBatch>) -> io::Result<()> {
         let mut count = 0;
@@ -105,7 +154,7 @@ impl TcpTransportReader {
 #[derive(Clone)]
 pub struct TcpTransportWriter {
     local_addr: SocketAddr,
-    conns: Arc<DashMap<SocketAddr, async_channel::Sender<Bytes>>>,
+    conns: ConnMap,
     writable_notifier: Arc<tokio::sync::Notify>,
 }
 
@@ -126,7 +175,7 @@ impl TcpTransportWriter {
             }
             let mut any_available = false;
             for c in self.conns.iter() {
-                if !c.is_full() {
+                if !c.value().0.is_full() {
                     any_available = true;
                     break;
                 }
@@ -135,99 +184,112 @@ impl TcpTransportWriter {
                 return Ok(());
             }
             let wait = self.writable_notifier.notified();
-
-            // Re-check logic to prevent race conditions
-            let mut any_available = false;
-            for c in self.conns.iter() {
-                if !c.is_full() {
-                    any_available = true;
-                    break;
-                }
-            }
-            if any_available {
-                return Ok(());
-            }
             wait.await;
         }
     }
 
     #[inline]
     pub fn try_send_batch(&self, batch: &SendPacketBatch) -> io::Result<bool> {
-        let Some(peer_tx) = self.conns.get(&batch.dst) else {
-            // If the peer is gone, we drop the packet (consistent with UDP behavior)
+        let Some(peer_entry) = self.conns.get(&batch.dst) else {
             return Ok(true);
         };
+        let (peer_tx, _) = peer_entry.value();
 
-        let required_slots = (batch.buf.len() + batch.segment_size - 1) / batch.segment_size;
+        let required_slots = batch.buf.len().div_ceil(batch.segment_size);
         if peer_tx.capacity().unwrap() - peer_tx.len() < required_slots {
             return Ok(false);
         }
 
         let mut offset = 0;
-        let total_len = batch.buf.len();
-        while offset < total_len {
-            let end = std::cmp::min(offset + batch.segment_size, total_len);
-            let segment = &batch.buf[offset..end];
-            let _ = peer_tx.try_send(Bytes::copy_from_slice(segment));
+        while offset < batch.buf.len() {
+            let end = std::cmp::min(offset + batch.segment_size, batch.buf.len());
+            let _ = peer_tx.try_send(Bytes::copy_from_slice(&batch.buf[offset..end]));
             offset = end;
         }
         Ok(true)
     }
 }
 
-/// Background connection handler
 fn handle_new_connection(
     stream: TcpStream,
     peer_addr: SocketAddr,
     local_addr: SocketAddr,
     packet_tx: async_channel::Sender<RecvPacketBatch>,
-    conns: Arc<DashMap<SocketAddr, async_channel::Sender<Bytes>>>,
+    conns: ConnMap,
+    ip_counts: Arc<DashMap<IpAddr, usize>>,
     r_notify: Arc<tokio::sync::Notify>,
     w_notify: Arc<tokio::sync::Notify>,
+    permit: tokio::sync::OwnedSemaphorePermit,
 ) {
-    let (send_tx, send_rx) = async_channel::bounded::<Bytes>(8192);
-    conns.insert(peer_addr, send_tx);
+    let (send_tx, send_rx) = async_channel::bounded::<Bytes>(1024);
+    let cancel_token = CancellationToken::new();
 
-    let (mut reader, writer) = stream.into_split();
+    conns.insert(peer_addr, (send_tx, cancel_token.clone()));
+
+    let (mut tcp_reader, tcp_writer) = stream.into_split();
+    let peer_ip = peer_addr.ip();
 
     // Task: Receiver (RFC 4571 Un-framing)
+    let r_cancel = cancel_token.clone();
+    let r_conns = conns.clone();
+    let r_ip_counts = ip_counts.clone();
     tokio::spawn(async move {
-        let mut recv_buf = BytesMut::with_capacity(CHUNK_SIZE * 4);
-        while let Ok(n) = reader.read_buf(&mut recv_buf).await {
-            if n == 0 {
-                break;
-            }
-            let mut added = false;
-            while recv_buf.len() >= 2 {
-                let len = u16::from_be_bytes([recv_buf[0], recv_buf[1]]) as usize;
-                if recv_buf.len() < 2 + len {
-                    break;
+        // Guard to release semaphore and cleanup DashMap on task exit
+        let _guard = (permit, r_cancel);
+        let mut recv_buf = BytesMut::with_capacity(MAX_FRAME_SIZE + 2);
+
+        loop {
+            tokio::select! {
+                _ = _guard.1.cancelled() => break,
+                res = tokio::time::timeout(READ_TIMEOUT, tcp_reader.read_buf(&mut recv_buf)) => {
+                    let n = match res {
+                        Ok(Ok(n)) if n > 0 => n,
+                        _ => break, // Timeout, Error, or EOF
+                    };
+
+                    while recv_buf.len() >= 2 {
+                        let len = u16::from_be_bytes([recv_buf[0], recv_buf[1]]) as usize;
+
+                        if len > MAX_FRAME_SIZE || len == 0 {
+                            tracing::warn!(%peer_addr, len, "Invalid TCP frame size, dropping connection");
+                            return;
+                        }
+
+                        if recv_buf.len() < 2 + len { break; }
+
+                        recv_buf.advance(2);
+                        let data = recv_buf.split_to(len).freeze();
+
+                        // Use try_send to prevent reader task from blocking if SFU logic lags
+                        if let Err(_) = packet_tx.try_send(RecvPacketBatch {
+                            src: peer_addr,
+                            dst: local_addr,
+                            buf: data,
+                            stride: len,
+                            len,
+                            transport: Transport::Tcp,
+                        }) {
+                            tracing::debug!("TCP packet dropped: Global queue full");
+                        } else {
+                            r_notify.notify_waiters();
+                        }
+                    }
                 }
-                recv_buf.advance(2);
-                let data = recv_buf.split_to(len).freeze();
-                let _ = packet_tx
-                    .send(RecvPacketBatch {
-                        src: peer_addr,
-                        dst: local_addr,
-                        buf: data,
-                        stride: len,
-                        len,
-                        transport: Transport::Tcp,
-                    })
-                    .await;
-                added = true;
-            }
-            if added {
-                r_notify.notify_waiters();
             }
         }
-        conns.remove(&peer_addr);
+
+        // Final cleanup
+        r_conns.remove(&peer_addr);
+        if let Some(mut count) = r_ip_counts.get_mut(&peer_ip) {
+            *count = count.saturating_sub(1);
+        }
     });
 
-    // Task: Sender (RFC 4571 Framing + Syscall Batching)
+    // Task: Sender
     tokio::spawn(async move {
-        let mut write_buf = Vec::with_capacity(CHUNK_SIZE * 2);
-        let mut writer = writer;
+        let mut write_buf = Vec::with_capacity(MAX_FRAME_SIZE + 2);
+        let mut writer = tcp_writer;
+
         while let Ok(first) = send_rx.recv().await {
             write_buf.clear();
             write_buf.put_u16(first.len() as u16);
@@ -236,7 +298,7 @@ fn handle_new_connection(
             while let Ok(next) = send_rx.try_recv() {
                 write_buf.put_u16(next.len() as u16);
                 write_buf.put_slice(&next);
-                if write_buf.len() > 65535 {
+                if write_buf.len() > 16384 {
                     break;
                 }
             }
@@ -248,6 +310,7 @@ fn handle_new_connection(
         }
     });
 }
+
 #[cfg(test)]
 mod tests {
     use super::*;

pulsebeam-runtime/src/net/tcp.rs

@@ -175,7 +175,10 @@ impl GatewayWorkerActor {
                             1
                         };
 
-                        self.demuxer.demux(batch).await;
+                        let src = batch.src;
+                        if !self.demuxer.demux(batch).await {
+                            self.socket.close_peer(&src);
+                        }
 
                         spent_budget += cost;