diff --git a/noq-proto/src/connection/assembler.rs b/noq-proto/src/connection/assembler.rs index ca0a8d9a06..e88ac6ec32 100644 --- a/noq-proto/src/connection/assembler.rs +++ b/noq-proto/src/connection/assembler.rs @@ -1,5 +1,5 @@ use std::{ - cmp::Ordering, + cmp::{Ordering, max}, collections::{BinaryHeap, binary_heap::PeekMut}, mem, }; @@ -107,30 +107,40 @@ impl Assembler { /// This makes sure we're not unnecessarily holding on to many larger allocations. /// We merge contiguous chunks in the process of doing so. fn defragment(&mut self) { + // Chunks smaller than min_chunk_size are merged regardless of fragmentation + let min_chunk_size = max(self.buffered.div_ceil(MAX_CHUNKS), MIN_RETAINED_CHUNK_SIZE); let new = BinaryHeap::with_capacity(self.data.len()); let old = mem::replace(&mut self.data, new); let mut buffers = old.into_sorted_vec(); self.buffered = 0; let mut fragmented_buffered = 0; - let mut offset = self.bytes_read; + // In unordered mode bytes_read counts bytes returned, not a stream offset, so it + // cannot mark where unread data starts. + let start = if self.state.is_ordered() { + self.bytes_read + } else { + 0 + }; + let mut offset = start; for chunk in buffers.iter_mut().rev() { chunk.try_mark_defragment(offset); let size = chunk.bytes.len(); offset = chunk.offset + size as u64; self.buffered += size; - if !chunk.defragmented { + if !chunk.defragmented || size < min_chunk_size { fragmented_buffered += size; } } self.allocated = self.buffered; let mut buffer = BytesMut::with_capacity(fragmented_buffered); - let mut offset = self.bytes_read; + let mut offset = start; for chunk in buffers.into_iter().rev() { - if chunk.defragmented { - // bytes might be empty after try_mark_defragment - if !chunk.bytes.is_empty() { - self.data.push(chunk); - } + // bytes might be empty after try_mark_defragment + if chunk.bytes.is_empty() { + continue; + } + if chunk.defragmented && chunk.bytes.len() >= min_chunk_size { + self.data.push(chunk); continue; } // Overlap is resolved by try_mark_defragment @@ -151,7 +161,12 @@ impl Assembler { // Note: If a packet contains many frames from the same stream, the estimated over-allocation // will be much higher because we are counting the same allocation multiple times. - pub(super) fn insert(&mut self, mut offset: u64, mut bytes: Bytes, allocation_size: usize) { + pub(super) fn insert( + &mut self, + mut offset: u64, + mut bytes: Bytes, + allocation_size: usize, + ) -> Result<(), TooManyChunks> { debug_assert!( bytes.len() <= allocation_size, "allocation_size less than bytes.len(): {:?} < {:?}", @@ -180,7 +195,7 @@ impl Assembler { recvd.insert(range); } else if offset < self.bytes_read { if (offset + bytes.len() as u64) <= self.bytes_read { - return; + return Ok(()); } else { let diff = self.bytes_read - offset; offset += diff; @@ -188,13 +203,13 @@ impl Assembler { } } - if bytes.is_empty() { - return; + // No early return when empty: the dedup loop above may already have pushed chunks. + if !bytes.is_empty() { + let buffer = Buffer::new(offset, bytes, allocation_size); + self.buffered += buffer.bytes.len(); + self.allocated += buffer.allocation_size; + self.data.push(buffer); } - let buffer = Buffer::new(offset, bytes, allocation_size); - self.buffered += buffer.bytes.len(); - self.allocated += buffer.allocation_size; - self.data.push(buffer); // `self.buffered` also counts duplicate bytes, therefore we use // `self.end - self.bytes_read` as an upper bound of buffered unique // bytes. This will cause a defragmentation if the amount of duplicate @@ -209,9 +224,16 @@ impl Assembler { // buffered data. The constants are chosen somewhat arbitrarily and try to // balance between defragmentation overhead and over-allocation. let threshold = 32768.max(buffered * 3 / 2); - if over_allocation > threshold { - self.defragment() + // Small gapped frames hold over-allocation below the threshold, so bound the count too. + if over_allocation > threshold || self.data.len() > COMPACT_THRESHOLD { + self.defragment(); + // ngtcp2 uses a threshold of 4000 -- try to be a little more conservative? + if self.data.len() > MAX_CHUNKS { + return Err(TooManyChunks); + } } + + Ok(()) } /// Number of bytes consumed by the application @@ -341,6 +363,28 @@ impl State { #[derive(Debug)] pub(crate) struct IllegalOrderedRead; +/// Error indicating that too many chunks are buffered due to maliciously small/gapped frames +#[derive(Debug)] +pub(crate) struct TooManyChunks; + +/// Bound on the number of distinct spans kept for a stream +/// +/// Independent of how much memory those spans over-allocate. A frame is rejected only +/// if compaction cannot get the count back down to this. +const MAX_CHUNKS: usize = 1024; + +/// Minimum size of a defragmented chunk that is retained without coalescing. +/// +/// Chunks below this size will be copied on every compaction. 128 bytes balances +/// being cheap to copy while reducing the overhead of many small buffer nodes. +const MIN_RETAINED_CHUNK_SIZE: usize = 128; + +/// Chunk count past which `insert` compacts before deciding whether to reject +/// +/// Above `MAX_CHUNKS` so a flood of mergeable frames cannot force a defragmentation +/// per frame. +const COMPACT_THRESHOLD: usize = 2 * MAX_CHUNKS; + #[cfg(test)] mod test { use super::*; @@ -350,13 +394,13 @@ mod test { fn assemble_ordered() { let mut x = Assembler::new(); assert_matches!(next(&mut x, 32), None); - x.insert(0, Bytes::from_static(b"123"), 3); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); assert_matches!(next(&mut x, 1), Some(ref y) if &y[..] == b"1"); assert_matches!(next(&mut x, 3), Some(ref y) if &y[..] == b"23"); - x.insert(3, Bytes::from_static(b"456"), 3); + x.insert(3, Bytes::from_static(b"456"), 3).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"456"); - x.insert(6, Bytes::from_static(b"789"), 3); - x.insert(9, Bytes::from_static(b"10"), 2); + x.insert(6, Bytes::from_static(b"789"), 3).unwrap(); + x.insert(9, Bytes::from_static(b"10"), 2).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"789"); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"10"); assert_matches!(next(&mut x, 32), None); @@ -366,9 +410,9 @@ mod test { fn assemble_unordered() { let mut x = Assembler::new(); x.ensure_ordering(false).unwrap(); - x.insert(3, Bytes::from_static(b"456"), 3); + x.insert(3, Bytes::from_static(b"456"), 3).unwrap(); assert_matches!(next(&mut x, 32), None); - x.insert(0, Bytes::from_static(b"123"), 3); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123"); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"456"); assert_matches!(next(&mut x, 32), None); @@ -377,8 +421,8 @@ mod test { #[test] fn assemble_duplicate() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"123"), 3); - x.insert(0, Bytes::from_static(b"123"), 3); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123"); assert_matches!(next(&mut x, 32), None); } @@ -386,8 +430,8 @@ mod test { #[test] fn assemble_duplicate_compact() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"123"), 3); - x.insert(0, Bytes::from_static(b"123"), 3); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); x.defragment(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123"); assert_matches!(next(&mut x, 32), None); @@ -396,8 +440,8 @@ mod test { #[test] fn assemble_contained() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"12345"), 5); - x.insert(1, Bytes::from_static(b"234"), 3); + x.insert(0, Bytes::from_static(b"12345"), 5).unwrap(); + x.insert(1, Bytes::from_static(b"234"), 3).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345"); assert_matches!(next(&mut x, 32), None); } @@ -405,8 +449,8 @@ mod test { #[test] fn assemble_contained_compact() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"12345"), 5); - x.insert(1, Bytes::from_static(b"234"), 3); + x.insert(0, Bytes::from_static(b"12345"), 5).unwrap(); + x.insert(1, Bytes::from_static(b"234"), 3).unwrap(); x.defragment(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345"); assert_matches!(next(&mut x, 32), None); @@ -415,8 +459,8 @@ mod test { #[test] fn assemble_contains() { let mut x = Assembler::new(); - x.insert(1, Bytes::from_static(b"234"), 3); - x.insert(0, Bytes::from_static(b"12345"), 5); + x.insert(1, Bytes::from_static(b"234"), 3).unwrap(); + x.insert(0, Bytes::from_static(b"12345"), 5).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345"); assert_matches!(next(&mut x, 32), None); } @@ -424,8 +468,8 @@ mod test { #[test] fn assemble_contains_compact() { let mut x = Assembler::new(); - x.insert(1, Bytes::from_static(b"234"), 3); - x.insert(0, Bytes::from_static(b"12345"), 5); + x.insert(1, Bytes::from_static(b"234"), 3).unwrap(); + x.insert(0, Bytes::from_static(b"12345"), 5).unwrap(); x.defragment(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345"); assert_matches!(next(&mut x, 32), None); @@ -434,8 +478,8 @@ mod test { #[test] fn assemble_overlapping() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"123"), 3); - x.insert(1, Bytes::from_static(b"234"), 3); + x.insert(0, Bytes::from_static(b"123"), 3).unwrap(); + x.insert(1, Bytes::from_static(b"234"), 3).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123"); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"4"); assert_matches!(next(&mut x, 32), None); @@ -444,8 +488,8 @@ mod test { #[test] fn assemble_overlapping_compact() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"123"), 4); - x.insert(1, Bytes::from_static(b"234"), 4); + x.insert(0, Bytes::from_static(b"123"), 4).unwrap(); + x.insert(1, Bytes::from_static(b"234"), 4).unwrap(); x.defragment(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"1234"); assert_matches!(next(&mut x, 32), None); @@ -454,10 +498,10 @@ mod test { #[test] fn assemble_complex() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"1"), 1); - x.insert(2, Bytes::from_static(b"3"), 1); - x.insert(4, Bytes::from_static(b"5"), 1); - x.insert(0, Bytes::from_static(b"123456"), 6); + x.insert(0, Bytes::from_static(b"1"), 1).unwrap(); + x.insert(2, Bytes::from_static(b"3"), 1).unwrap(); + x.insert(4, Bytes::from_static(b"5"), 1).unwrap(); + x.insert(0, Bytes::from_static(b"123456"), 6).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123456"); assert_matches!(next(&mut x, 32), None); } @@ -465,10 +509,10 @@ mod test { #[test] fn assemble_complex_compact() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"1"), 1); - x.insert(2, Bytes::from_static(b"3"), 1); - x.insert(4, Bytes::from_static(b"5"), 1); - x.insert(0, Bytes::from_static(b"123456"), 6); + x.insert(0, Bytes::from_static(b"1"), 1).unwrap(); + x.insert(2, Bytes::from_static(b"3"), 1).unwrap(); + x.insert(4, Bytes::from_static(b"5"), 1).unwrap(); + x.insert(0, Bytes::from_static(b"123456"), 6).unwrap(); x.defragment(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123456"); assert_matches!(next(&mut x, 32), None); @@ -477,19 +521,19 @@ mod test { #[test] fn assemble_old() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"1234"), 4); + x.insert(0, Bytes::from_static(b"1234"), 4).unwrap(); assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"1234"); - x.insert(0, Bytes::from_static(b"1234"), 4); + x.insert(0, Bytes::from_static(b"1234"), 4).unwrap(); assert_matches!(next(&mut x, 32), None); } #[test] fn compact() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"abc"), 4); - x.insert(3, Bytes::from_static(b"def"), 4); - x.insert(9, Bytes::from_static(b"jkl"), 4); - x.insert(12, Bytes::from_static(b"mno"), 4); + x.insert(0, Bytes::from_static(b"abc"), 4).unwrap(); + x.insert(3, Bytes::from_static(b"def"), 4).unwrap(); + x.insert(9, Bytes::from_static(b"jkl"), 4).unwrap(); + x.insert(12, Bytes::from_static(b"mno"), 4).unwrap(); x.defragment(); assert_eq!( next_unordered(&mut x), @@ -504,7 +548,7 @@ mod test { #[test] fn defrag_with_missing_prefix() { let mut x = Assembler::new(); - x.insert(3, Bytes::from_static(b"def"), 3); + x.insert(3, Bytes::from_static(b"def"), 3).unwrap(); x.defragment(); assert_eq!( next_unordered(&mut x), @@ -515,17 +559,17 @@ mod test { #[test] fn defrag_read_chunk() { let mut x = Assembler::new(); - x.insert(3, Bytes::from_static(b"def"), 4); - x.insert(0, Bytes::from_static(b"abc"), 4); - x.insert(7, Bytes::from_static(b"hij"), 4); - x.insert(11, Bytes::from_static(b"lmn"), 4); + x.insert(3, Bytes::from_static(b"def"), 4).unwrap(); + x.insert(0, Bytes::from_static(b"abc"), 4).unwrap(); + x.insert(7, Bytes::from_static(b"hij"), 4).unwrap(); + x.insert(11, Bytes::from_static(b"lmn"), 4).unwrap(); x.defragment(); assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"abcdef"); - x.insert(5, Bytes::from_static(b"fghijklmn"), 9); + x.insert(5, Bytes::from_static(b"fghijklmn"), 9).unwrap(); assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"ghijklmn"); - x.insert(13, Bytes::from_static(b"nopq"), 4); + x.insert(13, Bytes::from_static(b"nopq"), 4).unwrap(); assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"opq"); - x.insert(15, Bytes::from_static(b"pqrs"), 4); + x.insert(15, Bytes::from_static(b"pqrs"), 4).unwrap(); assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"rs"); assert_matches!(x.read(usize::MAX, true), None); } @@ -534,13 +578,13 @@ mod test { fn unordered_happy_path() { let mut x = Assembler::new(); x.ensure_ordering(false).unwrap(); - x.insert(0, Bytes::from_static(b"abc"), 3); + x.insert(0, Bytes::from_static(b"abc"), 3).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(0, Bytes::from_static(b"abc")) ); assert_eq!(x.read(usize::MAX, false), None); - x.insert(3, Bytes::from_static(b"def"), 3); + x.insert(3, Bytes::from_static(b"def"), 3).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(3, Bytes::from_static(b"def")) @@ -552,15 +596,15 @@ mod test { fn unordered_dedup() { let mut x = Assembler::new(); x.ensure_ordering(false).unwrap(); - x.insert(3, Bytes::from_static(b"def"), 3); + x.insert(3, Bytes::from_static(b"def"), 3).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(3, Bytes::from_static(b"def")) ); assert_eq!(x.read(usize::MAX, false), None); - x.insert(0, Bytes::from_static(b"a"), 1); - x.insert(0, Bytes::from_static(b"abcdefghi"), 9); - x.insert(0, Bytes::from_static(b"abcd"), 4); + x.insert(0, Bytes::from_static(b"a"), 1).unwrap(); + x.insert(0, Bytes::from_static(b"abcdefghi"), 9).unwrap(); + x.insert(0, Bytes::from_static(b"abcd"), 4).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(0, Bytes::from_static(b"a")) @@ -574,30 +618,30 @@ mod test { Chunk::new(6, Bytes::from_static(b"ghi")) ); assert_eq!(x.read(usize::MAX, false), None); - x.insert(8, Bytes::from_static(b"ijkl"), 4); + x.insert(8, Bytes::from_static(b"ijkl"), 4).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(9, Bytes::from_static(b"jkl")) ); assert_eq!(x.read(usize::MAX, false), None); - x.insert(12, Bytes::from_static(b"mno"), 3); + x.insert(12, Bytes::from_static(b"mno"), 3).unwrap(); assert_eq!( next_unordered(&mut x), Chunk::new(12, Bytes::from_static(b"mno")) ); assert_eq!(x.read(usize::MAX, false), None); - x.insert(2, Bytes::from_static(b"cde"), 3); + x.insert(2, Bytes::from_static(b"cde"), 3).unwrap(); assert_eq!(x.read(usize::MAX, false), None); } #[test] fn chunks_dedup() { let mut x = Assembler::new(); - x.insert(3, Bytes::from_static(b"def"), 3); + x.insert(3, Bytes::from_static(b"def"), 3).unwrap(); assert_eq!(x.read(usize::MAX, true), None); - x.insert(0, Bytes::from_static(b"a"), 1); - x.insert(1, Bytes::from_static(b"bcdefghi"), 9); - x.insert(0, Bytes::from_static(b"abcd"), 4); + x.insert(0, Bytes::from_static(b"a"), 1).unwrap(); + x.insert(1, Bytes::from_static(b"bcdefghi"), 9).unwrap(); + x.insert(0, Bytes::from_static(b"abcd"), 4).unwrap(); assert_eq!( x.read(usize::MAX, true), Some(Chunk::new(0, Bytes::from_static(b"abcd"))) @@ -607,34 +651,34 @@ mod test { Some(Chunk::new(4, Bytes::from_static(b"efghi"))) ); assert_eq!(x.read(usize::MAX, true), None); - x.insert(8, Bytes::from_static(b"ijkl"), 4); + x.insert(8, Bytes::from_static(b"ijkl"), 4).unwrap(); assert_eq!( x.read(usize::MAX, true), Some(Chunk::new(9, Bytes::from_static(b"jkl"))) ); assert_eq!(x.read(usize::MAX, true), None); - x.insert(12, Bytes::from_static(b"mno"), 3); + x.insert(12, Bytes::from_static(b"mno"), 3).unwrap(); assert_eq!( x.read(usize::MAX, true), Some(Chunk::new(12, Bytes::from_static(b"mno"))) ); assert_eq!(x.read(usize::MAX, true), None); - x.insert(2, Bytes::from_static(b"cde"), 3); + x.insert(2, Bytes::from_static(b"cde"), 3).unwrap(); assert_eq!(x.read(usize::MAX, true), None); } #[test] fn ordered_eager_discard() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"abc"), 3); + x.insert(0, Bytes::from_static(b"abc"), 3).unwrap(); assert_eq!(x.data.len(), 1); assert_eq!( x.read(usize::MAX, true), Some(Chunk::new(0, Bytes::from_static(b"abc"))) ); - x.insert(0, Bytes::from_static(b"ab"), 2); + x.insert(0, Bytes::from_static(b"ab"), 2).unwrap(); assert_eq!(x.data.len(), 0); - x.insert(2, Bytes::from_static(b"cd"), 2); + x.insert(2, Bytes::from_static(b"cd"), 2).unwrap(); assert_eq!( x.data.peek(), Some(&Buffer::new(3, Bytes::from_static(b"d"), 2)) @@ -644,8 +688,8 @@ mod test { #[test] fn ordered_insert_unordered_read() { let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"abc"), 3); - x.insert(0, Bytes::from_static(b"abc"), 3); + x.insert(0, Bytes::from_static(b"abc"), 3).unwrap(); + x.insert(0, Bytes::from_static(b"abc"), 3).unwrap(); x.ensure_ordering(false).unwrap(); assert_eq!( x.read(3, false), @@ -659,8 +703,8 @@ mod test { // Regression test: bytes read in ordered mode should not be returned again in unordered // mode let mut x = Assembler::new(); - x.insert(0, Bytes::from_static(b"a"), 1); - x.insert(0, Bytes::from_static(b"a"), 1); // duplicate + x.insert(0, Bytes::from_static(b"a"), 1).unwrap(); + x.insert(0, Bytes::from_static(b"a"), 1).unwrap(); // duplicate assert_eq!( x.read(1, true), Some(Chunk::new(0, Bytes::from_static(b"a"))) @@ -669,6 +713,125 @@ mod test { assert_eq!(x.read(1, false), None); // should be None, byte 0 already returned } + #[test] + fn ordered_lossless_stream_is_not_rejected() { + const FRAME_LEN: usize = 32; + + let mut x = Assembler::new(); + let data = Bytes::from(vec![0; FRAME_LEN]); + let mut offset = 0u64; + while offset < (2 * COMPACT_THRESHOLD * FRAME_LEN) as u64 { + x.insert(offset, data.clone(), FRAME_LEN) + .expect("contiguous lossless data must not be rejected"); + offset += FRAME_LEN as u64; + } + } + + #[test] + fn defrag_respects_max_chunks_at_rounding_boundary() { + let mut x = Assembler::new(); + let data = Bytes::from(vec![0; MIN_RETAINED_CHUNK_SIZE]); + let mut offset = 0; + // Fill with minimium sized chunks. + for _ in 0..MAX_CHUNKS { + x.insert(offset, data.clone(), data.len()).unwrap(); + offset += data.len() as u64; + } + // Add 1 extra chunk, defragment must coalesce the previous + // chunks to stay within MAX_CHUNKS + x.insert(offset, Bytes::from_static(b"x"), 1).unwrap(); + x.defragment(); + assert!(x.data.len() <= MAX_CHUNKS); + } + + #[test] + fn bounded_chunks_under_low_over_allocation() { + // Gapped frames whose `allocation_size` equals their length hold + // `over_allocation` at zero, so that trigger never fires. + let mut x = Assembler::new(); + // Withhold offset 0 so an ordered reader can never drain anything. + let mut offset = 1u64; + let mut result = Ok(()); + for _ in 0..(MAX_CHUNKS * 8) { + result = x.insert(offset, Bytes::from_static(b"gap"), 3); + if result.is_err() { + break; + } + offset += 3 + 1; // 3 data bytes, 1 byte gap + } + assert_matches!(result, Err(TooManyChunks)); + assert!( + x.data.len() <= COMPACT_THRESHOLD + 1, + "chunk count {} exceeded the bound", + x.data.len() + ); + } + + #[test] + fn bounded_chunks_unordered_overlap_flood() { + // Overlapping frames whose tail is already received: the dedup loop pushes + // the fresh head byte and leaves `bytes` empty, and `end` never rises, so + // the flood costs the peer no flow control. + let mut x = Assembler::new(); + x.ensure_ordering(false).unwrap(); + let top = 1_000_000u64; + x.insert(top, Bytes::from_static(b"ab"), 2).unwrap(); + for k in 0..(4 * MAX_CHUNKS as u64) { + x.insert(top - k - 1, Bytes::from_static(b"ab"), 2).unwrap(); + assert!( + x.data.len() <= COMPACT_THRESHOLD + 1, + "chunk count {} exceeded the bound at k={k}", + x.data.len() + ); + } + } + + /// Count-triggered compaction keeps unread data in unordered mode, where bytes_read is + /// not a stream offset. + #[test] + fn unordered_compaction_keeps_unread_data() { + let mut x = Assembler::new(); + x.ensure_ordering(false).unwrap(); + x.insert(100_000, Bytes::from(vec![0u8; 10_000]), 10_000) + .unwrap(); + assert_eq!(x.read(usize::MAX, false).unwrap().bytes.len(), 10_000); + for i in 0..(COMPACT_THRESHOLD as u64 + 1) { + x.insert(i, Bytes::from_static(b"a"), 1).unwrap(); + } + let mut got = 0; + while let Some(c) = x.read(usize::MAX, false) { + got += c.bytes.len(); + } + assert_eq!(got, COMPACT_THRESHOLD + 1); + } + + #[test] + fn bounded_chunks_duplicate_flood() { + // Duplicates against a stream already at the cap. Ordered mode does not dedup, + // so each one pushes a chunk; they must not be rejected, or compact every frame. + let mut x = Assembler::new(); + // Withhold offset 0 so nothing can be drained. + for i in 0..MAX_CHUNKS as u64 { + x.insert(1 + i * 4, Bytes::from_static(b"abc"), 3).unwrap(); + } + let mut max_len = x.data.len(); + for _ in 0..(3 * MAX_CHUNKS) { + x.insert(1, Bytes::from_static(b"abc"), 3) + .expect("duplicate flood must not be rejected"); + max_len = max_len.max(x.data.len()); + assert!( + x.data.len() <= COMPACT_THRESHOLD + 1, + "chunk count {} exceeded the bound", + x.data.len() + ); + } + // Compacting on every frame would pin the count at `MAX_CHUNKS`. + assert!( + max_len > MAX_CHUNKS, + "buffer compacted on every frame (max observed len {max_len})" + ); + } + fn next_unordered(x: &mut Assembler) -> Chunk { x.read(usize::MAX, false).unwrap() } @@ -779,7 +942,7 @@ mod proptests { match op { Op::Insert { offset, len } => { let bytes = get_slice(&data, offset, len); - asm.insert(offset, bytes, len); + asm.insert(offset, bytes, len).unwrap(); reference.insert(offset, len); } Op::Read { max_len } => { @@ -833,11 +996,7 @@ mod proptests { let expected = reference.ensure_ordering(ordered); prop_assert_eq!(actual, expected, "ensure_ordering result mismatch"); } - Op::Defragment => { - if asm.state.is_ordered() { - asm.defragment(); - } - } + Op::Defragment => asm.defragment(), } } diff --git a/noq-proto/src/connection/mod.rs b/noq-proto/src/connection/mod.rs index 835f022b2f..988258c526 100644 --- a/noq-proto/src/connection/mod.rs +++ b/noq-proto/src/connection/mod.rs @@ -4121,7 +4121,9 @@ impl Connection { crypto_space .crypto_stream - .insert(crypto.offset, crypto.data.clone(), payload_len); + .insert(crypto.offset, crypto.data.clone(), payload_len) + .map_err(|_| TransportError::INTERNAL_ERROR("too many gaps in crypto stream buffer"))?; + while let Some(chunk) = crypto_space.crypto_stream.read(usize::MAX, true) { trace!("consumed {} CRYPTO bytes", chunk.bytes.len()); if self.crypto_state.session.read_handshake(&chunk.bytes)? { diff --git a/noq-proto/src/connection/streams/recv.rs b/noq-proto/src/connection/streams/recv.rs index 462adeb563..1209d1daad 100644 --- a/noq-proto/src/connection/streams/recv.rs +++ b/noq-proto/src/connection/streams/recv.rs @@ -79,7 +79,9 @@ impl Recv { // Don't bother storing data or releasing stream-level flow control credit if the stream's // already stopped if !self.stopped { - self.assembler.insert(frame.offset, frame.data, payload_len); + self.assembler + .insert(frame.offset, frame.data, payload_len) + .map_err(|_| TransportError::INTERNAL_ERROR("too many gaps in stream buffer"))?; } Ok((new_bytes, frame.fin && self.stopped))