Acknowledge received data as it arrives

src/stream/tcb.rs

@@ -1,10 +1,13 @@
 use super::seqnum::SeqNum;
 use etherparse::TcpHeader;
-use std::{collections::BTreeMap, time::Duration};
+use std::{
+    collections::{BTreeMap, VecDeque},
+    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 READ_CHUNK: usize = 8192; // 8KB, bytes drained from the received queue per handoff
 pub(super) const MAX_COUNT_FOR_DUP_ACK: usize = 3; // Maximum number of duplicate ACKs before retransmission
 
 /// Retransmission timeout, and the floor RFC 6298 §2.4 rounds a configured one up to
@@ -49,8 +52,8 @@ pub(super) enum PacketType {
 /// - `inflight_packets` is prerepresented bytes stream from upstream application,
 ///   which have been sent to the lower device but not yet acknowledged.
 /// - `unordered_packets` is the bytes stream received from the lower device,
-///   which can be acknowledged and extracted by `consume_unordered_packets` method
-///   then can be read by upstream application via `Tcp::poll_read` method.
+///   whose in-sequence part `acknowledge_unordered_packets` acknowledges and moves to `received`.
+/// - `received` is the acknowledged bytes stream, which `take_received` hands to the upstream reader.
 /// - `send_window_shift` is the peer's window scale, applied to every window the peer advertises,
 ///   and `recv_window_shift` is this stack's own, applied to every window this stack advertises.
 ///   Both are settled by the SYN exchange, and `recv_window_shift` is `None` when the peer's SYN
@@ -67,6 +70,7 @@ pub(crate) struct Tcb {
     state: TcpState,
     inflight_packets: BTreeMap<SeqNum, InflightPacket>,
     unordered_packets: BTreeMap<SeqNum, Vec<u8>>,
+    received: VecDeque<u8>,
     duplicate_ack_count: usize,
     duplicate_ack_count_helper: SeqNum,
     max_unacked_bytes: u32,
@@ -128,6 +132,7 @@ impl Tcb {
             state: TcpState::Listen,
             inflight_packets: BTreeMap::new(),
             unordered_packets: BTreeMap::new(),
+            received: VecDeque::new(),
             duplicate_ack_count: 0,
             duplicate_ack_count_helper: seq.into(),
             max_unacked_bytes,
@@ -159,23 +164,28 @@ impl Tcb {
         self.duplicate_ack_count >= self.max_count_for_dup_ack
     }
 
-    pub(super) fn add_unordered_packet(&mut self, seq: SeqNum, buf: Vec<u8>) {
+    pub(super) fn add_unordered_packet(&mut self, seq: SeqNum, mut buf: Vec<u8>) {
         if seq < self.ack {
             #[rustfmt::skip]
             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 {
+        // A segment starting beyond a closed receive window is dropped for the peer's RTO to resend.
+        if seq != self.ack && self.get_available_read_buffer_size() == 0 {
             #[rustfmt::skip]
             log::warn!("{:?}: Receive window full, dropping packet seq {seq}, len = {}", self.state, buf.len());
             return;
         }
-        self.unordered_packets.insert(seq, buf);
+        // Bytes past the right edge of the receive window are dropped for the peer to resend.
+        let window = self.read_buffer_size.saturating_sub(self.received.len());
+        buf.truncate(window.saturating_sub(seq.distance(self.ack) as usize));
+        if !buf.is_empty() {
+            self.unordered_packets.insert(seq, buf);
+        }
     }
     pub(super) fn get_available_read_buffer_size(&self) -> usize {
-        self.read_buffer_size.saturating_sub(self.get_unordered_packets_total_len())
+        self.read_buffer_size
+            .saturating_sub(self.get_unordered_packets_total_len() + self.received.len())
     }
     #[inline]
     pub(crate) fn get_unordered_packets_total_len(&self) -> usize {
@@ -228,6 +238,19 @@ impl Tcb {
         if data.is_empty() { None } else { Some(data) }
     }
 
+    /// Acknowledges the in-sequence data of the reassembly buffer, moving it to `received`.
+    pub(super) fn acknowledge_unordered_packets(&mut self) {
+        if let Some(data) = self.consume_unordered_packets(usize::MAX) {
+            self.received.extend(data);
+        }
+    }
+
+    /// Takes up to `max_bytes` of acknowledged data, in sequence order.
+    pub(super) fn take_received(&mut self, max_bytes: usize) -> Option<Vec<u8>> {
+        let len = max_bytes.min(self.received.len());
+        (len > 0).then(|| self.received.drain(..len).collect())
+    }
+
     pub(super) fn increase_seq(&mut self) {
         self.seq += 1;
     }
@@ -535,17 +558,13 @@ mod tests {
             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]);
+        // fill the receive buffer to its limit
+        tcb.add_unordered_packet(SeqNum(1000), 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]

src/stream/tcp.rs

@@ -312,6 +312,14 @@ impl AsyncRead for IpStackTcpStream {
 
         let state = self.tcb.lock().unwrap().get_state();
         if state == TcpState::Closed {
+            // Data on hand reaches the reader before end-of-stream.
+            let data = self.data_rx.try_recv().ok();
+            if let Some(data) = data.or_else(|| self.tcb.lock().unwrap().take_received(usize::MAX)) {
+                let len = std::cmp::min(buf.remaining(), data.len());
+                buf.put_slice(&data[..len]);
+                self.temp_read_buffer.extend_from_slice(&data[len..]);
+                return Poll::Ready(Ok(()));
+            }
             self.shutdown.lock().unwrap().ready();
             self.write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(());
             return Poll::Ready(Ok(()));
@@ -671,6 +679,7 @@ async fn tcp_main_logic_loop(
         let network_packet = tokio::select! {
             _ = exit_monitor.recv() => {
                 log::debug!("{network_tuple} task exited due to exit signal");
+                read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(());
                 break;
             }
             _ = drain_notify.notified() => {
@@ -943,7 +952,9 @@ async fn tcp_main_logic_loop(
                         extract_data_n_write_upstream(&up_packet_sender, &mut tcb, network_tuple, &data_tx, &read_notify)?;
                         write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(());
                     }
-                    if flags & FIN == FIN {
+                    if flags & FIN == FIN && tcb.get_ack() == incoming_seq + len as u32 {
+                        tcb.increase_ack();
+                        write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
                         tcb.change_state(TcpState::TimeWait);
                         tokio::spawn(task_wait_to_close(tcb_clone.clone(), exit_notifier, network_tuple, config.two_msl));
                         let new_state = tcb.get_state();
@@ -994,8 +1005,9 @@ fn extract_data_n_write_upstream(
         return Ok(());
     }
 
-    // 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.
+    tcb.acknowledge_unordered_packets();
+
+    // Reserve the handoff slot before taking acknowledged data.
     let permit = match data_tx.try_reserve() {
         Ok(permit) => permit,
         Err(tokio::sync::mpsc::error::TrySendError::Full(())) => {
@@ -1007,7 +1019,7 @@ fn extract_data_n_write_upstream(
         }
     };
 
-    if let Some(data) = tcb.consume_unordered_packets(READ_CHUNK) {
+    if let Some(data) = tcb.take_received(READ_CHUNK) {
         let hint = if state == TcpState::Established { "normally" } else { "still" };
         log::trace!("{network_tuple} {state:?}: {l_info} {hint} receiving data, len = {}", data.len());
         permit.send(data);
@@ -1336,18 +1348,18 @@ mod tests {
         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
+        // channel is full: extract advances ack over the tail and holds it for the reader
         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);
+        assert_eq!(tcb.get_ack(), SeqNum(2500));
+        assert_eq!(tcb.get_available_read_buffer_size(), READ_BUFFER_SIZE - 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);
+        assert_eq!(tcb.get_available_read_buffer_size(), READ_BUFFER_SIZE);
     }
 
     /// Opens a connection with a SYN advertising `window` and `syn_options`, returning the stream