feat(s2n-quic-core): use scatter buffers for encryption

common/s2n-codec/src/encoder/buffer.rs

@@ -23,7 +23,11 @@ impl<'a> EncoderBuffer<'a> {
     /// Panics when `position > capacity`
     #[inline]
     pub fn set_position(&mut self, position: usize) {
-        debug_assert!(position <= self.capacity());
+        debug_assert!(
+            position <= self.capacity(),
+            "position {position} exceeded capacity of {}",
+            self.capacity()
+        );
         self.position = position;
     }
 
@@ -56,7 +60,7 @@ impl<'a> EncoderBuffer<'a> {
     }
 
     #[inline]
-    fn assert_capacity(&self, len: usize) {
+    pub(crate) fn assert_capacity(&self, len: usize) {
         debug_assert!(
             len <= self.remaining_capacity(),
             "not enough buffer capacity. wanted: {}, available: {}",

common/s2n-codec/src/encoder/mod.rs

@@ -3,6 +3,7 @@
 
 pub mod buffer;
 pub mod estimator;
+pub mod scatter;
 pub mod value;
 
 pub use buffer::*;

common/s2n-codec/src/encoder/scatter.rs

@@ -0,0 +1,165 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{Encoder, EncoderBuffer};
+
+pub struct Buffer<'a> {
+    inner: EncoderBuffer<'a>,
+    #[cfg(feature = "bytes")]
+    extra: Option<bytes::Bytes>,
+}
+
+impl<'a> Buffer<'a> {
+    /// Ensures an extra bytes are written into the main EncoderBuffer.
+    #[inline]
+    pub fn flatten(&mut self) -> &mut EncoderBuffer<'a> {
+        self.flush();
+        &mut self.inner
+    }
+
+    /// Resets the encoder to its initial state while also dropping any extra bytes at the end
+    #[inline]
+    pub fn clear(&mut self) {
+        self.inner.set_position(0);
+        self.extra = None;
+    }
+}
+
+/// Implement a version with `bytes` enabled
+#[cfg(feature = "bytes")]
+impl<'a> Buffer<'a> {
+    /// Initializes a buffer without any extra bytes at the end
+    #[inline]
+    pub fn new(inner: EncoderBuffer<'a>) -> Self {
+        Self { inner, extra: None }
+    }
+
+    /// Initializes the buffer with extra bytes
+    ///
+    /// NOTE the EncoderBuffer position should not include the extra bytes. This ensures the bytes
+    /// can be "flushed" into the EncoderBuffer if another write happens or `flatten` is called.
+    #[inline]
+    pub fn new_with_extra(inner: EncoderBuffer<'a>, extra: Option<bytes::Bytes>) -> Self {
+        Self { inner, extra }
+    }
+
+    /// Converts the buffer into its inner parts
+    ///
+    /// NOTE: the EncoderBuffer position will not include the extra bytes. The caller will need to
+    /// account for this data.
+    #[inline]
+    pub fn into_inner(self) -> (EncoderBuffer<'a>, Option<bytes::Bytes>) {
+        (self.inner, self.extra)
+    }
+
+    /// Returns the inner EncoderBuffer, along with the current extra bytes
+    #[inline]
+    pub fn inner_mut(&mut self) -> (&mut EncoderBuffer<'a>, &Option<bytes::Bytes>) {
+        (&mut self.inner, &self.extra)
+    }
+
+    #[inline]
+    fn flush(&mut self) {
+        // move the extra bytes into the main buffer
+        if let Some(extra) = self.extra.take() {
+            self.inner.write_slice(&extra);
+        }
+    }
+}
+
+/// Include a functional implementation when `bytes` is not available
+#[cfg(not(feature = "bytes"))]
+impl<'a> Buffer<'a> {
+    #[inline]
+    pub fn new(inner: EncoderBuffer<'a>) -> Self {
+        Self { inner }
+    }
+
+    #[inline]
+    pub fn new_with_extra(inner: EncoderBuffer<'a>, extra: Option<&'static [u8]>) -> Self {
+        debug_assert!(extra.is_none());
+        Self { inner }
+    }
+
+    #[inline]
+    pub fn into_inner(self) -> (EncoderBuffer<'a>, Option<&'static [u8]>) {
+        (self.inner, None)
+    }
+
+    #[inline]
+    pub fn inner_mut(&mut self) -> (&mut EncoderBuffer<'a>, &Option<&'static [u8]>) {
+        (&mut self.inner, &None)
+    }
+
+    #[inline(always)]
+    fn flush(&mut self) {}
+}
+
+impl<'a> Encoder for Buffer<'a> {
+    /// We have special handling for writes that include `Bytes` so signal that
+    #[cfg(feature = "bytes")]
+    const SPECIALIZES_BYTES: bool = true;
+
+    #[inline]
+    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, write: F) {
+        // we need to flush the extra bytes if we have them
+        self.flush();
+        self.inner.write_sized(len, write)
+    }
+
+    #[inline]
+    fn write_slice(&mut self, slice: &[u8]) {
+        // we need to flush the extra bytes if we have them
+        self.flush();
+        self.inner.write_slice(slice);
+    }
+
+    #[inline]
+    #[cfg(feature = "bytes")]
+    fn write_bytes(&mut self, bytes: bytes::Bytes) {
+        // we need to flush the extra bytes if we have them and replace them with the new one
+        self.flush();
+        // ensure the underlying buffer is big enough for the write
+        self.inner.assert_capacity(bytes.len());
+        // store the extra bytes and defer the copy
+        self.extra = Some(bytes);
+    }
+
+    #[inline]
+    fn write_zerocopy<
+        T: zerocopy::AsBytes + zerocopy::FromBytes + zerocopy::Unaligned,
+        F: FnOnce(&mut T),
+    >(
+        &mut self,
+        write: F,
+    ) {
+        // we need to flush the extra bytes if we have them
+        self.flush();
+        self.inner.write_zerocopy(write);
+    }
+
+    #[inline]
+    fn write_repeated(&mut self, count: usize, value: u8) {
+        // we need to flush the extra bytes if we have them
+        self.flush();
+        self.inner.write_repeated(count, value)
+    }
+
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.inner.capacity()
+    }
+
+    #[inline]
+    fn len(&self) -> usize {
+        let mut len = self.inner.len();
+
+        // make sure our len includes the extra bytes if we have them
+        #[cfg(feature = "bytes")]
+        if let Some(extra) = self.extra.as_ref() {
+            len += extra.len();
+        }
+
+        len
+    }
+}

quic/s2n-quic-bench/src/crypto/aesgcm.rs

@@ -3,6 +3,7 @@
 
 use core::convert::TryInto;
 use criterion::{black_box, BenchmarkId, Criterion, Throughput};
+use s2n_codec::{encoder::scatter, EncoderBuffer};
 use s2n_quic_crypto::testing::{
     self,
     aesgcm::{Key, NONCE_LEN, TAG_LEN},
@@ -52,10 +53,11 @@ pub fn benchmarks(c: &mut Criterion) {
                         let payload_len = input.len();
                         input.extend_from_slice(&[0; TAG_LEN]);
 
-                        let (payload, tag) = input.split_at_mut(payload_len);
-                        let tag: &mut [u8; TAG_LEN] = tag.try_into().unwrap();
                         b.iter(|| {
-                            let _ = key.encrypt(&nonce, &aad, payload, tag);
+                            let mut payload = EncoderBuffer::new(&mut input);
+                            payload.advance_position(payload_len);
+                            let mut payload = scatter::Buffer::new(payload);
+                            let _ = key.encrypt(&nonce, &aad, &mut payload);
                         });
                     },
                 );
@@ -74,11 +76,16 @@ pub fn benchmarks(c: &mut Criterion) {
                     let payload_len = input.len();
                     input.extend_from_slice(&[0; TAG_LEN]);
 
+                    {
+                        let mut payload = EncoderBuffer::new(&mut input);
+                        payload.advance_position(payload_len);
+                        let mut payload = scatter::Buffer::new(payload);
+                        // create a valid encrypted payload
+                        key.encrypt(&nonce, &aad, &mut payload).unwrap();
+                    }
+
                     let (payload, tag) = input.split_at_mut(payload_len);
                     let tag: &mut [u8; TAG_LEN] = tag.try_into().unwrap();
-
-                    // create a valid encrypted payload
-                    key.encrypt(&nonce, &aad, payload, tag).unwrap();
                     let tag = &*tag;
 
                     b.iter_batched(

quic/s2n-quic-core/src/crypto/key.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::crypto::CryptoError;
+use s2n_codec::encoder::scatter;
 
 /// A trait for crypto keys
 pub trait Key: Send {
@@ -18,7 +19,7 @@ pub trait Key: Send {
         &self,
         packet_number: u64,
         header: &[u8],
-        payload: &mut [u8],
+        payload: &mut scatter::Buffer,
     ) -> Result<(), CryptoError>;
 
     /// Length of the appended tag
@@ -37,7 +38,7 @@ pub trait Key: Send {
 pub mod testing {
     use crate::crypto::{
         retry::{IntegrityTag, INTEGRITY_TAG_LEN},
-        CryptoError, HandshakeHeaderKey, HandshakeKey, HeaderKey as CryptoHeaderKey,
+        scatter, CryptoError, HandshakeHeaderKey, HandshakeKey, HeaderKey as CryptoHeaderKey,
         HeaderProtectionMask, InitialHeaderKey, InitialKey, OneRttHeaderKey, OneRttKey, RetryKey,
         ZeroRttHeaderKey, ZeroRttKey,
     };
@@ -89,8 +90,10 @@ pub mod testing {
             &self,
             _packet_number: u64,
             _header: &[u8],
-            _payload: &mut [u8],
+            payload: &mut scatter::Buffer,
         ) -> Result<(), CryptoError> {
+            // copy any bytes into the final slice
+            payload.flatten();
             Ok(())
         }
 

quic/s2n-quic-core/src/crypto/mod.rs

@@ -176,6 +176,7 @@ pub trait CryptoSuite {
 use crate::packet::number::{
     PacketNumber, PacketNumberLen, PacketNumberSpace, TruncatedPacketNumber,
 };
+pub use s2n_codec::encoder::scatter;
 use s2n_codec::{DecoderBufferMut, DecoderError, Encoder, EncoderBuffer};
 
 /// Protects an `EncryptedPayload` into a `ProtectedPayload`
@@ -211,12 +212,21 @@ pub fn encrypt<'a, K: Key>(
     packet_number: PacketNumber,
     packet_number_len: PacketNumberLen,
     header_len: usize,
-    mut payload: EncoderBuffer<'a>,
+    payload: scatter::Buffer<'a>,
 ) -> Result<(EncryptedPayload<'a>, EncoderBuffer<'a>), CryptoError> {
     let header_with_pn_len = packet_number_len.bytesize() + header_len;
 
-    // Make space for the key tag
-    payload.write_repeated(key.tag_len(), 0);
+    let (mut payload, extra) = payload.into_inner();
+
+    let inline_len = payload.len() - header_with_pn_len;
+
+    // reserve bytes for the extra chunk at the end
+    if let Some(extra) = extra.as_ref() {
+        payload.advance_position(extra.len());
+    }
+
+    // reserve bytes for the tag
+    payload.advance_position(key.tag_len());
 
     let (payload, remaining) = payload.split_off();
 
@@ -227,7 +237,10 @@ pub fn encrypt<'a, K: Key>(
         payload.len()
     );
     let (header, body) = payload.split_at_mut(header_with_pn_len);
-    key.encrypt(packet_number.as_crypto_nonce(), header, body)?;
+    let mut body = EncoderBuffer::new(body);
+    body.advance_position(inline_len);
+    let mut body = scatter::Buffer::new_with_extra(body, extra);
+    key.encrypt(packet_number.as_crypto_nonce(), header, &mut body)?;
 
     let encrypted_payload = EncryptedPayload::new(header_len, packet_number_len, payload);
     let remaining = EncoderBuffer::new(remaining);

quic/s2n-quic-core/src/crypto/tests.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    crypto::{CryptoError, HeaderKey, HeaderProtectionMask, Key, ProtectedPayload},
+    crypto::{scatter, CryptoError, HeaderKey, HeaderProtectionMask, Key, ProtectedPayload},
     packet::number::{PacketNumber, PacketNumberSpace},
     varint::VarInt,
 };
@@ -80,6 +80,7 @@ fn fuzz_protect(
     let payload_len = input.len();
     let mut payload = EncoderBuffer::new(input);
     payload.set_position(payload_len);
+    let payload = scatter::Buffer::new(payload);
 
     let truncated_packet_number = packet_number.truncate(largest_packet_number).unwrap();
     let packet_number_len = truncated_packet_number.len();
@@ -119,12 +120,16 @@ impl Key for FuzzCrypto {
         &self,
         packet_number: u64,
         _header: &[u8],
-        payload: &mut [u8],
+        payload: &mut scatter::Buffer,
     ) -> Result<(), CryptoError> {
+        let payload = payload.flatten();
+        let (payload, _) = payload.split_mut();
+
         let mask = packet_number as u8;
         for byte in payload.iter_mut() {
             *byte ^= mask;
         }
+
         Ok(())
     }
 

quic/s2n-quic-core/src/crypto/tls/null.rs

@@ -291,6 +291,7 @@ pub mod server {
 
 mod key {
     use super::*;
+    use crate::crypto::scatter;
 
     #[derive(Debug)]
     pub struct NoCrypto;
@@ -312,9 +313,10 @@ mod key {
             &self,
             _packet_number: u64,
             _header: &[u8],
-            _payload: &mut [u8],
+            payload: &mut scatter::Buffer,
         ) -> Result<(), crypto::CryptoError> {
-            // Do nothing
+            // copy any extra bytes into the slice
+            payload.flatten();
             Ok(())
         }