feat(s2n-quic-core): add buffer reader/writer traits

quic/s2n-quic-bench/src/buffer.rs

@@ -2,7 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use criterion::{black_box, BenchmarkId, Criterion, Throughput};
-use s2n_quic_core::{buffer::ReceiveBuffer, varint::VarInt};
+use s2n_quic_core::{
+    buffer::{reader::Storage as _, writer, Reassembler},
+    varint::VarInt,
+};
 
 pub fn benchmarks(c: &mut Criterion) {
     let mut group = c.benchmark_group("buffer");
@@ -13,19 +16,21 @@ pub fn benchmarks(c: &mut Criterion) {
         group.throughput(Throughput::Bytes(input.len() as _));
 
         group.bench_with_input(BenchmarkId::new("skip", size), &input, |b, _input| {
-            let mut buffer = ReceiveBuffer::new();
+            let mut buffer = Reassembler::new();
+            let size = VarInt::try_from(size).unwrap();
             b.iter(move || {
                 buffer.skip(black_box(size)).unwrap();
             });
         });
 
         group.bench_with_input(BenchmarkId::new("write_at", size), &input, |b, input| {
-            let mut buffer = ReceiveBuffer::new();
+            let mut buffer = Reassembler::new();
             let mut offset = VarInt::from_u8(0);
             let len = VarInt::new(input.len() as _).unwrap();
             b.iter(move || {
                 buffer.write_at(offset, input).unwrap();
-                buffer.copy_into_buf(&mut NoOpBuf);
+                // Avoid oversampling the `pop` implementation
+                buffer.copy_into(&mut writer::storage::Discard).unwrap();
                 offset += len;
             });
         });
@@ -36,42 +41,19 @@ pub fn benchmarks(c: &mut Criterion) {
             BenchmarkId::new("write_at_fragmented", size),
             &input,
             |b, input| {
-                let mut buffer = ReceiveBuffer::new();
+                let mut buffer = Reassembler::new();
                 let mut offset = VarInt::from_u8(0);
                 let len = VarInt::new(input.len() as _).unwrap();
                 b.iter(move || {
                     let first_offset = offset + len;
                     buffer.write_at(first_offset, input).unwrap();
                     let second_offset = offset;
                     buffer.write_at(second_offset, input).unwrap();
-                    buffer.copy_into_buf(&mut NoOpBuf);
+                    // Avoid oversampling the `pop` implementation
+                    buffer.copy_into(&mut writer::storage::Discard).unwrap();
                     offset = first_offset + len;
                 });
             },
         );
     }
 }
-
-/// A BufMut implementation that doesn't actually copy data into it
-///
-/// This is used to avoid oversampling the `pop` implementation for
-/// `write_at` benchmarks.
-struct NoOpBuf;
-
-unsafe impl bytes::BufMut for NoOpBuf {
-    #[inline]
-    fn remaining_mut(&self) -> usize {
-        usize::MAX
-    }
-
-    #[inline]
-    unsafe fn advance_mut(&mut self, _cnt: usize) {}
-
-    #[inline]
-    fn put_slice(&mut self, _slice: &[u8]) {}
-
-    #[inline]
-    fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
-        unimplemented!()
-    }
-}

quic/s2n-quic-core/src/buffer.rs

@@ -0,0 +1,14 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod duplex;
+mod error;
+pub mod reader;
+pub mod reassembler;
+pub mod writer;
+
+pub use duplex::Duplex;
+pub use error::Error;
+pub use reader::Reader;
+pub use reassembler::Reassembler;
+pub use writer::Writer;

quic/s2n-quic-core/src/buffer/duplex.rs

@@ -0,0 +1,22 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::{Error, Reader, Writer};
+use crate::varint::VarInt;
+
+mod split;
+
+pub use split::Split;
+
+/// A buffer that is capable of both reading and writing
+pub trait Duplex: Reader + Writer {}
+
+impl<T: Reader + Writer> Duplex for T {}
+
+/// A buffer that is capable of both skipping a write and read with a given amount.
+///
+/// This can be used for scenarios where the buffer was written somewhere else but still needed to
+/// be tracked.
+pub trait Skip: Duplex {
+    fn skip(&mut self, len: VarInt, final_offset: Option<VarInt>) -> Result<(), Error>;
+}

quic/s2n-quic-core/src/buffer/duplex/split.rs

@@ -0,0 +1,170 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{
+    buffer::{
+        duplex,
+        reader::{self, Reader, Storage as _},
+        writer::{self, Writer},
+        Error,
+    },
+    varint::VarInt,
+};
+use core::convert::Infallible;
+
+/// A split duplex that tries to write as much as possible to `chunk`, while falling back to
+/// `duplex`.
+pub struct Split<'a, C, D>
+where
+    C: writer::Storage + ?Sized,
+    D: duplex::Skip<Error = Infallible> + ?Sized,
+{
+    chunk: &'a mut C,
+    duplex: &'a mut D,
+}
+
+impl<'a, C, D> Split<'a, C, D>
+where
+    C: writer::Storage + ?Sized,
+    D: duplex::Skip<Error = Infallible> + ?Sized,
+{
+    #[inline]
+    pub fn new(chunk: &'a mut C, duplex: &'a mut D) -> Self {
+        Self { chunk, duplex }
+    }
+}
+
+/// Delegates to the inner Duplex
+impl<'a, C, D> reader::Storage for Split<'a, C, D>
+where
+    C: writer::Storage + ?Sized,
+    D: duplex::Skip<Error = Infallible> + ?Sized,
+{
+    type Error = D::Error;
+
+    #[inline]
+    fn buffered_len(&self) -> usize {
+        self.duplex.buffered_len()
+    }
+
+    #[inline]
+    fn buffer_is_empty(&self) -> bool {
+        self.duplex.buffer_is_empty()
+    }
+
+    #[inline]
+    fn read_chunk(&mut self, watermark: usize) -> Result<reader::storage::Chunk<'_>, Self::Error> {
+        self.duplex.read_chunk(watermark)
+    }
+
+    #[inline]
+    fn partial_copy_into<Dest>(
+        &mut self,
+        dest: &mut Dest,
+    ) -> Result<reader::storage::Chunk<'_>, Self::Error>
+    where
+        Dest: writer::Storage + ?Sized,
+    {
+        self.duplex.partial_copy_into(dest)
+    }
+
+    #[inline]
+    fn copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<(), Self::Error>
+    where
+        Dest: writer::Storage + ?Sized,
+    {
+        self.duplex.copy_into(dest)
+    }
+}
+
+/// Delegates to the inner Duplex
+impl<'a, C, D> Reader for Split<'a, C, D>
+where
+    C: writer::Storage + ?Sized,
+    D: duplex::Skip<Error = Infallible> + ?Sized,
+{
+    #[inline]
+    fn current_offset(&self) -> VarInt {
+        self.duplex.current_offset()
+    }
+
+    #[inline]
+    fn final_offset(&self) -> Option<VarInt> {
+        self.duplex.final_offset()
+    }
+
+    #[inline]
+    fn has_buffered_fin(&self) -> bool {
+        self.duplex.has_buffered_fin()
+    }
+
+    #[inline]
+    fn is_consumed(&self) -> bool {
+        self.duplex.is_consumed()
+    }
+}
+
+impl<'a, C, D> Writer for Split<'a, C, D>
+where
+    C: writer::Storage + ?Sized,
+    D: duplex::Skip<Error = Infallible> + ?Sized,
+{
+    #[inline]
+    fn copy_from<R>(&mut self, reader: &mut R) -> Result<(), Error<R::Error>>
+    where
+        R: Reader + ?Sized,
+    {
+        // enable reader checks
+        let mut reader = reader.with_checks();
+        let reader = &mut reader;
+
+        let final_offset = reader.final_offset();
+
+        {
+            // if the chunk specializes writing zero-copy Bytes/BytesMut, then just write to the
+            // receive buffer, since that's what it stores
+            let mut should_delegate = C::SPECIALIZES_BYTES || C::SPECIALIZES_BYTES_MUT;
+
+            // if this packet is non-contiguous, then delegate to the wrapped writer
+            should_delegate |= reader.current_offset() != self.duplex.current_offset();
+
+            // if the chunk has less than half of the payload, then delegate
+            should_delegate |= self.chunk.remaining_capacity() < (reader.buffered_len() / 2);
+
+            if should_delegate {
+                self.duplex.copy_from(reader)?;
+
+                // don't copy into chunk here - let the caller do that later if they'd like
+
+                return Ok(());
+            }
+        }
+
+        debug_assert!(
+            self.chunk.has_remaining_capacity(),
+            "this code should only be executed if the chunk is empty"
+        );
+
+        {
+            // track the number of consumed bytes
+            let mut reader = reader.tracked();
+
+            reader.copy_into(self.chunk)?;
+
+            let write_len = reader.consumed_len();
+            let write_len = VarInt::try_from(write_len).map_err(|_| Error::OutOfRange)?;
+
+            // notify the duplex that we bypassed it and should skip
+            self.duplex
+                .skip(write_len, final_offset)
+                .map_err(Error::mapped)?;
+        }
+
+        // if we still have some remaining bytes consume the rest in the duplex
+        if !reader.buffer_is_empty() {
+            self.duplex.copy_from(reader)?;
+        }
+
+        Ok(())
+    }
+}

quic/s2n-quic-core/src/buffer/error.rs

@@ -0,0 +1,59 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum Error<Reader = core::convert::Infallible> {
+    /// An invalid data range was provided
+    OutOfRange,
+    /// The provided final size was invalid for the buffer's state
+    InvalidFin,
+    /// The provided reader failed
+    ReaderError(Reader),
+}
+
+impl<Reader> From<Reader> for Error<Reader> {
+    #[inline]
+    fn from(reader: Reader) -> Self {
+        Self::ReaderError(reader)
+    }
+}
+
+impl Error {
+    #[inline]
+    pub fn mapped<Reader>(error: Error) -> Error<Reader> {
+        match error {
+            Error::OutOfRange => Error::OutOfRange,
+            Error::InvalidFin => Error::InvalidFin,
+            Error::ReaderError(_) => unreachable!(),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl<Reader: std::error::Error> std::error::Error for Error<Reader> {}
+
+impl<Reader: core::fmt::Display> core::fmt::Display for Error<Reader> {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        match self {
+            Self::OutOfRange => write!(f, "write extends out of the maximum possible offset"),
+            Self::InvalidFin => write!(
+                f,
+                "write modifies the final offset in a non-compliant manner"
+            ),
+            Self::ReaderError(reader) => write!(f, "the provided reader failed with: {reader}"),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl<Reader: 'static + std::error::Error + Send + Sync> From<Error<Reader>> for std::io::Error {
+    #[inline]
+    fn from(error: Error<Reader>) -> Self {
+        let kind = match &error {
+            Error::OutOfRange => std::io::ErrorKind::InvalidData,
+            Error::InvalidFin => std::io::ErrorKind::InvalidData,
+            Error::ReaderError(_) => std::io::ErrorKind::Other,
+        };
+        Self::new(kind, error)
+    }
+}