Only allow access to the current datagram

flub · flub · commit 4b0dcbf70e74 · 2025-05-02T13:16:46.000+02:00
The TransmitBuilder now only allows access to the current datagram
instead of to the entire buffer.

This buffer is implemented as the DatagramBuffer struct.  It is
defined by the BufSlice trait however: code paths outside of
poll_transmit also encode QUIC packets into buffers and they use a
simple Vec buffer.  This trait could potentially later be useful for
other buffers as well, e.g. a packet buffer.

Most importantly the logic of the minimum and maximum datagra size
in the PacketBuilder had to be adjusted to track offsets relative to
the datagram instead of relative to the entire buffer.
diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
@@ -89,6 +89,7 @@ mod timer;
 use timer::{Timer, TimerTable};
 
 mod transmit_builder;
+pub(crate) use transmit_builder::BufSlice;
 use transmit_builder::TransmitBuilder;
 
 /// Protocol state and logic for a single QUIC connection
@@ -537,13 +538,13 @@ impl Connection {
                 ack_eliciting |= self.can_send_1rtt(frame_space_1rtt);
             }
 
-            // Can we append more data into the current buffer?
-            // It is not safe to assume that `buf.len()` is the end of the data,
-            // since the last packet might not have been finished.
+            // Can we append more data into the current datagram?
+            // It is not safe to assume that `transmits.datagram().len()` is the end of the
+            // data, since the last packet might not have been finished.
             let buf_end = if let Some(builder) = &builder_storage {
-                transmits.len().max(builder.min_size) + builder.tag_len
+                transmits.datagram().len().max(builder.min_size) + builder.tag_len
             } else {
-                transmits.len()
+                transmits.datagram().len()
             };
 
             let tag_len = if let Some(ref crypto) = self.spaces[space_id].crypto {
@@ -559,7 +560,9 @@ impl Connection {
             // We are NOT coalescing (the default is we are, so this was turned off in an
             // earlier iteration) OR there is not enough space for another *packet* in this
             // datagram (buf_capacity - buf_end == unused space in datagram).
-            if !coalesce || transmits.datagram_max_offset() - buf_end < MIN_PACKET_SPACE + tag_len {
+            if !coalesce
+                || transmits.datagram_mut().capacity() - buf_end < MIN_PACKET_SPACE + tag_len
+            {
                 // We need to send 1 more datagram and extend the buffer for that.
 
                 // Is 1 more datagram allowed?
@@ -586,7 +589,7 @@ impl Connection {
                 if ack_eliciting && self.spaces[space_id].loss_probes == 0 {
                     // Assume the current packet will get padded to fill the segment
                     let untracked_bytes = if let Some(builder) = &builder_storage {
-                        transmits.datagram_max_offset() - builder.partial_encode.start
+                        transmits.datagram().len() - builder.partial_encode.start
                     } else {
                         0
                     } as u64;
@@ -640,9 +643,10 @@ impl Connection {
                         // MTU. Loss probes are the only packets for which we might grow
                         // `buf_capacity` by less than `segment_size`.
                         const MAX_PADDING: usize = 16;
-                        let packet_len_unpadded = cmp::max(builder.min_size, transmits.len())
-                            - transmits.datagram_start_offset()
-                            + builder.tag_len;
+                        let packet_len_unpadded =
+                            cmp::max(builder.min_size, transmits.datagram().len())
+                                - builder.partial_encode.start
+                                + builder.tag_len;
                         if packet_len_unpadded + MAX_PADDING < transmits.segment_size()
                             || transmits.datagram_start_offset() + transmits.segment_size()
                                 > transmits.datagram_max_offset()
@@ -767,11 +771,14 @@ impl Connection {
                 // to encode the ConnectionClose frame too. However we still have the
                 // check here to prevent crashes if something changes.
                 debug_assert!(
-                    transmits.len() + frame::ConnectionClose::SIZE_BOUND < builder.max_size,
-                    "ACKs should leave space for ConnectionClose"
+                    transmits.datagram().len() + frame::ConnectionClose::SIZE_BOUND
+                        < builder.max_size,
+                    "ACKS should leave space for ConnectionClose"
                 );
-                if transmits.len() + frame::ConnectionClose::SIZE_BOUND < builder.max_size {
-                    let max_frame_size = builder.max_size - transmits.len();
+                if transmits.datagram().len() + frame::ConnectionClose::SIZE_BOUND
+                    < builder.max_size
+                {
+                    let max_frame_size = builder.max_size - transmits.datagram().len();
                     match self.state {
                         State::Closed(state::Closed { ref reason }) => {
                             if space_id == SpaceId::Data || reason.is_transport_layer() {
diff --git a/quinn-proto/src/connection/packet_builder.rs b/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::ConnectionSide,
+    connection::{BufSlice, ConnectionSide},
     frame::{self, Close},
     packet::{FIXED_BIT, Header, InitialHeader, LongType, PacketNumber, PartialEncode, SpaceId},
 };
@@ -17,11 +17,12 @@ pub(super) struct PacketBuilder {
     pub(super) ack_eliciting: bool,
     pub(super) exact_number: u64,
     pub(super) short_header: bool,
-    /// Smallest absolute position in the associated buffer that must be occupied by this packet's
-    /// frames
+    /// The smallest datagram offset that must be occupied by this packet's frames
+    ///
+    /// This is the smallest offset into the datagram this packet is being written into,
+    /// that must contain frames for this packet.
     pub(super) min_size: usize,
-    /// Largest absolute position in the associated buffer that may be occupied by this packet's
-    /// frames
+    /// The largest datagram offset that may be occupied by this packet's frames
     pub(super) max_size: usize,
     pub(super) tag_len: usize,
     pub(super) _span: tracing::span::EnteredSpan,
@@ -122,7 +123,7 @@ impl PacketBuilder {
         };
         let partial_encode = header.encode(&mut transmits.datagram_mut());
         if conn.peer_params.grease_quic_bit && conn.rng.random() {
-            transmits.as_mut_slice()[partial_encode.start] ^= FIXED_BIT;
+            transmits.datagram_mut()[partial_encode.start] ^= FIXED_BIT;
         }
 
         let (sample_size, tag_len) = if let Some(ref crypto) = space.crypto {
@@ -146,10 +147,10 @@ impl PacketBuilder {
         // pn_len + payload_len + tag_len >= sample_size + 4
         // payload_len >= sample_size + 4 - pn_len - tag_len
         let min_size = Ord::max(
-            transmits.len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
+            transmits.datagram().len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
             partial_encode.start + dst_cid.len() + 6,
         );
-        let max_size = transmits.datagram_max_offset() - tag_len;
+        let max_size = transmits.datagram_mut().capacity() - tag_len;
         debug_assert!(max_size >= min_size);
 
         Some(Self {
@@ -172,10 +173,7 @@ impl PacketBuilder {
         // The datagram might already have a larger minimum size than the caller is requesting, if
         // e.g. we're coalescing packets and have populated more than `min_size` bytes with packets
         // already.
-        self.min_size = Ord::max(
-            self.min_size,
-            self.datagram_start + (min_size as usize) - self.tag_len,
-        );
+        self.min_size = Ord::max(self.min_size, (min_size as usize) - self.tag_len);
     }
 
     pub(super) fn finish_and_track(
@@ -231,9 +229,9 @@ impl PacketBuilder {
         conn: &mut Connection,
         transmits: &mut TransmitBuilder<'_>,
     ) -> (usize, bool) {
-        let pad = transmits.len() < self.min_size;
+        let pad = self.min_size > transmits.datagram().len();
         if pad {
-            let padding_bytes = self.min_size - transmits.len();
+            let padding_bytes = self.min_size - transmits.datagram().len();
             trace!("PADDING * {padding_bytes}");
             transmits.datagram_mut().put_bytes(0, padding_bytes);
         }
@@ -258,13 +256,14 @@ impl PacketBuilder {
             .datagram_mut()
             .put_bytes(0, packet_crypto.tag_len());
         let encode_start = self.partial_encode.start;
-        let packet_buf = &mut transmits.as_mut_slice()[encode_start..];
+        let mut datagram_buf = transmits.datagram_mut();
+        let packet_buf = &mut datagram_buf[encode_start..];
         self.partial_encode.finish(
             packet_buf,
             header_crypto,
             Some((self.exact_number, packet_crypto)),
         );
 
-        (transmits.len() - encode_start, pad)
+        (transmits.datagram().len() - encode_start, pad)
     }
 }
diff --git a/quinn-proto/src/connection/transmit_builder.rs b/quinn-proto/src/connection/transmit_builder.rs
@@ -1,5 +1,9 @@
+use std::ops::{Deref, DerefMut};
+
 use bytes::BufMut;
 
+use super::BufLen;
+
 /// The buffer in which to write datagrams for [`Connection::poll_transmit`]
 ///
 /// The `poll_transmit` function writes zero or more datagrams to a buffer. Multiple
@@ -140,14 +144,24 @@ impl<'a> TransmitBuilder<'a> {
         self.buf_capacity = self.buf.len();
     }
 
-    /// Returns a buffer into which the current datagram can be written
-    pub(super) fn datagram_mut(&mut self) -> bytes::buf::Limit<&mut Vec<u8>> {
-        self.buf.limit(self.buf_capacity)
+    /// Returns a mutable buffer for the current datagram
+    ///
+    /// 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,
+            self.datagram_start,
+            self.buf_capacity - self.datagram_start,
+        )
     }
 
-    /// Returns the already written bytes in the buffer
-    pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
-        self.buf.as_mut_slice()
+    /// Returns the bytes written into the current datagram
+    pub(super) fn datagram(&self) -> &[u8] {
+        &self.buf[self.datagram_start..]
     }
 
     /// Returns the GSO segment size
@@ -196,3 +210,98 @@ impl<'a> TransmitBuilder<'a> {
         self.buf.len()
     }
 }
+
+/// 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> {
+    /// The underlying storage the datagram buffer exists in
+    buf: &'a mut Vec<u8>,
+    /// The start offset of the datagram in the underlying buffer
+    start_offset: usize,
+    /// The maximum write offset in the underlying buffer for this datagram
+    max_offset: usize,
+}
+
+impl<'a> DatagramBuffer<'a> {
+    pub(crate) fn new(buf: &'a mut Vec<u8>, start_offset: usize, max_size: usize) -> Self {
+        let max_offset = start_offset + max_size;
+        DatagramBuffer {
+            buf,
+            start_offset,
+            max_offset,
+        }
+    }
+}
+
+unsafe impl BufMut for DatagramBuffer<'_> {
+    fn remaining_mut(&self) -> usize {
+        self.max_offset - self.buf.len()
+    }
+
+    unsafe fn advance_mut(&mut self, cnt: usize) {
+        self.buf.advance_mut(cnt);
+    }
+
+    fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
+        self.buf.chunk_mut()
+    }
+}
+
+impl Deref for DatagramBuffer<'_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        &self.buf[self.start_offset..]
+    }
+}
+
+impl DerefMut for DatagramBuffer<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.buf[self.start_offset..]
+    }
+}
+
+impl BufSlice for DatagramBuffer<'_> {
+    /// Returns the number of already written bytes in the buffer
+    fn len(&self) -> usize {
+        self.buf.len() - self.start_offset
+    }
+
+    /// Returns the maximum size of the buffer
+    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)
+    }
+}
diff --git a/quinn-proto/src/packet.rs b/quinn-proto/src/packet.rs
@@ -6,7 +6,7 @@ use thiserror::Error;
 use crate::{
     ConnectionId,
     coding::{self, BufExt, BufMutExt},
-    connection::BufLen,
+    connection::BufSlice,
     crypto,
 };
 
@@ -282,7 +282,11 @@ pub(crate) enum Header {
 }
 
 impl Header {
-    pub(crate) fn encode(&self, w: &mut (impl BufMut + BufLen)) -> PartialEncode {
+    /// Encodes the QUIC packet header into the buffer
+    ///
+    /// 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 {
         use Header::*;
         let start = w.len();
         match *self {
