From 3aa89b4ccfee8e080fa756507935541816586e7e Mon Sep 17 00:00:00 2001
From: Floris Bruynooghe <flub@n0.computer>
Date: Thu, 10 Apr 2025 14:50:13 +0200
Subject: [PATCH] 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.
---
 quinn-proto/src/connection/mod.rs            | 25 ++++++--------------
 quinn-proto/src/connection/packet_builder.rs | 14 +++++------
 2 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index 1f33c218c..2a73f088a 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -822,9 +822,7 @@ impl Connection {
                 space_id,
                 path_id,
                 self.rem_cids.get(&path_id).unwrap().active(),
-                buf.buf,
-                buf.buf_capacity,
-                buf.datagram_start,
+                &mut buf,
                 ack_eliciting,
                 self,
             )?);
@@ -1004,6 +1002,7 @@ impl Connection {
             buf.buf_capacity = probe_size as usize;
             buf.buf.reserve(buf.buf_capacity);
 
+            debug_assert_eq!(buf.datagram_start, 0);
             // TODO(flub): I'm not particularly happy about this unwrap.  But let's leave it
             //    for now until more stuff is settled.  We probably should check earlier on
             //    in poll_transmit that we have a valid CID to use.
@@ -1012,9 +1011,7 @@ impl Connection {
                 space_id,
                 path_id,
                 self.rem_cids.get(&path_id).unwrap().active(),
-                buf.buf,
-                buf.buf_capacity,
-                0,
+                &mut buf,
                 true,
                 self,
             )?;
@@ -1096,24 +1093,16 @@ 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,
-            path_id,
-            *prev_cid,
-            buf.buf,
-            buf_capacity,
-            0,
-            false,
-            self,
-        )?;
+        debug_assert_eq!(buf.datagram_start, 0);
+        let mut builder =
+            PacketBuilder::new(now, SpaceId::Data, path_id, *prev_cid, buf, false, self)?;
         trace!("validating previous path with PATH_CHALLENGE {:08x}", token);
         buf.write(frame::FrameType::PATH_CHALLENGE);
         buf.write(token);
diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs
index aab4a15be..78a300737 100644
--- a/quinn-proto/src/connection/packet_builder.rs
+++ b/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, PathId, SentFrames, spaces::SentPacket};
+use super::{Connection, PathId, SentFrames, TransmitBuf, spaces::SentPacket};
 use crate::{
     ConnectionId, Instant, TransportError, TransportErrorCode,
     connection::ConnectionSide,
@@ -38,9 +38,7 @@ impl PacketBuilder {
         space_id: SpaceId,
         path_id: PathId,
         dst_cid: ConnectionId,
-        buffer: &mut Vec<u8>,
-        buffer_capacity: usize,
-        datagram_start: usize,
+        buffer: &mut TransmitBuf<'_>,
         ack_eliciting: bool,
         conn: &mut Connection,
     ) -> Option<Self> {
@@ -123,9 +121,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 {
@@ -152,11 +150,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,
             path: path_id,
             partial_encode,