From 5fcc0ba7f46189aaf471820ae5941c496c0a9a4d Mon Sep 17 00:00:00 2001 From: ospab Date: Sun, 12 Jul 2026 01:25:34 +0300 Subject: [PATCH] =?UTF-8?q?fix(server):=20UoT=20=E2=80=94=20set=20TCP=5FNO?= =?UTF-8?q?DELAY=20+=20tear=20down=20half-open=20connections?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two UoT (UDP-over-TCP) correctness issues: - The accepted UoT stream never had TCP_NODELAY set (the client sets it on its end, the server didn't). Nagle's algorithm then batched server->client writes and interacted with the client's delayed ACKs, adding tens-to- hundreds of ms of stall per burst — throttling the download direction badly for streaming/video. Every TCP-tunnel proxy disables Nagle; now the server matches the client. - handle_tcp_connection join!ed the reader and writer tasks, so a half-open connection (client's read side gone, no outbound data pending) parked the writer on rx.recv() forever, leaking the task and a stale tcp_map entry. Rewrote it with select! so either half closing cancels the other and the tcp_map entry is always removed. Added duplex-stream tests covering inbound reassembly across segment boundaries, outbound framing, and teardown-on-close. --- ostp-server/src/lib.rs | 9 +++ ostp-server/src/transport/uot.rs | 105 +++++++++++++++++++++++++++---- 2 files changed, 101 insertions(+), 13 deletions(-) diff --git a/ostp-server/src/lib.rs b/ostp-server/src/lib.rs index 2177eb9..ad6957d 100644 --- a/ostp-server/src/lib.rs +++ b/ostp-server/src/lib.rs @@ -410,6 +410,15 @@ async fn run_server_loop( loop { if let Ok((stream, peer_addr)) = listener.accept().await { + // Disable Nagle's algorithm on the UoT carrier. Without + // this the server→client (download) direction batches + // small writes and interacts with the client's delayed + // ACKs, adding tens-to-hundreds of ms of stall per burst + // — which throttles throughput badly for streaming/video. + // The client already sets nodelay on its end; the server + // must match. (Every TCP-tunnel proxy sets TCP_NODELAY.) + let _ = stream.set_nodelay(true); + // Rate limit check let peer_ip = peer_addr.ip(); let allowed = { diff --git a/ostp-server/src/transport/uot.rs b/ostp-server/src/transport/uot.rs index 86749e8..b8b6778 100644 --- a/ostp-server/src/transport/uot.rs +++ b/ostp-server/src/transport/uot.rs @@ -23,36 +23,115 @@ where tcp_map.write().await.insert(peer_addr, tx); } - // Process streams let (mut read_half, mut write_half) = tokio::io::split(stream); - // Spawn writer task - let peer_clone = peer_addr; - let tcp_map_clone = tcp_map.clone(); - let writer_task = tokio::spawn(async move { + // Writer: length-prefix (u16 BE) each outbound datagram. OSTP datagrams are + // always well under 64 KiB (MTU-bounded), so the u16 prefix never truncates. + let writer = async move { while let Some(packet) = rx.recv().await { let mut out = BytesMut::with_capacity(2 + packet.len()); out.put_u16(packet.len() as u16); out.put_slice(&packet); if write_half.write_all(&out).await.is_err() { break; } } - let _ = tcp_map_clone.write().await.remove(&peer_clone); - }); + }; - // Spawn reader task - let reader_task = tokio::spawn(async move { + // Reader: reassemble length-prefixed frames off the TCP stream (read_exact + // handles TCP segmentation) and forward each to the UDP dispatch path. + let reader = async move { let mut len_buf = [0u8; 2]; loop { if read_half.read_exact(&mut len_buf).await.is_err() { break; } let len = u16::from_be_bytes(len_buf) as usize; - if len > 65536 { break; } let mut data = vec![0u8; len]; if read_half.read_exact(&mut data).await.is_err() { break; } - if udp_tx.send((Bytes::from(data), peer_clone)).await.is_err() { return; } + if udp_tx.send((Bytes::from(data), peer_addr)).await.is_err() { break; } } - }); + }; - let _ = tokio::join!(writer_task, reader_task); + // Either half completing means the connection is dead in that direction, so + // tear the whole thing down: `select!` drops (cancels) the other half. The + // old code `join!`ed both, so a half-open connection (client's read side + // gone but no outbound data pending) left the writer parked on rx.recv() + // forever, leaking the task and its stale tcp_map entry. + tokio::select! { + _ = writer => {}, + _ = reader => {}, + } + + tcp_map.write().await.remove(&peer_addr); tracing::debug!("UoT client disconnected: {}", peer_addr); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn peer() -> SocketAddr { + "127.0.0.1:40000".parse().unwrap() + } + + /// Inbound framing: length-prefixed frames written by the client are + /// reassembled (even when split across reads) and forwarded to udp_tx. + #[tokio::test] + async fn reader_reassembles_framed_datagrams() { + let (mut client, server) = tokio::io::duplex(64 * 1024); + let tcp_map: Arc>>> = + Arc::new(RwLock::new(HashMap::new())); + let (udp_tx, mut udp_rx) = mpsc::channel(16); + let handle = tokio::spawn(handle_tcp_connection(server, peer(), tcp_map.clone(), udp_tx)); + + // Two datagrams; write the second one byte-at-a-time to exercise the + // read_exact reassembly across TCP segment boundaries. + let d1 = b"hello".to_vec(); + client.write_all(&(d1.len() as u16).to_be_bytes()).await.unwrap(); + client.write_all(&d1).await.unwrap(); + + let d2 = vec![0xAB_u8; 1400]; + let framed2: Vec = (d2.len() as u16).to_be_bytes().iter().chain(d2.iter()).copied().collect(); + for b in &framed2 { + client.write_all(&[*b]).await.unwrap(); + } + + let (got1, _) = udp_rx.recv().await.unwrap(); + assert_eq!(got1.as_ref(), d1.as_slice()); + let (got2, from) = udp_rx.recv().await.unwrap(); + assert_eq!(got2.as_ref(), d2.as_slice()); + assert_eq!(from, peer()); + + drop(client); + let _ = handle.await; + } + + /// Outbound framing + teardown: a datagram sent via the tcp_map sender is + /// written to the wire with its u16 length prefix, and once the client + /// hangs up the connection is torn down and its tcp_map entry removed. + #[tokio::test] + async fn writer_frames_outbound_and_cleans_up_on_close() { + let (mut client, server) = tokio::io::duplex(64 * 1024); + let tcp_map: Arc>>> = + Arc::new(RwLock::new(HashMap::new())); + let (udp_tx, _udp_rx) = mpsc::channel(16); + let handle = tokio::spawn(handle_tcp_connection(server, peer(), tcp_map.clone(), udp_tx)); + + // Wait for registration, then push an outbound datagram. + let tx = loop { + if let Some(tx) = tcp_map.read().await.get(&peer()).cloned() { break tx; } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + }; + tx.send(Bytes::from_static(b"world!")).await.unwrap(); + + let mut len_buf = [0u8; 2]; + client.read_exact(&mut len_buf).await.unwrap(); + assert_eq!(u16::from_be_bytes(len_buf) as usize, 6); + let mut body = [0u8; 6]; + client.read_exact(&mut body).await.unwrap(); + assert_eq!(&body, b"world!"); + + // Client hangs up -> reader hits EOF -> select! tears down -> entry gone. + drop(client); + let _ = handle.await; + assert!(!tcp_map.read().await.contains_key(&peer()), "tcp_map entry must be removed on close"); + } +}