Handle overlapping stream receive chunks

src/bin/dosh-client.rs

@@ -8124,48 +8124,14 @@ fn accept_stream_data(
     stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
     data: StreamData,
 ) -> (Vec<Vec<u8>>, usize, u64) {
-    let stream_id = data.stream_id;
-    let expected = stream_next_recv_offset.entry(stream_id).or_insert(0);
-    if !dosh::transport::stream_data_within_receive_window(
-        *expected,
+    dosh::transport::accept_stream_data_chunks(
+        stream_next_recv_offset,
+        stream_recv_pending,
+        data.stream_id,
         data.offset,
-        data.bytes.len(),
+        data.bytes,
         STREAM_INITIAL_WINDOW,
-    ) {
-        return (Vec::new(), 0, *expected);
-    }
-    if data.offset < *expected {
-        return (Vec::new(), 0, *expected);
-    }
-    if data.offset > *expected {
-        stream_recv_pending
-            .entry(stream_id)
-            .or_default()
-            .entry(data.offset)
-            .or_insert(data.bytes);
-        return (Vec::new(), 0, *expected);
-    }
-
-    let mut writes = vec![data.bytes];
-    let mut consumed = writes[0].len();
-    *expected = expected.saturating_add(consumed as u64);
-
-    while let Some(bytes) = stream_recv_pending
-        .entry(stream_id)
-        .or_default()
-        .remove(expected)
-    {
-        consumed = consumed.saturating_add(bytes.len());
-        *expected = expected.saturating_add(bytes.len() as u64);
-        writes.push(bytes);
-    }
-    if stream_recv_pending
-        .get(&stream_id)
-        .is_some_and(BTreeMap::is_empty)
-    {
-        stream_recv_pending.remove(&stream_id);
-    }
-    (writes, consumed, *expected)
+    )
 }
 
 #[allow(clippy::too_many_arguments)]

src/bin/dosh-server.rs

@@ -3523,51 +3523,14 @@ async fn flush_stream_pending_data_to_client(
 }
 
 fn accept_stream_data(client: &mut ClientState, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
-    let stream_id = data.stream_id;
-    let expected = client.stream_next_recv_offset.entry(stream_id).or_insert(0);
-    if !dosh::transport::stream_data_within_receive_window(
-        *expected,
+    dosh::transport::accept_stream_data_chunks(
+        &mut client.stream_next_recv_offset,
+        &mut client.stream_recv_pending,
+        data.stream_id,
         data.offset,
-        data.bytes.len(),
+        data.bytes,
         STREAM_INITIAL_WINDOW,
-    ) {
-        return (Vec::new(), 0, *expected);
-    }
-    if data.offset < *expected {
-        return (Vec::new(), 0, *expected);
-    }
-    if data.offset > *expected {
-        client
-            .stream_recv_pending
-            .entry(stream_id)
-            .or_default()
-            .entry(data.offset)
-            .or_insert(data.bytes);
-        return (Vec::new(), 0, *expected);
-    }
-
-    let mut writes = vec![data.bytes];
-    let mut consumed = writes[0].len();
-    *expected = expected.saturating_add(consumed as u64);
-
-    while let Some(bytes) = client
-        .stream_recv_pending
-        .entry(stream_id)
-        .or_default()
-        .remove(expected)
-    {
-        consumed += bytes.len();
-        *expected = expected.saturating_add(bytes.len() as u64);
-        writes.push(bytes);
-    }
-    if client
-        .stream_recv_pending
-        .get(&stream_id)
-        .is_some_and(BTreeMap::is_empty)
-    {
-        client.stream_recv_pending.remove(&stream_id);
-    }
-    (writes, consumed, *expected)
+    )
 }
 
 fn cleanup_client_stream(client: &mut ClientState, stream_id: u64) {

src/transport.rs

@@ -654,49 +654,14 @@ impl StreamMux {
     }
 
     fn accept_data(&mut self, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
-        let stream_id = data.stream_id;
-        let expected = self.next_recv_offset.entry(stream_id).or_insert(0);
-        if !stream_data_within_receive_window(
-            *expected,
+        accept_stream_data_chunks(
+            &mut self.next_recv_offset,
+            &mut self.recv_pending,
+            data.stream_id,
             data.offset,
-            data.bytes.len(),
+            data.bytes,
             self.config.initial_window,
-        ) {
-            return (Vec::new(), 0, *expected);
-        }
-        if data.offset < *expected {
-            return (Vec::new(), 0, *expected);
-        }
-        if data.offset > *expected {
-            self.recv_pending
-                .entry(stream_id)
-                .or_default()
-                .entry(data.offset)
-                .or_insert(data.bytes);
-            return (Vec::new(), 0, *expected);
-        }
-
-        let mut chunks = vec![data.bytes];
-        let mut consumed = chunks[0].len();
-        *expected = expected.saturating_add(consumed as u64);
-        while let Some(bytes) = self
-            .recv_pending
-            .entry(stream_id)
-            .or_default()
-            .remove(expected)
-        {
-            consumed = consumed.saturating_add(bytes.len());
-            *expected = expected.saturating_add(bytes.len() as u64);
-            chunks.push(bytes);
-        }
-        if self
-            .recv_pending
-            .get(&stream_id)
-            .is_some_and(BTreeMap::is_empty)
-        {
-            self.recv_pending.remove(&stream_id);
-        }
-        (chunks, consumed, *expected)
+        )
     }
 
     fn ack_data(&mut self, stream_id: u64, received_offset: u64) -> usize {
@@ -822,6 +787,92 @@ pub fn stream_data_within_receive_window(
     end <= expected_offset.saturating_add(window)
 }
 
+fn trim_stream_data_to_offset(
+    offset: u64,
+    bytes: Vec<u8>,
+    expected_offset: u64,
+) -> Option<(u64, Vec<u8>)> {
+    let len = u64::try_from(bytes.len()).ok()?;
+    let end = offset.checked_add(len)?;
+    if end <= expected_offset {
+        return None;
+    }
+    if offset >= expected_offset {
+        return Some((offset, bytes));
+    }
+    let skip = usize::try_from(expected_offset - offset).ok()?;
+    Some((expected_offset, bytes.get(skip..)?.to_vec()))
+}
+
+fn insert_pending_stream_data(pending: &mut BTreeMap<u64, Vec<u8>>, offset: u64, bytes: Vec<u8>) {
+    match pending.entry(offset) {
+        std::collections::btree_map::Entry::Vacant(entry) => {
+            entry.insert(bytes);
+        }
+        std::collections::btree_map::Entry::Occupied(mut entry) => {
+            if bytes.len() > entry.get().len() {
+                entry.insert(bytes);
+            }
+        }
+    }
+}
+
+pub fn accept_stream_data_chunks(
+    next_recv_offset: &mut HashMap<u64, u64>,
+    recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    stream_id: u64,
+    offset: u64,
+    bytes: Vec<u8>,
+    receive_window: usize,
+) -> (Vec<Vec<u8>>, usize, u64) {
+    let expected = next_recv_offset.entry(stream_id).or_insert(0);
+    let Some((offset, bytes)) = trim_stream_data_to_offset(offset, bytes, *expected) else {
+        return (Vec::new(), 0, *expected);
+    };
+    if !stream_data_within_receive_window(*expected, offset, bytes.len(), receive_window) {
+        return (Vec::new(), 0, *expected);
+    }
+    if offset > *expected {
+        insert_pending_stream_data(recv_pending.entry(stream_id).or_default(), offset, bytes);
+        return (Vec::new(), 0, *expected);
+    }
+
+    let mut chunks = vec![bytes];
+    let mut consumed = chunks[0].len();
+    *expected = expected.saturating_add(consumed as u64);
+    while let Some((&pending_offset, _)) = recv_pending
+        .get(&stream_id)
+        .and_then(|pending| pending.range(..=*expected).next_back())
+    {
+        let Some(bytes) = recv_pending
+            .get_mut(&stream_id)
+            .and_then(|pending| pending.remove(&pending_offset))
+        else {
+            break;
+        };
+        let Some((trimmed_offset, bytes)) =
+            trim_stream_data_to_offset(pending_offset, bytes, *expected)
+        else {
+            continue;
+        };
+        if trimmed_offset > *expected {
+            insert_pending_stream_data(
+                recv_pending.entry(stream_id).or_default(),
+                trimmed_offset,
+                bytes,
+            );
+            break;
+        }
+        consumed = consumed.saturating_add(bytes.len());
+        *expected = expected.saturating_add(bytes.len() as u64);
+        chunks.push(bytes);
+    }
+    if recv_pending.get(&stream_id).is_some_and(BTreeMap::is_empty) {
+        recv_pending.remove(&stream_id);
+    }
+    (chunks, consumed, *expected)
+}
+
 pub fn stream_data_chunk_limit(receive_window: usize) -> usize {
     MAX_STREAM_DATA_BYTES.min(receive_window.max(1))
 }
@@ -1327,6 +1378,79 @@ mod tests {
         assert_eq!(oversized_current.received_offset, 0);
     }
 
+    #[test]
+    fn overlapping_out_of_order_data_trims_deliverable_suffixes() {
+        let mut mux = StreamMux::new(TransportConfig {
+            initial_window: 16,
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 3,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+
+        assert!(
+            mux.handle_data(StreamData {
+                stream_id: 3,
+                offset: 5,
+                bytes: b"FGHIJ".to_vec(),
+            })
+            .unwrap()
+            .unwrap()
+            .chunks
+            .is_empty()
+        );
+        assert!(
+            mux.handle_data(StreamData {
+                stream_id: 3,
+                offset: 3,
+                bytes: b"DEFGH".to_vec(),
+            })
+            .unwrap()
+            .unwrap()
+            .chunks
+            .is_empty()
+        );
+
+        let delivered = mux
+            .handle_data(StreamData {
+                stream_id: 3,
+                offset: 0,
+                bytes: b"ABC".to_vec(),
+            })
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(
+            delivered.chunks.concat(),
+            b"ABCDEFGHIJ",
+            "overlapping pending suffix should be trimmed and delivered"
+        );
+        assert_eq!(delivered.consumed, 10);
+        assert_eq!(delivered.received_offset, 10);
+        assert!(!mux.recv_pending.contains_key(&3));
+    }
+
+    #[test]
+    fn same_offset_pending_data_keeps_longest_range() {
+        let mut next = HashMap::from([(1, 0)]);
+        let mut pending = HashMap::new();
+
+        let first = accept_stream_data_chunks(&mut next, &mut pending, 1, 3, b"DE".to_vec(), 16);
+        assert!(first.0.is_empty());
+        let longer = accept_stream_data_chunks(&mut next, &mut pending, 1, 3, b"DEFG".to_vec(), 16);
+        assert!(longer.0.is_empty());
+
+        let delivered =
+            accept_stream_data_chunks(&mut next, &mut pending, 1, 0, b"ABC".to_vec(), 16);
+        assert_eq!(delivered.0.concat(), b"ABCDEFG");
+        assert_eq!(delivered.1, 7);
+        assert_eq!(delivered.2, 7);
+        assert!(pending.is_empty());
+    }
+
     #[test]
     fn receive_window_helper_rejects_overflowing_ranges() {
         assert!(stream_data_within_receive_window(0, 4, 1, 5));