From eb32d2a930f1cdd09ff8627fa134673861ffb73c Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:58:46 +0200 Subject: [PATCH] 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 | 102 ++++++++++++++++++++------------ easytier/src/utils/buf.rs | 62 ++++++++++++++----- 2 files changed, 113 insertions(+), 51 deletions(-) diff --git a/easytier-core/src/tunnel/buf.rs b/easytier-core/src/tunnel/buf.rs index 689d4ff6..f6b3b85e 100644 --- a/easytier-core/src/tunnel/buf.rs +++ b/easytier-core/src/tunnel/buf.rs @@ -15,23 +15,24 @@ impl BufList { } } - #[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 { 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 BufList { impl Extend for BufList { fn extend>(&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 Buf for BufList { #[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 Buf for BufList { #[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::() < 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 Buf for BufList { 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"); + } } diff --git a/easytier/src/utils/buf.rs b/easytier/src/utils/buf.rs index 31173141..5bc32902 100644 --- a/easytier/src/utils/buf.rs +++ b/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] { - 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 FixedBufPool { #[derive(Debug, Deref, DerefMut, AsRef, AsMut)] pub struct FixedBufGuard<'p, const SIZE: usize> { pool: &'p FixedBufPool, - #[deref] - #[deref_mut] + #[deref(forward)] + #[deref_mut(forward)] #[as_ref([u8])] #[as_mut([u8])] buf: Vec,