From 667230f398f21916a4a59beade7ef0057d4af74c Mon Sep 17 00:00:00 2001 From: IntellyCode <47359527+IntellyCode@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:04:55 +0300 Subject: [PATCH 1/2] Enforce the TCP receive window to bound per-session memory The receive window ipstack advertises now holds. A segment beyond it is dropped for the peer to resend, the head-of-line segment always admitted so the stream advances. The handoff channel to the reader is bounded and filled by reserving a slot before consuming, so buffered data leaves the reassembly map only once it has a home; a reader that frees space wakes the loop to flush more and the follow-up ACK carries the reopened window. The window is advertised as zero below one segment, so a stalled reader puts the peer into persist mode until space frees. Consuming trims a stale head entry a re-segmented retransmission left below the ack, and a FIN is accepted only once the data before it has been consumed, so a full channel never strands the tail. --- src/stream/tcb.rs | 70 ++++++++++++++++++++++++++++++++++++++++++++++- src/stream/tcp.rs | 57 ++++++++++++++++++++++++++++++-------- 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/src/stream/tcb.rs b/src/stream/tcb.rs index 1d496d3..9446aaa 100644 --- a/src/stream/tcb.rs +++ b/src/stream/tcb.rs @@ -4,6 +4,7 @@ use std::{collections::BTreeMap, time::Duration}; pub(super) const MAX_UNACK: u32 = 1024 * 16; // 16KB pub(super) const READ_BUFFER_SIZE: usize = 1024 * 16; // 16KB +pub(super) const READ_CHUNK: usize = 8192; // 8KB, bytes drained from the reassembly buffer per handoff pub(super) const MAX_COUNT_FOR_DUP_ACK: usize = 3; // Maximum number of duplicate ACKs before retransmission /// Retransmission timeout @@ -121,6 +122,13 @@ impl Tcb { log::warn!("{:?}: Received packet seq {seq} < self ack {}, len = {}", self.state, self.ack, buf.len()); return; } + // The head-of-line segment always advances the stream, so it is admitted even at the limit; + // any other segment beyond the receive window is dropped for the peer's RTO to resend. + if seq != self.ack && self.get_unordered_packets_total_len() >= self.read_buffer_size { + #[rustfmt::skip] + log::warn!("{:?}: Receive window full, dropping packet seq {seq}, len = {}", self.state, buf.len()); + return; + } self.unordered_packets.insert(seq, buf); } pub(super) fn get_available_read_buffer_size(&self) -> usize { @@ -137,10 +145,21 @@ impl Tcb { while remaining_bytes > 0 { if let Some(seq) = self.unordered_packets.keys().next().copied() { - if seq != self.ack { + if seq > self.ack { break; // sequence number is not continuous, stop extracting } + if seq < self.ack { + // A retransmission re-segmented across `ack` left a stale head entry; trim the + // part already delivered so consumption can continue from `ack`. + let payload = self.unordered_packets.remove(&seq).unwrap(); + let consumed = self.ack.distance(seq) as usize; + if consumed < payload.len() { + self.unordered_packets.insert(self.ack, payload[consumed..].to_vec()); + } + continue; + } + // remove and get the first packet let mut payload = self.unordered_packets.remove(&seq).unwrap(); let payload_len = payload.len(); @@ -408,6 +427,55 @@ mod tests { assert!(data.is_none()); } + #[test] + fn test_add_unordered_packet_enforces_read_buffer() { + let mut tcb = Tcb::new( + SeqNum(1000), + 1500, + MAX_UNACK, + READ_BUFFER_SIZE, + MAX_COUNT_FOR_DUP_ACK, + RTO, + MAX_RETRANSMIT_COUNT, + ); + + // fill the receive buffer to its limit with an out-of-order gap held open + tcb.add_unordered_packet(SeqNum(1000 + READ_BUFFER_SIZE as u32), vec![7; READ_BUFFER_SIZE]); + assert_eq!(tcb.get_unordered_packets_total_len(), READ_BUFFER_SIZE); + + // a further out-of-order segment is dropped, keeping the buffer bounded + tcb.add_unordered_packet(SeqNum(1000 + 2 * READ_BUFFER_SIZE as u32), vec![8; 500]); + assert_eq!(tcb.get_unordered_packets_total_len(), READ_BUFFER_SIZE); + + // the head-of-line segment is admitted even at the limit, so the stream advances + tcb.add_unordered_packet(SeqNum(1000), vec![9; 500]); + assert_eq!(tcb.unordered_packets.get(&SeqNum(1000)).unwrap().len(), 500); + } + + #[test] + fn test_consume_trims_overlapping_head_entry() { + let mut tcb = Tcb::new( + SeqNum(1000), + 1500, + MAX_UNACK, + READ_BUFFER_SIZE, + MAX_COUNT_FOR_DUP_ACK, + RTO, + MAX_RETRANSMIT_COUNT, + ); + + // an out-of-order segment stored ahead of ack + tcb.add_unordered_packet(SeqNum(1200), vec![2; 300]); + // the gap-filler that a retransmission re-segmented to overlap the stored one + tcb.add_unordered_packet(SeqNum(1000), vec![1; 400]); + + // consuming pulls [1000..1400), advancing ack into the stored entry keyed at 1200 + let data = tcb.consume_unordered_packets(10_000).unwrap(); + assert_eq!(data.len(), 500); // 400 + the 100 bytes of the stored entry past ack + assert_eq!(tcb.ack, SeqNum(1500)); + assert_eq!(tcb.unordered_packets.len(), 0); + } + #[test] fn test_update_inflight_packet_queue() { let mut tcb = Tcb::new( diff --git a/src/stream/tcp.rs b/src/stream/tcp.rs index 2be4bfe..ae5672a 100644 --- a/src/stream/tcp.rs +++ b/src/stream/tcp.rs @@ -7,7 +7,7 @@ use crate::{ tcp_flags::{ACK, FIN, PSH, RST, SYN}, tcp_header_flags, tcp_header_fmt, }, - stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, PacketType, READ_BUFFER_SIZE, RTO, Tcb, TcpState}, + stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, PacketType, READ_BUFFER_SIZE, READ_CHUNK, RTO, Tcb, TcpState}, }; use etherparse::{IpNumber, Ipv4Header, Ipv6FlowLabel, TcpHeader, TcpOptionElement}; use std::{ @@ -165,9 +165,10 @@ pub struct IpStackTcpStream { write_notify: std::sync::Arc>>, destroy_messenger: Option<::tokio::sync::oneshot::Sender<()>>, timeout: Pin>, - data_tx: tokio::sync::mpsc::UnboundedSender>, - data_rx: tokio::sync::mpsc::UnboundedReceiver>, + data_tx: tokio::sync::mpsc::Sender>, + data_rx: tokio::sync::mpsc::Receiver>, read_notify: std::sync::Arc>>, + drain_notify: Arc, task_handle: Option>>, exit_notifier: Option>, temp_read_buffer: Vec, @@ -205,7 +206,8 @@ impl IpStackTcpStream { } let (stream_sender, stream_receiver) = tokio::sync::mpsc::unbounded_channel::(); - let (data_tx, data_rx) = tokio::sync::mpsc::unbounded_channel::>(); + let data_channel_len = config.read_buffer_size.div_ceil(READ_CHUNK).max(1); + let (data_tx, data_rx) = tokio::sync::mpsc::channel::>(data_channel_len); let deadline = tokio::time::Instant::now() + config.timeout; let mut stream = IpStackTcpStream { @@ -222,6 +224,7 @@ impl IpStackTcpStream { data_tx, data_rx, read_notify: std::sync::Arc::new(std::sync::Mutex::new(None)), + drain_notify: Arc::new(tokio::sync::Notify::new()), task_handle: None, exit_notifier: None, temp_read_buffer: Vec::new(), @@ -330,6 +333,8 @@ impl AsyncRead for IpStackTcpStream { buf.put_slice(&data[..capacity]); self.temp_read_buffer.extend_from_slice(&data[capacity..]); } + // A channel slot just freed, so wake the loop to flush more and reopen the window. + self.drain_notify.notify_one(); Poll::Ready(Ok(())) } Poll::Ready(None) => Poll::Ready(Ok(())), @@ -471,6 +476,7 @@ impl IpStackTcpStream { let write_notify = self.write_notify.clone(); let read_notify = self.read_notify.clone(); let data_tx = self.data_tx.clone(); + let drain_notify = self.drain_notify.clone(); let destroy_messenger = self.destroy_messenger.take(); let (exit_task_notifier, exit_monitor) = tokio::sync::mpsc::channel::<()>(10); @@ -489,6 +495,7 @@ impl IpStackTcpStream { write_notify, read_notify, data_tx, + drain_notify, exit_monitor, ) .await; @@ -516,7 +523,8 @@ async fn tcp_main_logic_loop( network_tuple: NetworkTuple, write_notify: std::sync::Arc>>, read_notify: std::sync::Arc>>, - data_tx: tokio::sync::mpsc::UnboundedSender>, + data_tx: tokio::sync::mpsc::Sender>, + drain_notify: Arc, mut exit_monitor: tokio::sync::mpsc::Receiver<()>, ) -> std::io::Result<()> { { @@ -642,6 +650,13 @@ async fn tcp_main_logic_loop( log::debug!("{network_tuple} task exited due to exit signal"); break; } + _ = drain_notify.notified() => { + // The upstream reader freed channel space, so flush whatever is buffered and + // let the follow-up ACK carry the reopened window. + let mut tcb = tcb.lock().unwrap(); + extract_data_n_write_upstream(&up_packet_sender, &mut tcb, network_tuple, &data_tx, &read_notify)?; + continue; + } network_packet = stream_receiver.recv() => network_packet, }; @@ -743,7 +758,7 @@ async fn tcp_main_logic_loop( } PacketType::Invalid => {} } - } else if flags == (ACK | FIN) { + } else if flags == (ACK | FIN) && tcb.get_ack() == incoming_seq { // The other side is closing the connection, we need to send an ACK and change state to CloseWait tcb.increase_ack(); write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?; @@ -838,7 +853,7 @@ async fn tcp_main_logic_loop( log::trace!("{network_tuple} {state:?}: Received final ACK, transitioned to {new_state:?}"); } TcpState::FinWait1 => { - if flags & (ACK | FIN) == (ACK | FIN) && len == 0 { + if flags & (ACK | FIN) == (ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq { // If the received packet is an ACK with FIN, we need to send an ACK and change state to TimeWait directly, not to FinWait2 tcb.increase_ack(); write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?; @@ -863,7 +878,7 @@ async fn tcp_main_logic_loop( } } TcpState::FinWait2 => { - if flags & (ACK | FIN) == (ACK | FIN) && len == 0 { + if flags & (ACK | FIN) == (ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq { tcb.increase_ack(); write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?; tcb.change_state(TcpState::TimeWait); @@ -914,7 +929,7 @@ fn extract_data_n_write_upstream( up_packet_sender: &PacketSender, tcb: &mut Tcb, network_tuple: NetworkTuple, - data_tx: &tokio::sync::mpsc::UnboundedSender>, + data_tx: &tokio::sync::mpsc::Sender>, read_notify: &std::sync::Arc>>, ) -> std::io::Result<()> { let (state, seq, ack) = (tcb.get_state(), tcb.get_seq(), tcb.get_ack()); @@ -924,10 +939,23 @@ fn extract_data_n_write_upstream( return Ok(()); } - if let Some(data) = tcb.consume_unordered_packets(8192) { + // Reserve the handoff slot before consuming, so buffered data is removed only once it has a + // guaranteed home; the reserved permit shrinks the advertised window until the reader drains it. + let permit = match data_tx.try_reserve() { + Ok(permit) => permit, + Err(tokio::sync::mpsc::error::TrySendError::Full(())) => { + write_packet_to_device(up_packet_sender, network_tuple, tcb, None, ACK, None, None)?; + return Ok(()); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(())) => { + return Err(std::io::Error::new(BrokenPipe, "data channel closed")); + } + }; + + if let Some(data) = tcb.consume_unordered_packets(READ_CHUNK) { let hint = if state == TcpState::Established { "normally" } else { "still" }; log::trace!("{network_tuple} {state:?}: {l_info} {hint} receiving data, len = {}", data.len()); - data_tx.send(data).map_err(|e| std::io::Error::new(BrokenPipe, e))?; + permit.send(data); read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); write_packet_to_device(up_packet_sender, network_tuple, tcb, None, ACK, None, None)?; } @@ -947,7 +975,11 @@ pub(crate) fn write_packet_to_device( ) -> std::io::Result { use std::io::Error; let seq = seq.unwrap_or(tcb.get_seq()).0; - let (ack, window_size) = (tcb.get_ack().0, tcb.get_recv_window().max(tcb.get_mtu())); + // Silly-window-syndrome avoidance: advertise a real window only when a full segment fits, + // otherwise advertise zero so the peer enters persist mode until the reader frees space. + let recv_window = tcb.get_recv_window(); + let window_size = if recv_window >= tcb.get_mtu() { recv_window } else { 0 }; + let ack = tcb.get_ack().0; let (src, dst) = (tuple.dst, tuple.src); // Note: The address is reversed here let calc = |ip_header_len: usize, tcp_header_len: usize| tcb.calculate_payload_max_len(ip_header_len, tcp_header_len); let packet = create_raw_packet( @@ -1048,3 +1080,4 @@ pub(crate) fn create_raw_packet( payload: Some(payload), }) } + From d0785d660debc1aa919f749d69eadc12add4c946 Mon Sep 17 00:00:00 2001 From: IntellyCode <47359527+IntellyCode@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:05:18 +0300 Subject: [PATCH 2/2] Cover reserve-before-consume with an async test The stream's first tokio test drives extract_data_n_write_upstream against a full handoff channel: buffered data stays in the reassembly map and the ack holds until the reader drains a slot, then the tail flushes and the ack advances. --- src/stream/tcp.rs | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/src/stream/tcp.rs b/src/stream/tcp.rs index ae5672a..d766f16 100644 --- a/src/stream/tcp.rs +++ b/src/stream/tcp.rs @@ -1081,3 +1081,46 @@ pub(crate) fn create_raw_packet( }) } +#[cfg(test)] +mod tests { + use super::*; + use crate::stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, READ_BUFFER_SIZE, RTO}; + + #[tokio::test] + async fn extract_reserves_before_consuming() { + let (up_tx, _up_rx) = tokio::sync::mpsc::unbounded_channel::(); + let (data_tx, mut data_rx) = tokio::sync::mpsc::channel::>(1); + let read_notify = std::sync::Arc::new(std::sync::Mutex::new(None)); + let nt = NetworkTuple::new("1.1.1.1:1".parse().unwrap(), "2.2.2.2:2".parse().unwrap(), true); + + let mut tcb = Tcb::new( + SeqNum(1000), + 1500, + MAX_UNACK, + READ_BUFFER_SIZE, + MAX_COUNT_FOR_DUP_ACK, + RTO, + MAX_RETRANSMIT_COUNT, + ); + tcb.change_state(TcpState::Established); + tcb.add_unordered_packet(SeqNum(1000), vec![1; 500]); + tcb.add_unordered_packet(SeqNum(1500), vec![2; 500]); + + // first extract fills the single channel slot and advances ack over the first chunk + extract_data_n_write_upstream(&up_tx, &mut tcb, nt, &data_tx, &read_notify).unwrap(); + assert_eq!(tcb.get_ack(), SeqNum(2000)); + + // channel is full: extract leaves the remaining data in the map and does not advance ack + tcb.add_unordered_packet(SeqNum(2000), vec![3; 500]); + extract_data_n_write_upstream(&up_tx, &mut tcb, nt, &data_tx, &read_notify).unwrap(); + assert_eq!(tcb.get_ack(), SeqNum(2000)); + assert_eq!(tcb.get_unordered_packets_total_len(), 500); + + // draining the reader frees a slot, and the next extract flushes the tail + let first = data_rx.recv().await.unwrap(); + assert_eq!(first.len(), 1000); + extract_data_n_write_upstream(&up_tx, &mut tcb, nt, &data_tx, &read_notify).unwrap(); + assert_eq!(tcb.get_ack(), SeqNum(2500)); + assert_eq!(tcb.get_unordered_packets_total_len(), 0); + } +}