Pass the TransmitBuf directly to the PacketBuilder

This removes the extra arguments from Packetbuilder::new that tell the
builder about datagram boundaries in favour of using the state of
TransmitBuf directly.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -759,9 +759,7 @@ impl Connection {
                 now,
                 space_id,
                 self.rem_cids.active(),
-                buf.buf,
-                buf.buf_capacity,
-                buf.datagram_start,
+                &mut buf,
                 ack_eliciting,
                 self,
             )?);
@@ -926,19 +924,12 @@ impl Connection {
                 .mtud
                 .poll_transmit(now, self.packet_number_filter.peek(&self.spaces[space_id]))?;
 
-            let buf_capacity = probe_size as usize;
-            buf.buf.reserve(buf_capacity);
+            buf.buf_capacity = probe_size as usize;
+            buf.buf.reserve(buf.buf_capacity);
 
-            let mut builder = PacketBuilder::new(
-                now,
-                space_id,
-                self.rem_cids.active(),
-                buf.buf,
-                buf_capacity,
-                0,
-                true,
-                self,
-            )?;
+            debug_assert_eq!(buf.datagram_start, 0);
+            let mut builder =
+                PacketBuilder::new(now, space_id, self.rem_cids.active(), &mut buf, true, self)?;
 
             // We implement MTU probes as ping packets padded up to the probe size
             buf.write(frame::FrameType::PING);
@@ -995,7 +986,11 @@ impl Connection {
     }
 
     /// Send PATH_CHALLENGE for a previous path if necessary
-    fn send_path_challenge(&mut self, now: Instant, buf: &mut TransmitBuf<'_>) -> Option<Transmit> {
+    fn send_path_challenge(
+        &mut self,
+        now: Instant,
+        mut buf: &mut TransmitBuf<'_>,
+    ) -> Option<Transmit> {
         let (prev_cid, prev_path) = self.prev_path.as_mut()?;
         if !prev_path.challenge_pending {
             return None;
@@ -1012,23 +1007,15 @@ impl Connection {
         );
         buf.buf.reserve(MIN_INITIAL_SIZE as usize);
 
-        let buf_capacity = buf.buf.capacity();
+        buf.buf_capacity = buf.buf.capacity();
 
         // Use the previous CID to avoid linking the new path with the previous path. We
         // don't bother accounting for possible retirement of that prev_cid because this is
         // sent once, immediately after migration, when the CID is known to be valid. Even
         // if a post-migration packet caused the CID to be retired, it's fair to pretend
         // this is sent first.
-        let mut builder = PacketBuilder::new(
-            now,
-            SpaceId::Data,
-            *prev_cid,
-            buf.buf,
-            buf_capacity,
-            0,
-            false,
-            self,
-        )?;
+        debug_assert_eq!(buf.datagram_start, 0);
+        let mut builder = PacketBuilder::new(now, SpaceId::Data, *prev_cid, buf, false, self)?;
         trace!("validating previous path with PATH_CHALLENGE {:08x}", token);
         buf.write(frame::FrameType::PATH_CHALLENGE);
         buf.write(token);

quinn-proto/src/connection/packet_builder.rs

@@ -2,7 +2,7 @@ use bytes::Bytes;
 use rand::Rng;
 use tracing::{trace, trace_span};
 
-use super::{Connection, SentFrames, spaces::SentPacket};
+use super::{Connection, SentFrames, TransmitBuf, spaces::SentPacket};
 use crate::{
     ConnectionId, Instant, TransportError, TransportErrorCode,
     connection::ConnectionSide,
@@ -36,9 +36,7 @@ impl PacketBuilder {
         now: Instant,
         space_id: SpaceId,
         dst_cid: ConnectionId,
-        buffer: &mut Vec<u8>,
-        buffer_capacity: usize,
-        datagram_start: usize,
+        buffer: &mut TransmitBuf<'_>,
         ack_eliciting: bool,
         conn: &mut Connection,
     ) -> Option<Self> {
@@ -122,9 +120,9 @@ impl PacketBuilder {
                 version,
             }),
         };
-        let partial_encode = header.encode(buffer);
+        let partial_encode = header.encode(buffer.buf);
         if conn.peer_params.grease_quic_bit && conn.rng.random() {
-            buffer[partial_encode.start] ^= FIXED_BIT;
+            buffer.buf[partial_encode.start] ^= FIXED_BIT;
         }
 
         let (sample_size, tag_len) = if let Some(ref crypto) = space.crypto {
@@ -151,11 +149,11 @@ impl PacketBuilder {
             buffer.len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
             partial_encode.start + dst_cid.len() + 6,
         );
-        let max_size = buffer_capacity - tag_len;
+        let max_size = buffer.buf_capacity - tag_len;
         debug_assert!(max_size >= min_size);
 
         Some(Self {
-            datagram_start,
+            datagram_start: buffer.datagram_start,
             space: space_id,
             partial_encode,
             exact_number,