quinn_proto/connection/
assembler.rs

1use std::{
2    cmp::{Ordering, max},
3    collections::{BinaryHeap, binary_heap::PeekMut},
4    mem,
5};
6
7use bytes::{Buf, Bytes, BytesMut};
8
9use crate::range_set::RangeSet;
10
11/// Helper to assemble unordered stream frames into an ordered stream
12#[derive(Debug, Default)]
13pub(super) struct Assembler {
14    state: State,
15    data: BinaryHeap<Buffer>,
16    /// Total number of buffered bytes, including duplicates in ordered mode.
17    buffered: usize,
18    /// Estimated number of allocated bytes, will never be less than `buffered`.
19    allocated: usize,
20    /// Number of bytes read by the application. When only ordered reads have been used, this is the
21    /// length of the contiguous prefix of the stream which has been consumed by the application,
22    /// aka the stream offset.
23    bytes_read: u64,
24    end: u64,
25}
26
27impl Assembler {
28    pub(super) fn new() -> Self {
29        Self::default()
30    }
31
32    /// Reset to the initial state
33    pub(super) fn reinit(&mut self) {
34        let old_data = mem::take(&mut self.data);
35        *self = Self::default();
36        self.data = old_data;
37        self.data.clear();
38    }
39
40    pub(super) fn ensure_ordering(&mut self, ordered: bool) -> Result<(), IllegalOrderedRead> {
41        if ordered && !self.state.is_ordered() {
42            return Err(IllegalOrderedRead);
43        } else if !ordered && self.state.is_ordered() {
44            // Enter unordered mode
45            if !self.data.is_empty() {
46                // Get rid of possible duplicates
47                self.defragment();
48            }
49            let mut recvd = RangeSet::new();
50            recvd.insert(0..self.bytes_read);
51            for chunk in &self.data {
52                recvd.insert(chunk.offset..chunk.offset + chunk.bytes.len() as u64);
53            }
54            self.state = State::Unordered { recvd };
55        }
56        Ok(())
57    }
58
59    /// Get the the next chunk
60    pub(super) fn read(&mut self, max_length: usize, ordered: bool) -> Option<Chunk> {
61        loop {
62            let mut chunk = self.data.peek_mut()?;
63
64            if ordered {
65                if chunk.offset > self.bytes_read {
66                    // Next chunk is after current read index
67                    return None;
68                } else if (chunk.offset + chunk.bytes.len() as u64) <= self.bytes_read {
69                    // Next chunk is useless as the read index is beyond its end
70                    self.buffered -= chunk.bytes.len();
71                    self.allocated -= chunk.allocation_size;
72                    PeekMut::pop(chunk);
73                    continue;
74                }
75
76                // Determine `start` and `len` of the slice of useful data in chunk
77                let start = (self.bytes_read - chunk.offset) as usize;
78                if start > 0 {
79                    chunk.bytes.advance(start);
80                    chunk.offset += start as u64;
81                    self.buffered -= start;
82                }
83            }
84
85            return Some(if max_length < chunk.bytes.len() {
86                self.bytes_read += max_length as u64;
87                let offset = chunk.offset;
88                chunk.offset += max_length as u64;
89                self.buffered -= max_length;
90                Chunk::new(offset, chunk.bytes.split_to(max_length))
91            } else {
92                self.bytes_read += chunk.bytes.len() as u64;
93                self.buffered -= chunk.bytes.len();
94                self.allocated -= chunk.allocation_size;
95                let chunk = PeekMut::pop(chunk);
96                Chunk::new(chunk.offset, chunk.bytes)
97            });
98        }
99    }
100
101    /// Copy fragmented chunk data to new chunks backed by a single buffer
102    ///
103    /// This makes sure we're not unnecessarily holding on to many larger allocations.
104    /// We merge contiguous chunks in the process of doing so.
105    fn defragment(&mut self) {
106        // Chunks smaller than min_chunk_size are merged regardless of fragmentation
107        let min_chunk_size = max(self.buffered.div_ceil(MAX_CHUNKS), MIN_RETAINED_CHUNK_SIZE);
108        let new = BinaryHeap::with_capacity(self.data.len());
109        let old = mem::replace(&mut self.data, new);
110        let mut buffers = old.into_sorted_vec();
111        self.buffered = 0;
112        let mut fragmented_buffered = 0;
113        let mut offset = 0;
114        for chunk in buffers.iter_mut().rev() {
115            chunk.try_mark_defragment(offset);
116            let size = chunk.bytes.len();
117            offset = chunk.offset + size as u64;
118            self.buffered += size;
119            if !chunk.defragmented || size < min_chunk_size {
120                fragmented_buffered += size;
121            }
122        }
123        self.allocated = self.buffered;
124        let mut buffer = BytesMut::with_capacity(fragmented_buffered);
125        let mut offset = 0;
126        for chunk in buffers.into_iter().rev() {
127            // bytes might be empty after try_mark_defragment
128            if chunk.bytes.is_empty() {
129                continue;
130            }
131            if chunk.defragmented && chunk.bytes.len() >= min_chunk_size {
132                self.data.push(chunk);
133                continue;
134            }
135            // Overlap is resolved by try_mark_defragment
136            if chunk.offset != offset + (buffer.len() as u64) {
137                if !buffer.is_empty() {
138                    self.data
139                        .push(Buffer::new_defragmented(offset, buffer.split().freeze()));
140                }
141                offset = chunk.offset;
142            }
143            buffer.extend_from_slice(&chunk.bytes);
144        }
145        if !buffer.is_empty() {
146            self.data
147                .push(Buffer::new_defragmented(offset, buffer.split().freeze()));
148        }
149    }
150
151    // Note: If a packet contains many frames from the same stream, the estimated over-allocation
152    // will be much higher because we are counting the same allocation multiple times.
153    pub(super) fn insert(
154        &mut self,
155        mut offset: u64,
156        mut bytes: Bytes,
157        allocation_size: usize,
158    ) -> Result<(), TooManyChunks> {
159        debug_assert!(
160            bytes.len() <= allocation_size,
161            "allocation_size less than bytes.len(): {:?} < {:?}",
162            allocation_size,
163            bytes.len()
164        );
165        self.end = self.end.max(offset + bytes.len() as u64);
166        if let State::Unordered { ref mut recvd } = self.state {
167            // Discard duplicate data
168            for duplicate in recvd.replace(offset..offset + bytes.len() as u64) {
169                if duplicate.start > offset {
170                    let buffer = Buffer::new(
171                        offset,
172                        bytes.split_to((duplicate.start - offset) as usize),
173                        allocation_size,
174                    );
175                    self.buffered += buffer.bytes.len();
176                    self.allocated += buffer.allocation_size;
177                    self.data.push(buffer);
178                    offset = duplicate.start;
179                }
180                bytes.advance((duplicate.end - offset) as usize);
181                offset = duplicate.end;
182            }
183        } else if offset < self.bytes_read {
184            if (offset + bytes.len() as u64) <= self.bytes_read {
185                return Ok(());
186            } else {
187                let diff = self.bytes_read - offset;
188                offset += diff;
189                bytes.advance(diff as usize);
190            }
191        }
192
193        // No early return when empty: the dedup loop above may already have pushed chunks.
194        if !bytes.is_empty() {
195            let buffer = Buffer::new(offset, bytes, allocation_size);
196            self.buffered += buffer.bytes.len();
197            self.allocated += buffer.allocation_size;
198            self.data.push(buffer);
199        }
200        // `self.buffered` also counts duplicate bytes, therefore we use
201        // `self.end - self.bytes_read` as an upper bound of buffered unique
202        // bytes. This will cause a defragmentation if the amount of duplicate
203        // bytes exceedes a proportion of the receive window size.
204        let buffered = self.buffered.min((self.end - self.bytes_read) as usize);
205        let over_allocation = self.allocated - buffered;
206        // Rationale: on the one hand, we want to defragment rarely, ideally never
207        // in non-pathological scenarios. However, a pathological or malicious
208        // peer could send us one-byte frames, and since we use reference-counted
209        // buffers in order to prevent copying, this could result in keeping a lot
210        // of memory allocated. This limits over-allocation in proportion to the
211        // buffered data. The constants are chosen somewhat arbitrarily and try to
212        // balance between defragmentation overhead and over-allocation.
213        let threshold = 32768.max(buffered * 3 / 2);
214        // Small gapped frames hold over-allocation below the threshold, so bound the count too.
215        if over_allocation > threshold || self.data.len() > COMPACT_THRESHOLD {
216            self.defragment();
217            // ngtcp2 uses a threshold of 4000 -- try to be a little more conservative?
218            if self.data.len() > MAX_CHUNKS {
219                return Err(TooManyChunks);
220            }
221        }
222
223        Ok(())
224    }
225
226    /// Number of bytes consumed by the application
227    pub(super) fn bytes_read(&self) -> u64 {
228        self.bytes_read
229    }
230
231    /// Discard all buffered data
232    pub(super) fn clear(&mut self) {
233        self.data.clear();
234        self.buffered = 0;
235        self.allocated = 0;
236    }
237}
238
239/// A chunk of data from the receive stream
240#[derive(Debug, PartialEq, Eq)]
241pub struct Chunk {
242    /// The offset in the stream
243    pub offset: u64,
244    /// The contents of the chunk
245    pub bytes: Bytes,
246}
247
248impl Chunk {
249    fn new(offset: u64, bytes: Bytes) -> Self {
250        Self { offset, bytes }
251    }
252}
253
254#[derive(Debug, Eq)]
255struct Buffer {
256    offset: u64,
257    bytes: Bytes,
258    /// Size of the allocation behind `bytes`, if `defragmented == false`.
259    /// Otherwise this will be set to `bytes.len()` by `try_mark_defragment`.
260    /// Will never be less than `bytes.len()`.
261    allocation_size: usize,
262    defragmented: bool,
263}
264
265impl Buffer {
266    /// Constructs a new fragmented Buffer
267    fn new(offset: u64, bytes: Bytes, allocation_size: usize) -> Self {
268        Self {
269            offset,
270            bytes,
271            allocation_size,
272            defragmented: false,
273        }
274    }
275
276    /// Constructs a new defragmented Buffer
277    fn new_defragmented(offset: u64, bytes: Bytes) -> Self {
278        let allocation_size = bytes.len();
279        Self {
280            offset,
281            bytes,
282            allocation_size,
283            defragmented: true,
284        }
285    }
286
287    /// Discards data before `offset` and flags `self` as defragmented if it has good utilization
288    fn try_mark_defragment(&mut self, offset: u64) {
289        let duplicate = offset.saturating_sub(self.offset) as usize;
290        self.offset = self.offset.max(offset);
291        if duplicate >= self.bytes.len() {
292            // All bytes are duplicate
293            self.bytes = Bytes::new();
294            self.defragmented = true;
295            self.allocation_size = 0;
296            return;
297        }
298        self.bytes.advance(duplicate);
299        // Make sure that fragmented buffers with high utilization become defragmented and
300        // defragmented buffers remain defragmented
301        self.defragmented = self.defragmented || self.bytes.len() * 6 / 5 >= self.allocation_size;
302        if self.defragmented {
303            // Make sure that defragmented buffers do not contribute to over-allocation
304            self.allocation_size = self.bytes.len();
305        }
306    }
307}
308
309impl Ord for Buffer {
310    // Invert ordering based on offset (max-heap, min offset first),
311    // prioritize longer chunks at the same offset.
312    fn cmp(&self, other: &Self) -> Ordering {
313        self.offset
314            .cmp(&other.offset)
315            .reverse()
316            .then(self.bytes.len().cmp(&other.bytes.len()))
317    }
318}
319
320impl PartialOrd for Buffer {
321    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
322        Some(self.cmp(other))
323    }
324}
325
326impl PartialEq for Buffer {
327    fn eq(&self, other: &Self) -> bool {
328        (self.offset, self.bytes.len()) == (other.offset, other.bytes.len())
329    }
330}
331
332#[derive(Debug, Default)]
333enum State {
334    #[default]
335    Ordered,
336    Unordered {
337        /// The set of offsets that have been received from the peer, including portions not yet
338        /// read by the application.
339        recvd: RangeSet,
340    },
341}
342
343impl State {
344    fn is_ordered(&self) -> bool {
345        matches!(self, Self::Ordered)
346    }
347}
348
349/// Error indicating that an ordered read was performed on a stream after an unordered read
350#[derive(Debug)]
351pub struct IllegalOrderedRead;
352
353/// Error indicating that too many chunks are buffered due to maliciously small/gapped frames
354#[derive(Debug)]
355pub(crate) struct TooManyChunks;
356
357/// Bound on the number of distinct spans kept for a stream
358///
359/// Independent of how much memory those spans over-allocate. A frame is rejected only
360/// if compaction cannot get the count back down to this.
361const MAX_CHUNKS: usize = 1024;
362
363/// Minimum size of a defragmented chunk that is retained without coalescing.
364///
365/// Chunks below this size will be copied on every compaction. 128 bytes balances
366/// being cheap to copy while reducing the overhead of many small buffer nodes.
367const MIN_RETAINED_CHUNK_SIZE: usize = 128;
368
369/// Chunk count past which `insert` compacts before deciding whether to reject
370///
371/// Above `MAX_CHUNKS` so a flood of mergeable frames cannot force a defragmentation
372/// per frame.
373const COMPACT_THRESHOLD: usize = 2 * MAX_CHUNKS;
374
375#[cfg(test)]
376mod test {
377    use super::*;
378    use assert_matches::assert_matches;
379
380    #[test]
381    fn assemble_ordered() {
382        let mut x = Assembler::new();
383        assert_matches!(next(&mut x, 32), None);
384        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
385        assert_matches!(next(&mut x, 1), Some(ref y) if &y[..] == b"1");
386        assert_matches!(next(&mut x, 3), Some(ref y) if &y[..] == b"23");
387        x.insert(3, Bytes::from_static(b"456"), 3).unwrap();
388        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"456");
389        x.insert(6, Bytes::from_static(b"789"), 3).unwrap();
390        x.insert(9, Bytes::from_static(b"10"), 2).unwrap();
391        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"789");
392        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"10");
393        assert_matches!(next(&mut x, 32), None);
394    }
395
396    #[test]
397    fn assemble_unordered() {
398        let mut x = Assembler::new();
399        x.ensure_ordering(false).unwrap();
400        x.insert(3, Bytes::from_static(b"456"), 3).unwrap();
401        assert_matches!(next(&mut x, 32), None);
402        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
403        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123");
404        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"456");
405        assert_matches!(next(&mut x, 32), None);
406    }
407
408    #[test]
409    fn assemble_duplicate() {
410        let mut x = Assembler::new();
411        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
412        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
413        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123");
414        assert_matches!(next(&mut x, 32), None);
415    }
416
417    #[test]
418    fn assemble_duplicate_compact() {
419        let mut x = Assembler::new();
420        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
421        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
422        x.defragment();
423        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123");
424        assert_matches!(next(&mut x, 32), None);
425    }
426
427    #[test]
428    fn assemble_contained() {
429        let mut x = Assembler::new();
430        x.insert(0, Bytes::from_static(b"12345"), 5).unwrap();
431        x.insert(1, Bytes::from_static(b"234"), 3).unwrap();
432        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345");
433        assert_matches!(next(&mut x, 32), None);
434    }
435
436    #[test]
437    fn assemble_contained_compact() {
438        let mut x = Assembler::new();
439        x.insert(0, Bytes::from_static(b"12345"), 5).unwrap();
440        x.insert(1, Bytes::from_static(b"234"), 3).unwrap();
441        x.defragment();
442        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345");
443        assert_matches!(next(&mut x, 32), None);
444    }
445
446    #[test]
447    fn assemble_contains() {
448        let mut x = Assembler::new();
449        x.insert(1, Bytes::from_static(b"234"), 3).unwrap();
450        x.insert(0, Bytes::from_static(b"12345"), 5).unwrap();
451        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345");
452        assert_matches!(next(&mut x, 32), None);
453    }
454
455    #[test]
456    fn assemble_contains_compact() {
457        let mut x = Assembler::new();
458        x.insert(1, Bytes::from_static(b"234"), 3).unwrap();
459        x.insert(0, Bytes::from_static(b"12345"), 5).unwrap();
460        x.defragment();
461        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"12345");
462        assert_matches!(next(&mut x, 32), None);
463    }
464
465    #[test]
466    fn assemble_overlapping() {
467        let mut x = Assembler::new();
468        x.insert(0, Bytes::from_static(b"123"), 3).unwrap();
469        x.insert(1, Bytes::from_static(b"234"), 3).unwrap();
470        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123");
471        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"4");
472        assert_matches!(next(&mut x, 32), None);
473    }
474
475    #[test]
476    fn assemble_overlapping_compact() {
477        let mut x = Assembler::new();
478        x.insert(0, Bytes::from_static(b"123"), 4).unwrap();
479        x.insert(1, Bytes::from_static(b"234"), 4).unwrap();
480        x.defragment();
481        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"1234");
482        assert_matches!(next(&mut x, 32), None);
483    }
484
485    #[test]
486    fn assemble_complex() {
487        let mut x = Assembler::new();
488        x.insert(0, Bytes::from_static(b"1"), 1).unwrap();
489        x.insert(2, Bytes::from_static(b"3"), 1).unwrap();
490        x.insert(4, Bytes::from_static(b"5"), 1).unwrap();
491        x.insert(0, Bytes::from_static(b"123456"), 6).unwrap();
492        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123456");
493        assert_matches!(next(&mut x, 32), None);
494    }
495
496    #[test]
497    fn assemble_complex_compact() {
498        let mut x = Assembler::new();
499        x.insert(0, Bytes::from_static(b"1"), 1).unwrap();
500        x.insert(2, Bytes::from_static(b"3"), 1).unwrap();
501        x.insert(4, Bytes::from_static(b"5"), 1).unwrap();
502        x.insert(0, Bytes::from_static(b"123456"), 6).unwrap();
503        x.defragment();
504        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"123456");
505        assert_matches!(next(&mut x, 32), None);
506    }
507
508    #[test]
509    fn assemble_old() {
510        let mut x = Assembler::new();
511        x.insert(0, Bytes::from_static(b"1234"), 4).unwrap();
512        assert_matches!(next(&mut x, 32), Some(ref y) if &y[..] == b"1234");
513        x.insert(0, Bytes::from_static(b"1234"), 4).unwrap();
514        assert_matches!(next(&mut x, 32), None);
515    }
516
517    #[test]
518    fn compact() {
519        let mut x = Assembler::new();
520        x.insert(0, Bytes::from_static(b"abc"), 4).unwrap();
521        x.insert(3, Bytes::from_static(b"def"), 4).unwrap();
522        x.insert(9, Bytes::from_static(b"jkl"), 4).unwrap();
523        x.insert(12, Bytes::from_static(b"mno"), 4).unwrap();
524        x.defragment();
525        assert_eq!(
526            next_unordered(&mut x),
527            Chunk::new(0, Bytes::from_static(b"abcdef"))
528        );
529        assert_eq!(
530            next_unordered(&mut x),
531            Chunk::new(9, Bytes::from_static(b"jklmno"))
532        );
533    }
534
535    #[test]
536    fn defrag_with_missing_prefix() {
537        let mut x = Assembler::new();
538        x.insert(3, Bytes::from_static(b"def"), 3).unwrap();
539        x.defragment();
540        assert_eq!(
541            next_unordered(&mut x),
542            Chunk::new(3, Bytes::from_static(b"def"))
543        );
544    }
545
546    #[test]
547    fn defrag_read_chunk() {
548        let mut x = Assembler::new();
549        x.insert(3, Bytes::from_static(b"def"), 4).unwrap();
550        x.insert(0, Bytes::from_static(b"abc"), 4).unwrap();
551        x.insert(7, Bytes::from_static(b"hij"), 4).unwrap();
552        x.insert(11, Bytes::from_static(b"lmn"), 4).unwrap();
553        x.defragment();
554        assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"abcdef");
555        x.insert(5, Bytes::from_static(b"fghijklmn"), 9).unwrap();
556        assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"ghijklmn");
557        x.insert(13, Bytes::from_static(b"nopq"), 4).unwrap();
558        assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"opq");
559        x.insert(15, Bytes::from_static(b"pqrs"), 4).unwrap();
560        assert_matches!(x.read(usize::MAX, true), Some(ref y) if &y.bytes[..] == b"rs");
561        assert_matches!(x.read(usize::MAX, true), None);
562    }
563
564    #[test]
565    fn unordered_happy_path() {
566        let mut x = Assembler::new();
567        x.ensure_ordering(false).unwrap();
568        x.insert(0, Bytes::from_static(b"abc"), 3).unwrap();
569        assert_eq!(
570            next_unordered(&mut x),
571            Chunk::new(0, Bytes::from_static(b"abc"))
572        );
573        assert_eq!(x.read(usize::MAX, false), None);
574        x.insert(3, Bytes::from_static(b"def"), 3).unwrap();
575        assert_eq!(
576            next_unordered(&mut x),
577            Chunk::new(3, Bytes::from_static(b"def"))
578        );
579        assert_eq!(x.read(usize::MAX, false), None);
580    }
581
582    #[test]
583    fn unordered_dedup() {
584        let mut x = Assembler::new();
585        x.ensure_ordering(false).unwrap();
586        x.insert(3, Bytes::from_static(b"def"), 3).unwrap();
587        assert_eq!(
588            next_unordered(&mut x),
589            Chunk::new(3, Bytes::from_static(b"def"))
590        );
591        assert_eq!(x.read(usize::MAX, false), None);
592        x.insert(0, Bytes::from_static(b"a"), 1).unwrap();
593        x.insert(0, Bytes::from_static(b"abcdefghi"), 9).unwrap();
594        x.insert(0, Bytes::from_static(b"abcd"), 4).unwrap();
595        assert_eq!(
596            next_unordered(&mut x),
597            Chunk::new(0, Bytes::from_static(b"a"))
598        );
599        assert_eq!(
600            next_unordered(&mut x),
601            Chunk::new(1, Bytes::from_static(b"bc"))
602        );
603        assert_eq!(
604            next_unordered(&mut x),
605            Chunk::new(6, Bytes::from_static(b"ghi"))
606        );
607        assert_eq!(x.read(usize::MAX, false), None);
608        x.insert(8, Bytes::from_static(b"ijkl"), 4).unwrap();
609        assert_eq!(
610            next_unordered(&mut x),
611            Chunk::new(9, Bytes::from_static(b"jkl"))
612        );
613        assert_eq!(x.read(usize::MAX, false), None);
614        x.insert(12, Bytes::from_static(b"mno"), 3).unwrap();
615        assert_eq!(
616            next_unordered(&mut x),
617            Chunk::new(12, Bytes::from_static(b"mno"))
618        );
619        assert_eq!(x.read(usize::MAX, false), None);
620        x.insert(2, Bytes::from_static(b"cde"), 3).unwrap();
621        assert_eq!(x.read(usize::MAX, false), None);
622    }
623
624    #[test]
625    fn chunks_dedup() {
626        let mut x = Assembler::new();
627        x.insert(3, Bytes::from_static(b"def"), 3).unwrap();
628        assert_eq!(x.read(usize::MAX, true), None);
629        x.insert(0, Bytes::from_static(b"a"), 1).unwrap();
630        x.insert(1, Bytes::from_static(b"bcdefghi"), 9).unwrap();
631        x.insert(0, Bytes::from_static(b"abcd"), 4).unwrap();
632        assert_eq!(
633            x.read(usize::MAX, true),
634            Some(Chunk::new(0, Bytes::from_static(b"abcd")))
635        );
636        assert_eq!(
637            x.read(usize::MAX, true),
638            Some(Chunk::new(4, Bytes::from_static(b"efghi")))
639        );
640        assert_eq!(x.read(usize::MAX, true), None);
641        x.insert(8, Bytes::from_static(b"ijkl"), 4).unwrap();
642        assert_eq!(
643            x.read(usize::MAX, true),
644            Some(Chunk::new(9, Bytes::from_static(b"jkl")))
645        );
646        assert_eq!(x.read(usize::MAX, true), None);
647        x.insert(12, Bytes::from_static(b"mno"), 3).unwrap();
648        assert_eq!(
649            x.read(usize::MAX, true),
650            Some(Chunk::new(12, Bytes::from_static(b"mno")))
651        );
652        assert_eq!(x.read(usize::MAX, true), None);
653        x.insert(2, Bytes::from_static(b"cde"), 3).unwrap();
654        assert_eq!(x.read(usize::MAX, true), None);
655    }
656
657    #[test]
658    fn ordered_eager_discard() {
659        let mut x = Assembler::new();
660        x.insert(0, Bytes::from_static(b"abc"), 3).unwrap();
661        assert_eq!(x.data.len(), 1);
662        assert_eq!(
663            x.read(usize::MAX, true),
664            Some(Chunk::new(0, Bytes::from_static(b"abc")))
665        );
666        x.insert(0, Bytes::from_static(b"ab"), 2).unwrap();
667        assert_eq!(x.data.len(), 0);
668        x.insert(2, Bytes::from_static(b"cd"), 2).unwrap();
669        assert_eq!(
670            x.data.peek(),
671            Some(&Buffer::new(3, Bytes::from_static(b"d"), 2))
672        );
673    }
674
675    #[test]
676    fn ordered_insert_unordered_read() {
677        let mut x = Assembler::new();
678        x.insert(0, Bytes::from_static(b"abc"), 3).unwrap();
679        x.insert(0, Bytes::from_static(b"abc"), 3).unwrap();
680        x.ensure_ordering(false).unwrap();
681        assert_eq!(
682            x.read(3, false),
683            Some(Chunk::new(0, Bytes::from_static(b"abc")))
684        );
685        assert_eq!(x.read(3, false), None);
686    }
687
688    #[test]
689    fn ordered_lossless_stream_is_not_rejected() {
690        const FRAME_LEN: usize = 32;
691
692        let mut x = Assembler::new();
693        let data = Bytes::from(vec![0; FRAME_LEN]);
694        let mut offset = 0u64;
695        while offset < (2 * COMPACT_THRESHOLD * FRAME_LEN) as u64 {
696            x.insert(offset, data.clone(), FRAME_LEN)
697                .expect("contiguous lossless data must not be rejected");
698            offset += FRAME_LEN as u64;
699        }
700    }
701
702    #[test]
703    fn defrag_respects_max_chunks_at_rounding_boundary() {
704        let mut x = Assembler::new();
705        let data = Bytes::from(vec![0; MIN_RETAINED_CHUNK_SIZE]);
706        let mut offset = 0;
707        // Fill with minimium sized chunks.
708        for _ in 0..MAX_CHUNKS {
709            x.insert(offset, data.clone(), data.len()).unwrap();
710            offset += data.len() as u64;
711        }
712        // Add 1 extra chunk, defragment must coalesce the previous
713        // chunks to stay within MAX_CHUNKS
714        x.insert(offset, Bytes::from_static(b"x"), 1).unwrap();
715        x.defragment();
716        assert!(x.data.len() <= MAX_CHUNKS);
717    }
718
719    #[test]
720    fn bounded_chunks_under_low_over_allocation() {
721        // Gapped frames whose `allocation_size` equals their length hold
722        // `over_allocation` at zero, so that trigger never fires.
723        let mut x = Assembler::new();
724        // Withhold offset 0 so an ordered reader can never drain anything.
725        let mut offset = 1u64;
726        let mut result = Ok(());
727        for _ in 0..(MAX_CHUNKS * 8) {
728            result = x.insert(offset, Bytes::from_static(b"gap"), 3);
729            if result.is_err() {
730                break;
731            }
732            offset += 3 + 1; // 3 data bytes, 1 byte gap
733        }
734        assert_matches!(result, Err(TooManyChunks));
735        assert!(
736            x.data.len() <= COMPACT_THRESHOLD + 1,
737            "chunk count {} exceeded the bound",
738            x.data.len()
739        );
740    }
741
742    #[test]
743    fn bounded_chunks_unordered_overlap_flood() {
744        // Overlapping frames whose tail is already received: the dedup loop pushes
745        // the fresh head byte and leaves `bytes` empty, and `end` never rises, so
746        // the flood costs the peer no flow control.
747        let mut x = Assembler::new();
748        x.ensure_ordering(false).unwrap();
749        let top = 1_000_000u64;
750        x.insert(top, Bytes::from_static(b"ab"), 2).unwrap();
751        for k in 0..(4 * MAX_CHUNKS as u64) {
752            x.insert(top - k - 1, Bytes::from_static(b"ab"), 2).unwrap();
753            assert!(
754                x.data.len() <= COMPACT_THRESHOLD + 1,
755                "chunk count {} exceeded the bound at k={k}",
756                x.data.len()
757            );
758        }
759    }
760
761    #[test]
762    fn bounded_chunks_duplicate_flood() {
763        // Duplicates against a stream already at the cap. Ordered mode does not dedup,
764        // so each one pushes a chunk; they must not be rejected, or compact every frame.
765        let mut x = Assembler::new();
766        // Withhold offset 0 so nothing can be drained.
767        for i in 0..MAX_CHUNKS as u64 {
768            x.insert(1 + i * 4, Bytes::from_static(b"abc"), 3).unwrap();
769        }
770        let mut max_len = x.data.len();
771        for _ in 0..(3 * MAX_CHUNKS) {
772            x.insert(1, Bytes::from_static(b"abc"), 3)
773                .expect("duplicate flood must not be rejected");
774            max_len = max_len.max(x.data.len());
775            assert!(
776                x.data.len() <= COMPACT_THRESHOLD + 1,
777                "chunk count {} exceeded the bound",
778                x.data.len()
779            );
780        }
781        // Compacting on every frame would pin the count at `MAX_CHUNKS`.
782        assert!(
783            max_len > MAX_CHUNKS,
784            "buffer compacted on every frame (max observed len {max_len})"
785        );
786    }
787
788    fn next_unordered(x: &mut Assembler) -> Chunk {
789        x.read(usize::MAX, false).unwrap()
790    }
791
792    fn next(x: &mut Assembler, size: usize) -> Option<Bytes> {
793        x.read(size, true).map(|chunk| chunk.bytes)
794    }
795}
quinn_proto/connection/assembler.rs

quinn_proto/connection/
assembler.rs
