fix IoSlice cursor bug

aws · Feb 8, 2024 · 476e4da · 476e4da
1 parent d6d0ca6
commit 476e4da
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/quic/s2n-quic-core/src/buffer/reader/storage/buf.rs b/quic/s2n-quic-core/src/buffer/reader/storage/buf.rs
@@ -88,8 +88,13 @@ where
             match chunk_len.cmp(&dest.remaining_capacity()) {
                 // if there's more chunks left, then copy this one out and keep going
                 Ordering::Less if self.buf.remaining() > chunk_len => {
-                    dest.put_slice(self.buf.chunk());
-                    self.buf.advance(chunk_len);
+                    if Dest::SPECIALIZES_BYTES {
+                        let chunk = self.buf.copy_to_bytes(chunk_len);
+                        dest.put_bytes(chunk);
+                    } else {
+                        dest.put_slice(self.buf.chunk());
+                        self.buf.advance(chunk_len);
+                    }
                     continue;
                 }
                 Ordering::Less | Ordering::Equal => {
@@ -120,8 +125,13 @@ where
 
             ensure!(len > 0, Ok(()));
 
-            dest.put_slice(&chunk[..len]);
-            self.buf.advance(len);
+            if Dest::SPECIALIZES_BYTES {
+                let chunk = self.buf.copy_to_bytes(len);
+                dest.put_bytes(chunk);
+            } else {
+                dest.put_slice(&chunk[..len]);
+                self.buf.advance(len);
+            }
         }
     }
 }

diff --git a/quic/s2n-quic-core/src/buffer/reader/storage/io_slice.rs b/quic/s2n-quic-core/src/buffer/reader/storage/io_slice.rs
@@ -24,7 +24,7 @@ where
     #[inline]
     pub fn new(buf: &'a [T]) -> Self {
         let mut len = 0;
-        let mut first_non_empty = 0;
+        let mut first_non_empty = usize::MAX;
         let mut last_non_empty = 0;
 
         // find the total length and the first non-empty slice
@@ -56,6 +56,7 @@ where
             buf,
         };
         slice.advance_buf_once();
+        slice.invariants();
         slice
     }
 
@@ -87,15 +88,8 @@ where
 
     #[inline]
     fn set_len(&mut self, len: usize) {
-        if cfg!(debug_assertions) {
-            // make sure the computed len matches the actual remaining len
-            let mut computed = self.head.len();
-            for buf in self.buf.iter() {
-                computed += buf.len();
-            }
-            assert_eq!(len, computed);
-        }
         self.len = len;
+        self.invariants();
     }
 
     #[inline]
@@ -136,6 +130,24 @@ where
             }
         }
     }
+
+    #[inline(always)]
+    fn invariants(&self) {
+        #[cfg(debug_assertions)]
+        {
+            // make sure the computed len matches the actual remaining len
+            let mut computed = self.head.len();
+            for buf in self.buf.iter() {
+                computed += buf.len();
+            }
+            assert_eq!(self.len, computed);
+
+            if self.head.is_empty() {
+                assert!(self.buf.is_empty());
+                assert_eq!(self.len, 0);
+            }
+        }
+    }
 }
 
 impl<'a, T> bytes::Buf for IoSlice<'a, T>