quinn-proto: make sure defragmented buffers don't contain duplicate data

Author: geieredgar

quinn-proto/src/connection/assembler.rs

@@ -2,6 +2,7 @@ use std::{
     cmp::{max, min, Ordering},
     collections::{binary_heap::PeekMut, BinaryHeap},
     mem,
+    num::NonZeroUsize,
 };
 
 use bytes::{Buf, Bytes, BytesMut};
@@ -32,9 +33,9 @@ impl Assembler {
             return Err(IllegalOrderedRead);
         } else if !ordered && self.state.is_ordered() {
             // Enter unordered mode
-            let mut recvd = self.deduplicate();
-            recvd.insert(0..self.bytes_read);
-            self.state = State::Unordered { recvd };
+            self.state = State::Unordered {
+                recvd: self.deduplicate(),
+            };
         }
         Ok(())
     }
@@ -54,7 +55,7 @@ impl Assembler {
                 } else if (chunk.offset + chunk.bytes.len() as u64) <= self.bytes_read {
                     // Next chunk is useless as the read index is beyond its end
                     self.buffered -= chunk.size;
-                    self.allocated -= chunk.allocation_size;
+                    self.allocated -= chunk.allocation_size();
                     PeekMut::pop(chunk);
                     continue;
                 }
@@ -75,7 +76,7 @@ impl Assembler {
             } else {
                 self.bytes_read += chunk.bytes.len() as u64;
                 self.buffered -= chunk.size;
-                self.allocated -= chunk.allocation_size;
+                self.allocated -= chunk.allocation_size();
                 let chunk = PeekMut::pop(chunk);
                 Chunk::new(chunk.offset, chunk.bytes)
             });
@@ -88,10 +89,13 @@ impl Assembler {
     // counter to the new number of chunks left in the heap so that we can decide
     // when to defragment the queue again if necessary.
     fn defragment(&mut self) {
+        if self.state.is_ordered() {
+            self.deduplicate();
+        }
         let fragmented_buffered = self
             .data
             .iter()
-            .filter(|c| c.is_fragmented())
+            .filter(|c| c.should_defragment())
             .map(|c| c.bytes.len())
             .sum::<usize>();
         let mut buffer = BytesMut::with_capacity(fragmented_buffered);
@@ -104,24 +108,25 @@ impl Assembler {
 
         let new = BinaryHeap::with_capacity(self.data.len());
         let old = mem::replace(&mut self.data, new);
+        self.buffered = fragmented_buffered;
         for chunk in old.into_sorted_vec().into_iter().rev() {
-            if !chunk.is_fragmented() {
+            if chunk.is_defragmented() {
+                self.buffered += chunk.size;
                 self.data.push(chunk);
                 continue;
+            } else if !chunk.should_defragment() {
+                self.buffered += chunk.bytes.len();
+                self.data
+                    .push(Buffer::new_defragmented(chunk.offset, chunk.bytes));
+                continue;
             }
-            let end = offset + (buffer.len() as u64);
-            if let Some(overlap) = end.checked_sub(chunk.offset) {
-                if let Some(bytes) = chunk.bytes.get(overlap as usize..) {
-                    buffer.extend_from_slice(bytes);
-                }
-            } else {
+            if chunk.offset != offset + (buffer.len() as u64) {
                 self.data
                     .push(Buffer::new_defragmented(offset, buffer.split().freeze()));
                 offset = chunk.offset;
-                buffer.extend_from_slice(&chunk.bytes);
             }
+            buffer.extend_from_slice(&chunk.bytes);
         }
-
         self.data
             .push(Buffer::new_defragmented(offset, buffer.split().freeze()));
         self.allocated = self.buffered;
@@ -130,9 +135,19 @@ impl Assembler {
     fn deduplicate(&mut self) -> RangeSet {
         debug_assert!(self.state.is_ordered());
         let mut ranges = RangeSet::new();
+        ranges.insert(0..self.bytes_read);
         let new = BinaryHeap::with_capacity(self.data.len());
-        let old = mem::replace(&mut self.data, new);
-        for mut buffer in old.into_sorted_vec().into_iter().rev() {
+        let old = mem::replace(&mut self.data, new).into_vec();
+        for buffer in old.iter() {
+            if buffer.is_defragmented() {
+                ranges.insert(buffer.offset..buffer.offset + buffer.bytes.len() as u64);
+            }
+        }
+        for mut buffer in old.into_iter() {
+            if buffer.is_defragmented() {
+                self.data.push(buffer);
+                continue;
+            }
             for duplicate in
                 ranges.replace(buffer.offset..buffer.offset + buffer.bytes.len() as u64)
             {
@@ -142,9 +157,10 @@ impl Assembler {
                         buffer
                             .bytes
                             .split_to((duplicate.start - buffer.offset) as usize),
-                        buffer.allocation_size,
+                        buffer.allocation_size(),
+                        false,
                     );
-                    self.allocated += new_buffer.allocation_size;
+                    self.allocated += new_buffer.allocation_size();
                     self.data.push(new_buffer);
                 }
                 let discarded = (duplicate.end - duplicate.start) as usize;
@@ -154,13 +170,15 @@ impl Assembler {
                 buffer.offset = duplicate.end;
             }
             if buffer.bytes.is_empty() {
-                self.allocated -= buffer.allocation_size;
+                self.allocated -= buffer.allocation_size();
                 self.buffered -= buffer.size;
             } else {
+                let allocation_size = buffer.allocation_size();
                 self.data.push(Buffer::new(
                     buffer.offset,
                     buffer.bytes,
-                    buffer.allocation_size,
+                    allocation_size,
+                    false,
                 ));
             }
         }
@@ -185,9 +203,10 @@ impl Assembler {
                         offset,
                         bytes.split_to((duplicate.start - offset) as usize),
                         allocation_size,
+                        false,
                     );
                     self.buffered += buffer.size;
-                    self.allocated += buffer.allocation_size;
+                    self.allocated += buffer.allocation_size();
                     self.data.push(buffer);
                     offset = duplicate.start;
                 }
@@ -207,9 +226,9 @@ impl Assembler {
         if bytes.is_empty() {
             return;
         }
-        let buffer = Buffer::new(offset, bytes, allocation_size);
+        let buffer = Buffer::new(offset, bytes, allocation_size, self.state.is_ordered());
         self.buffered += buffer.size;
-        self.allocated += buffer.allocation_size;
+        self.allocated += buffer.allocation_size();
         self.data.push(buffer);
         // Rationale: on the one hand, we want to defragment rarely, ideally never
         // in non-pathological scenarios. However, a pathological or malicious
@@ -222,9 +241,6 @@ impl Assembler {
         let over_allocation = self.allocated - buffered;
         let threshold = max(buffered * 3 / 2, 32 * 1024);
         if over_allocation > threshold {
-            if self.state.is_ordered() && self.buffered > buffered {
-                self.deduplicate();
-            }
             self.defragment()
         }
     }
@@ -266,42 +282,53 @@ struct Buffer {
     offset: u64,
     bytes: Bytes,
     size: usize,
-    allocation_size: usize,
+    state: BufferState,
 }
 
 impl Buffer {
     /// Constructs a new, possibly fragmented Buffer
-    fn new(offset: u64, bytes: Bytes, allocation_size: usize) -> Self {
+    fn new(offset: u64, bytes: Bytes, allocation_size: usize, possibly_duplicate: bool) -> Self {
         let size = bytes.len();
-        // Treat buffers with small over-allocation as defragmented
-        let threshold = size * 6 / 5;
-        let allocation_size = if allocation_size > threshold {
-            allocation_size
-        } else {
-            size
-        };
         Self {
             offset,
             bytes,
             size,
-            allocation_size,
+            state: BufferState::new(size, allocation_size, possibly_duplicate),
         }
     }
 
-    /// Constructs a new Buffer that is not fragmented
+    /// Constructs a new Buffer that is considered to be defragmented
     fn new_defragmented(offset: u64, bytes: Bytes) -> Self {
         let size = bytes.len();
         Self {
             offset,
             bytes,
             size,
-            allocation_size: size,
+            state: BufferState::Defragmented,
+        }
+    }
+
+    /// Returns `true` if the buffer is defragmented
+    fn is_defragmented(&self) -> bool {
+        self.state == BufferState::Defragmented
+    }
+
+    /// Returns the size of the associated allocation
+    fn allocation_size(&self) -> usize {
+        match self.state {
+            BufferState::Fragmented(allocation_size) => allocation_size.get(),
+            BufferState::Defragmented => self.size,
         }
     }
 
-    /// Returns `true` if the buffer is fragmented
-    fn is_fragmented(&self) -> bool {
-        self.size < self.allocation_size
+    /// Returns `true` if the buffer should be defragmented
+    fn should_defragment(&self) -> bool {
+        match self.state {
+            BufferState::Fragmented(allocation_size) => {
+                BufferState::significant_over_allocation(self.size, allocation_size.get())
+            }
+            BufferState::Defragmented => false,
+        }
     }
 }
 
@@ -328,6 +355,30 @@ impl PartialEq for Buffer {
     }
 }
 
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
+enum BufferState {
+    Fragmented(NonZeroUsize),
+    Defragmented,
+}
+
+impl BufferState {
+    fn new(size: usize, allocation_size: usize, possibly_duplicate: bool) -> Self {
+        if possibly_duplicate || Self::significant_over_allocation(size, allocation_size) {
+            match NonZeroUsize::new(allocation_size) {
+                Some(value) => Self::Fragmented(value),
+                None => Self::Defragmented,
+            }
+        } else {
+            Self::Defragmented
+        }
+    }
+
+    /// Returns `true` if the buffer size and allocation size warrants a copy
+    fn significant_over_allocation(size: usize, allocation_size: usize) -> bool {
+        size * 6 / 5 < allocation_size
+    }
+}
+
 #[derive(Debug)]
 enum State {
     Ordered,
@@ -650,12 +701,7 @@ mod test {
         x.insert(2, Bytes::from_static(b"cd"), 2);
         assert_eq!(
             x.data.peek(),
-            Some(&Buffer {
-                offset: 3,
-                bytes: Bytes::from_static(b"d"),
-                size: 1,
-                allocation_size: 2
-            })
+            Some(&Buffer::new(3, Bytes::from_static(b"d"), 2, true))
         );
     }