Use DatagramBuffer directly, remove BufSlice trait

This uses the DatagramBuffer type directly, in favour of having the
BufSlice trait that is only used once.  This meanse Header::encode now
takes &mut DatagramBuffer and all the callers have to construct this.
Previously they used `&mut Vec<u8>` for this.

The DatagramBuffer type is still kept crate-internal.
Endpoint::handle and friends still use `&mut Vec<u8>` for the buffer
in which to build the response packet, to not change the public API.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -89,7 +89,7 @@ mod timer;
 use timer::{Timer, TimerTable};
 
 mod transmit_builder;
-pub(crate) use transmit_builder::BufSlice;
+pub(crate) use transmit_builder::DatagramBuffer;
 use transmit_builder::TransmitBuilder;
 
 /// Protocol state and logic for a single QUIC connection

quinn-proto/src/connection/packet_builder.rs

@@ -5,7 +5,7 @@ use tracing::{trace, trace_span};
 use super::{Connection, SentFrames, TransmitBuilder, spaces::SentPacket};
 use crate::{
     ConnectionId, Instant, TransportError, TransportErrorCode,
-    connection::{BufSlice, ConnectionSide},
+    connection::ConnectionSide,
     frame::{self, Close},
     packet::{FIXED_BIT, Header, InitialHeader, LongType, PacketNumber, PartialEncode, SpaceId},
 };

quinn-proto/src/connection/transmit_builder.rs

@@ -149,8 +149,6 @@ impl<'a> TransmitBuilder<'a> {
     /// This buffer implements [`BufSlice`] and thus allows writing into the buffer using
     /// [`BufMut`] and both reading and modifying the already written data in the buffer
     /// using [`Deref`] and [`DerefMut`].  The buffer also enforces a maximum size.
-    // This returns a concrete type since `impl BufSlice` would need a `+ use<'_>` lifetime
-    // bound, which is only supported since Rust 1.82.
     pub(super) fn datagram_mut(&mut self) -> DatagramBuffer<'_> {
         DatagramBuffer::new(
             self.buf,
@@ -196,32 +194,6 @@ impl<'a> TransmitBuilder<'a> {
     }
 }
 
-/// A writable buffer with a capacity and access to the already written bytes
-///
-/// This is a writable buffer, using the [`BufMut`] trait, which also has a maximum capacity.
-/// If more bytes are written than the buffer can contain a panic will occur.
-///
-/// The [`Deref`] impl must return the already written bytes as a slice.  [`DerefMut`] must
-/// allow modification of these already written bytes.
-pub(crate) trait BufSlice: BufMut + Deref<Target = [u8]> + DerefMut {
-    /// Returns the number of already written bytes in the buffer
-    fn len(&self) -> usize;
-
-    /// Returns the maximum size of the buffer
-    fn capacity(&self) -> usize;
-}
-
-impl BufSlice for Vec<u8> {
-    fn len(&self) -> usize {
-        self.len()
-    }
-
-    /// For `Vec<u8>` the capacity is the same as the maximum vec size
-    fn capacity(&self) -> usize {
-        isize::MAX as usize
-    }
-}
-
 /// A [`BufSlice`] implementation for a datagram
 #[derive(Debug)]
 pub(crate) struct DatagramBuffer<'a> {
@@ -235,6 +207,10 @@ pub(crate) struct DatagramBuffer<'a> {
 
 impl<'a> DatagramBuffer<'a> {
     pub(crate) fn new(buf: &'a mut Vec<u8>, start_offset: usize, max_size: usize) -> Self {
+        // Make sure that at least this datagram is allocated.  Does nothing if, like for a
+        // transmit, already more has been allocated.
+        buf.reserve(max_size);
+
         let max_offset = start_offset + max_size;
         DatagramBuffer {
             buf,
@@ -272,21 +248,16 @@ impl DerefMut for DatagramBuffer<'_> {
     }
 }
 
-impl BufSlice for DatagramBuffer<'_> {
-    /// Returns the number of already written bytes in the buffer
-    fn len(&self) -> usize {
-        self.buf.len() - self.start_offset
-    }
-
+impl DatagramBuffer<'_> {
     /// Returns the maximum size of the buffer
-    fn capacity(&self) -> usize {
+    pub(crate) fn capacity(&self) -> usize {
         self.max_offset - self.start_offset
     }
 }
 
 // Temporary compatibility with the BufLen trait.  To be removed in follow-up commits.
 impl BufLen for DatagramBuffer<'_> {
     fn len(&self) -> usize {
-        <Self as BufSlice>::len(self)
+        self.deref().len()
     }
 }

quinn-proto/src/endpoint.rs

@@ -20,7 +20,7 @@ use crate::{
     cid_generator::ConnectionIdGenerator,
     coding::BufMutExt,
     config::{ClientConfig, EndpointConfig, ServerConfig},
-    connection::{Connection, ConnectionError, SideArgs},
+    connection::{Connection, ConnectionError, DatagramBuffer, SideArgs},
     crypto::{self, Keys, UnsupportedVersion},
     frame,
     packet::{
@@ -146,6 +146,8 @@ impl Endpoint {
         data: BytesMut,
         buf: &mut Vec<u8>,
     ) -> Option<DatagramEvent> {
+        let start_pos = buf.len();
+        let buf = &mut DatagramBuffer::new(buf, start_pos, MIN_INITIAL_SIZE.into());
         // Partially decode packet or short-circuit if unable
         let datagram_len = data.len();
         let event = match PartialDecode::new(
@@ -264,7 +266,7 @@ impl Endpoint {
         inciting_dgram_len: usize,
         addresses: FourTuple,
         dst_cid: ConnectionId,
-        buf: &mut Vec<u8>,
+        buf: &mut DatagramBuffer<'_>,
     ) -> Option<Transmit> {
         if self
             .last_stateless_reset
@@ -303,11 +305,10 @@ impl Endpoint {
             self.rng
                 .random_range(IDEAL_MIN_PADDING_LEN..max_padding_len)
         };
-        buf.reserve(padding_len + RESET_TOKEN_SIZE);
-        buf.resize(padding_len, 0);
+        buf.put_bytes(0, padding_len);
         self.rng.fill_bytes(&mut buf[0..padding_len]);
         buf[0] = 0b0100_0000 | (buf[0] >> 2);
-        buf.extend_from_slice(&ResetToken::new(&*self.config.reset_key, dst_cid));
+        buf.put_slice(&ResetToken::new(&*self.config.reset_key, dst_cid));
 
         debug_assert!(buf.len() < inciting_dgram_len);
 
@@ -419,7 +420,7 @@ impl Endpoint {
         datagram_len: usize,
         event: DatagramConnectionEvent,
         addresses: FourTuple,
-        buf: &mut Vec<u8>,
+        buf: &mut DatagramBuffer<'_>,
     ) -> Option<DatagramEvent> {
         let dst_cid = event.first_decode.dst_cid();
         let header = event.first_decode.initial_header().unwrap();
@@ -525,6 +526,8 @@ impl Endpoint {
         buf: &mut Vec<u8>,
         server_config: Option<Arc<ServerConfig>>,
     ) -> Result<(ConnectionHandle, Connection), AcceptError> {
+        let buf_start = buf.len();
+        let mut buf = DatagramBuffer::new(buf, buf_start, MIN_INITIAL_SIZE.into());
         let remote_address_validated = incoming.remote_address_validated();
         incoming.improper_drop_warner.dismiss();
         let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
@@ -566,7 +569,7 @@ impl Endpoint {
                     &incoming.crypto,
                     &src_cid,
                     TransportError::CONNECTION_REFUSED(""),
-                    buf,
+                    &mut buf,
                 )),
             });
         }
@@ -664,7 +667,7 @@ impl Endpoint {
                         &incoming.crypto,
                         &src_cid,
                         e.clone(),
-                        buf,
+                        &mut buf,
                     )),
                     _ => None,
                 };
@@ -708,13 +711,15 @@ impl Endpoint {
         self.clean_up_incoming(&incoming);
         incoming.improper_drop_warner.dismiss();
 
+        let start_pos = buf.len();
+        let mut buf = DatagramBuffer::new(buf, start_pos, MIN_INITIAL_SIZE.into());
         self.initial_close(
             incoming.packet.header.version,
             incoming.addresses,
             &incoming.crypto,
             &incoming.packet.header.src_cid,
             TransportError::CONNECTION_REFUSED(""),
-            buf,
+            &mut buf,
         )
     }
 
@@ -752,14 +757,16 @@ impl Endpoint {
             version: incoming.packet.header.version,
         };
 
-        let encode = header.encode(buf);
+        let start_pos = buf.len();
+        let mut buf = DatagramBuffer::new(buf, start_pos, MIN_INITIAL_SIZE.into());
+        let encode = header.encode(&mut buf);
         buf.put_slice(&token);
-        buf.extend_from_slice(&server_config.crypto.retry_tag(
+        buf.put_slice(&server_config.crypto.retry_tag(
             incoming.packet.header.version,
             &incoming.packet.header.dst_cid,
-            buf,
+            &buf,
         ));
-        encode.finish(buf, &*incoming.crypto.header.local, None);
+        encode.finish(&mut buf, &*incoming.crypto.header.local, None);
 
         Ok(Transmit {
             destination: incoming.addresses.remote,
@@ -854,7 +861,7 @@ impl Endpoint {
         crypto: &Keys,
         remote_id: &ConnectionId,
         reason: TransportError,
-        buf: &mut Vec<u8>,
+        buf: &mut DatagramBuffer<'_>,
     ) -> Transmit {
         // We don't need to worry about CID collisions in initial closes because the peer
         // shouldn't respond, and if it does, and the CID collides, we'll just drop the
@@ -873,7 +880,7 @@ impl Endpoint {
         let max_len =
             INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
         frame::Close::from(reason).encode(buf, max_len);
-        buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
+        buf.put_bytes(0, crypto.packet.local.tag_len());
         partial_encode.finish(buf, &*crypto.header.local, Some((0, &*crypto.packet.local)));
         Transmit {
             destination: addresses.remote,

quinn-proto/src/packet.rs

@@ -6,7 +6,7 @@ use thiserror::Error;
 use crate::{
     ConnectionId,
     coding::{self, BufExt, BufMutExt},
-    connection::BufSlice,
+    connection::DatagramBuffer,
     crypto,
 };
 
@@ -286,7 +286,7 @@ impl Header {
     ///
     /// The current position of the buffer is stored in the [`PartialEncode`] as the start
     /// of the packet in the buffer.
-    pub(crate) fn encode(&self, w: &mut impl BufSlice) -> PartialEncode {
+    pub(crate) fn encode(&self, w: &mut DatagramBuffer<'_>) -> PartialEncode {
         use Header::*;
         let start = w.len();
         match *self {
@@ -943,8 +943,10 @@ mod tests {
     #[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
     #[test]
     fn header_encoding() {
-        use crate::Side;
+        use std::ops::Deref;
+
         use crate::crypto::rustls::{initial_keys, initial_suite_from_provider};
+        use crate::{MIN_INITIAL_SIZE, Side};
         #[cfg(all(feature = "rustls-aws-lc-rs", not(feature = "rustls-ring")))]
         use rustls::crypto::aws_lc_rs::default_provider;
         #[cfg(feature = "rustls-ring")]
@@ -957,6 +959,7 @@ mod tests {
         let suite = initial_suite_from_provider(&std::sync::Arc::new(provider)).unwrap();
         let client = initial_keys(Version::V1, dcid, Side::Client, &suite);
         let mut buf = Vec::new();
+        let mut buf = DatagramBuffer::new(&mut buf, 0, MIN_INITIAL_SIZE.into());
         let header = Header::Initial(InitialHeader {
             number: PacketNumber::U8(0),
             src_cid: ConnectionId::new(&[]),
@@ -965,15 +968,14 @@ mod tests {
             version: crate::DEFAULT_SUPPORTED_VERSIONS[0],
         });
         let encode = header.encode(&mut buf);
-        let header_len = buf.len();
-        buf.resize(header_len + 16 + client.packet.local.tag_len(), 0);
+        buf.put_bytes(0, 16 + client.packet.local.tag_len());
         encode.finish(
             &mut buf,
             &*client.header.local,
             Some((0, &*client.packet.local)),
         );
 
-        for byte in &buf {
+        for byte in buf.deref() {
             print!("{byte:02x}");
         }
         println!();
@@ -988,7 +990,7 @@ mod tests {
         let server = initial_keys(Version::V1, dcid, Side::Server, &suite);
         let supported_versions = crate::DEFAULT_SUPPORTED_VERSIONS.to_vec();
         let decode = PartialDecode::new(
-            buf.as_slice().into(),
+            buf.deref().into(),
             &FixedLengthConnectionIdParser::new(0),
             &supported_versions,
             false,