diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index d913b6d..4aa74dd 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -8124,48 +8124,14 @@ fn accept_stream_data( stream_recv_pending: &mut HashMap>>, data: StreamData, ) -> (Vec>, 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)] diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 305653c..9f007bd 100644 --- a/src/bin/dosh-server.rs +++ b/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>, 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) { diff --git a/src/transport.rs b/src/transport.rs index 4ac66f5..6e6d1dc 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -654,49 +654,14 @@ impl StreamMux { } fn accept_data(&mut self, data: StreamData) -> (Vec>, 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, + expected_offset: u64, +) -> Option<(u64, Vec)> { + 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>, offset: u64, bytes: Vec) { + 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, + recv_pending: &mut HashMap>>, + stream_id: u64, + offset: u64, + bytes: Vec, + receive_window: usize, +) -> (Vec>, 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));