mirror of https://github.com/ospab/ostp.git
fix(server): UoT — set TCP_NODELAY + tear down half-open connections
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.
This commit is contained in:
parent
7e3ada8d4d
commit
5fcc0ba7f4
|
|
@ -410,6 +410,15 @@ async fn run_server_loop(
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if let Ok((stream, peer_addr)) = listener.accept().await {
|
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
|
// Rate limit check
|
||||||
let peer_ip = peer_addr.ip();
|
let peer_ip = peer_addr.ip();
|
||||||
let allowed = {
|
let allowed = {
|
||||||
|
|
|
||||||
|
|
@ -23,36 +23,115 @@ where
|
||||||
tcp_map.write().await.insert(peer_addr, tx);
|
tcp_map.write().await.insert(peer_addr, tx);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Process streams
|
|
||||||
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
||||||
|
|
||||||
// Spawn writer task
|
// Writer: length-prefix (u16 BE) each outbound datagram. OSTP datagrams are
|
||||||
let peer_clone = peer_addr;
|
// always well under 64 KiB (MTU-bounded), so the u16 prefix never truncates.
|
||||||
let tcp_map_clone = tcp_map.clone();
|
let writer = async move {
|
||||||
let writer_task = tokio::spawn(async move {
|
|
||||||
while let Some(packet) = rx.recv().await {
|
while let Some(packet) = rx.recv().await {
|
||||||
let mut out = BytesMut::with_capacity(2 + packet.len());
|
let mut out = BytesMut::with_capacity(2 + packet.len());
|
||||||
out.put_u16(packet.len() as u16);
|
out.put_u16(packet.len() as u16);
|
||||||
out.put_slice(&packet);
|
out.put_slice(&packet);
|
||||||
if write_half.write_all(&out).await.is_err() { break; }
|
if write_half.write_all(&out).await.is_err() { break; }
|
||||||
}
|
}
|
||||||
let _ = tcp_map_clone.write().await.remove(&peer_clone);
|
};
|
||||||
});
|
|
||||||
|
|
||||||
// Spawn reader task
|
// Reader: reassemble length-prefixed frames off the TCP stream (read_exact
|
||||||
let reader_task = tokio::spawn(async move {
|
// handles TCP segmentation) and forward each to the UDP dispatch path.
|
||||||
|
let reader = async move {
|
||||||
let mut len_buf = [0u8; 2];
|
let mut len_buf = [0u8; 2];
|
||||||
loop {
|
loop {
|
||||||
if read_half.read_exact(&mut len_buf).await.is_err() { break; }
|
if read_half.read_exact(&mut len_buf).await.is_err() { break; }
|
||||||
let len = u16::from_be_bytes(len_buf) as usize;
|
let len = u16::from_be_bytes(len_buf) as usize;
|
||||||
if len > 65536 { break; }
|
|
||||||
let mut data = vec![0u8; len];
|
let mut data = vec![0u8; len];
|
||||||
if read_half.read_exact(&mut data).await.is_err() { break; }
|
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);
|
tracing::debug!("UoT client disconnected: {}", peer_addr);
|
||||||
Ok(())
|
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<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>> =
|
||||||
|
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<u8> = (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<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>> =
|
||||||
|
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");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue