fix(utils): fix some bugs in BufPoolWriter/BufList (#2634)

fix bugs introduced by feat(buf): add BufPool, BufMargins, and BufList for zero-copy packet buffering #2625
cf. perf(quic-proxy): use BufPool and BufMargins in QuicSocket #2626 (comment)

easytier-core/src/tunnel/buf.rs

@@ -15,23 +15,24 @@ impl<T: Buf> BufList<T> {
         }
     }
 
-    #[inline]
+    #[inline(always)]
     pub fn push(&mut self, buf: T) {
-        debug_assert!(buf.has_remaining());
-        self.bufs.push_back(buf);
+        if buf.has_remaining() {
+            self.bufs.push_back(buf);
+        }
     }
 
-    #[inline]
+    #[inline(always)]
     pub fn pop(&mut self) -> Option<T> {
         self.bufs.pop_front()
     }
 
-    #[inline]
+    #[inline(always)]
     pub fn len(&self) -> usize {
         self.bufs.len()
     }
 
-    #[inline]
+    #[inline(always)]
     pub fn is_empty(&self) -> bool {
         self.bufs.is_empty()
     }
@@ -39,17 +40,15 @@ impl<T: Buf> BufList<T> {
 
 impl<T: Buf> Extend<T> for BufList<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
-        self.bufs.extend(
-            iter.into_iter()
-                .inspect(|buf| debug_assert!(buf.has_remaining())),
-        );
+        self.bufs
+            .extend(iter.into_iter().filter(Buf::has_remaining));
     }
 }
 
 impl<T: Buf> Buf for BufList<T> {
     #[inline]
     fn remaining(&self) -> usize {
-        self.bufs.iter().map(|buf| buf.remaining()).sum()
+        self.bufs.iter().map(Buf::remaining).sum()
     }
 
     #[inline]
@@ -59,33 +58,30 @@ impl<T: Buf> Buf for BufList<T> {
 
     #[inline]
     fn chunks_vectored<'t>(&'t self, dst: &mut [IoSlice<'t>]) -> usize {
-        if dst.is_empty() {
-            return 0;
-        }
         let mut vecs = 0;
+
         for buf in &self.bufs {
-            vecs += buf.chunks_vectored(&mut dst[vecs..]);
-            if vecs == dst.len() {
+            let n = buf.chunks_vectored(&mut dst[vecs..]);
+            vecs += n;
+
+            if dst[vecs - n..vecs].iter().map(|s| s.len()).sum::<usize>() < buf.remaining() {
                 break;
             }
         }
+
         vecs
     }
 
     #[inline]
     fn advance(&mut self, mut cnt: usize) {
         while cnt > 0 {
-            {
-                let front = &mut self.bufs[0];
-                let rem = front.remaining();
-                if rem > cnt {
-                    front.advance(cnt);
-                    return;
-                } else {
-                    front.advance(rem);
-                    cnt -= rem;
-                }
+            let front = &mut self.bufs[0];
+            let rem = front.remaining();
+            front.advance(cnt.min(rem));
+            if rem > cnt {
+                return;
             }
+            cnt -= rem;
             self.bufs.pop_front();
         }
     }
@@ -94,20 +90,21 @@ impl<T: Buf> Buf for BufList<T> {
     fn copy_to_bytes(&mut self, len: usize) -> Bytes {
         // Our inner buffer may have an optimized version of copy_to_bytes, and if the whole
         // request can be fulfilled by the front buffer, we can take advantage.
-        match self.bufs.front_mut() {
-            Some(front) if front.remaining() == len => {
-                let b = front.copy_to_bytes(len);
+        if let Some(front) = self.bufs.front_mut()
+            && let rem = front.remaining()
+            && len <= rem
+        {
+            let bytes = front.copy_to_bytes(len);
+            if len == rem {
                 self.bufs.pop_front();
-                b
-            }
-            Some(front) if front.remaining() > len => front.copy_to_bytes(len),
-            _ => {
-                assert!(len <= self.remaining(), "`len` greater than remaining");
-                let mut bm = BytesMut::with_capacity(len);
-                bm.put(self.take(len));
-                bm.freeze()
             }
+            return bytes;
         }
+
+        assert!(len <= self.remaining());
+        let mut bytes = BytesMut::with_capacity(len);
+        bytes.put(self.take(len));
+        bytes.freeze()
     }
 }
 
@@ -125,4 +122,35 @@ mod tests {
         assert_eq!(&bytes[..], b"hello world");
         assert_eq!(list.remaining(), 0);
     }
+
+    #[test]
+    fn test_chunks_vectored_partial_inner_buf() {
+        #[derive(Debug)]
+        struct PartialBuf(&'static [u8]);
+
+        impl Buf for PartialBuf {
+            fn remaining(&self) -> usize {
+                self.0.len()
+            }
+
+            fn chunk(&self) -> &[u8] {
+                &self.0[..self.0.len().min(1)]
+            }
+
+            fn advance(&mut self, cnt: usize) {
+                self.0 = &self.0[cnt..];
+            }
+        }
+
+        let mut list = BufList::new();
+        list.push(PartialBuf(b"ab"));
+        list.push(PartialBuf(b"c"));
+
+        let mut dst = [IoSlice::new(&[]); 3];
+
+        let n = list.chunks_vectored(&mut dst);
+
+        assert_eq!(n, 1);
+        assert_eq!(&*dst[0], b"a");
+    }
 }

easytier/src/utils/buf.rs

@@ -14,10 +14,17 @@ pub struct BufMargins {
 impl BufMargins {
     #[inline(always)]
     pub fn size(&self) -> usize {
-        self.header + self.trailer
+        self.header.checked_add(self.trailer).unwrap()
     }
 }
 
+/// A reusable packet buffer pool.
+///
+/// # Initialization contract
+///
+/// Buffers may contain uninitialized bytes in their [`BufMargins::header`] and
+/// [`BufMargins::trailer`] regions. These regions must be fully initialized
+/// before they are read or otherwise consumed.
 #[derive(Debug)]
 pub struct BufPool {
     pool: BytesMut,
@@ -40,6 +47,10 @@ impl BufPool {
         }
     }
 
+    /// Returns the accumulated buffer.
+    ///
+    /// The caller must ensure all reserved header and trailer regions have been
+    /// initialized before reading or consuming the returned buffer.
     #[inline(always)]
     pub fn split(&mut self) -> BytesMut {
         self.pool.split()
@@ -47,8 +58,10 @@ impl BufPool {
 
     #[inline]
     pub fn write(&mut self, chunk: &[u8], margins: BufMargins) {
-        let len = margins.size() + chunk.len();
+        let len = margins.size().checked_add(chunk.len()).unwrap();
         self.reserve(len);
+        // Header and trailer are intentionally left uninitialized. Users of this
+        // buffer must initialize both regions before reading or consuming it.
         unsafe {
             copy_nonoverlapping(
                 chunk.as_ptr(),
@@ -59,6 +72,10 @@ impl BufPool {
         }
     }
 
+    /// Writes `chunk` with the requested margins and returns the resulting buffer.
+    ///
+    /// The header and trailer regions are left uninitialized and must be filled
+    /// before the returned buffer is read or consumed.
     #[inline(always)]
     pub fn buf(&mut self, chunk: &[u8], margins: BufMargins) -> BytesMut {
         self.write(chunk, margins);
@@ -77,6 +94,12 @@ impl BufPool {
     }
 }
 
+/// A writer into a [`BufPool`].
+///
+/// # Initialization contract
+///
+/// Committing data does not initialize the reserved header or trailer regions.
+/// They must be fully initialized before the resulting buffer is read or consumed.
 #[derive(Debug)]
 pub struct BufPoolWriter<'t> {
     pool: &'t mut BufPool,
@@ -87,32 +110,43 @@ pub struct BufPoolWriter<'t> {
 impl<'t> BufPoolWriter<'t> {
     #[inline(always)]
     pub fn reserve(&mut self, additional: usize) {
-        if self.capacity < additional {
-            self.pool.reserve(additional);
-            self.capacity += additional;
-        }
+        let capacity = self.capacity.checked_add(additional).unwrap();
+        self.pool.reserve(capacity);
+        self.capacity = capacity;
     }
 
+    /// Returns the accumulated buffer.
+    ///
+    /// Reserved header and trailer regions must be initialized before the returned
+    /// buffer is read or consumed.
     #[inline(always)]
     pub fn split(&mut self) -> BytesMut {
         self.pool.split()
     }
 
+    #[inline(always)]
+    pub fn remaining(&self) -> usize {
+        self.capacity.saturating_sub(self.margins.size())
+    }
+
     #[inline(always)]
     pub fn as_slice(&mut self) -> &mut [MaybeUninit<u8>] {
-        unsafe {
-            self.pool
-                .pool
-                .spare_capacity_mut()
-                .get_unchecked_mut(self.margins.header..self.capacity - self.margins.trailer)
+        let remaining = self.remaining();
+        if remaining == 0 {
+            &mut []
+        } else {
+            &mut self.pool.pool.spare_capacity_mut()
+                [self.margins.header..self.margins.header + remaining]
         }
     }
 
     #[inline(always)]
     pub fn commit(&mut self, written: usize) {
-        let len = self.margins.size() + written;
+        let len = self.margins.size().checked_add(written).unwrap();
         assert!(self.capacity >= len);
         self.capacity -= len;
+        // `commit` marks the margins as part of the buffer without initializing them.
+        // They must be initialized before the resulting buffer is consumed.
         unsafe {
             self.pool.pool.advance_mut(len);
         }
@@ -141,8 +175,8 @@ impl<const SIZE: usize> FixedBufPool<SIZE> {
 #[derive(Debug, Deref, DerefMut, AsRef, AsMut)]
 pub struct FixedBufGuard<'p, const SIZE: usize> {
     pool: &'p FixedBufPool<SIZE>,
-    #[deref]
-    #[deref_mut]
+    #[deref(forward)]
+    #[deref_mut(forward)]
     #[as_ref([u8])]
     #[as_mut([u8])]
     buf: Vec<u8>,