Switch PacketBuilder to use TransmitBuf

This allows making TransmitBuf::buf private at last.  Now all the
logic it handles is fully encapsulated.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -659,7 +659,7 @@ impl Connection {
                         builder.pad_to(buf.segment_size() as u16);
                     }
 
-                    builder.finish_and_track(now, self, sent_frames.take(), buf.buf);
+                    builder.finish_and_track(now, self, sent_frames.take(), &mut buf);
 
                     if buf.num_datagrams() == 1 && space_id == SpaceId::Data {
                         // Now that we know the size of the first datagram, check whether
@@ -702,7 +702,7 @@ impl Connection {
                 // datagram.
                 // Finish current packet without adding extra padding
                 if let Some(builder) = builder_storage.take() {
-                    builder.finish_and_track(now, self, sent_frames.take(), buf.buf);
+                    builder.finish_and_track(now, self, sent_frames.take(), &mut buf);
                 }
             }
 
@@ -826,7 +826,7 @@ impl Connection {
                             non_retransmits: true,
                             ..SentFrames::default()
                         }),
-                        buf.buf,
+                        &mut buf,
                     );
                     self.stats.udp_tx.on_sent(1, buf.len());
                     return Some(Transmit {
@@ -882,7 +882,7 @@ impl Connection {
                 builder.pad_to(MIN_INITIAL_SIZE);
             }
             let last_packet_number = builder.exact_number;
-            builder.finish_and_track(now, self, sent_frames, buf.buf);
+            builder.finish_and_track(now, self, sent_frames, &mut buf);
             self.path
                 .congestion
                 .on_sent(now, buf.len() as u64, last_packet_number);
@@ -920,7 +920,7 @@ impl Connection {
                 non_retransmits: true,
                 ..Default::default()
             };
-            builder.finish_and_track(now, self, Some(sent_frames), buf.buf);
+            builder.finish_and_track(now, self, Some(sent_frames), &mut buf);
 
             self.stats.path.sent_plpmtud_probes += 1;
 
@@ -994,7 +994,7 @@ impl Connection {
         // sending a datagram of this size
         builder.pad_to(MIN_INITIAL_SIZE);
 
-        builder.finish(self, buf.buf);
+        builder.finish(self, buf);
         self.stats.udp_tx.on_sent(1, buf.len());
 
         Some(Transmit {

quinn-proto/src/connection/packet_builder.rs

@@ -1,4 +1,4 @@
-use bytes::Bytes;
+use bytes::{BufMut, Bytes};
 use rand::Rng;
 use tracing::{trace, trace_span};
 
@@ -120,9 +120,9 @@ impl PacketBuilder {
                 version,
             }),
         };
-        let partial_encode = header.encode(buffer.buf);
+        let partial_encode = header.encode(buffer);
         if conn.peer_params.grease_quic_bit && conn.rng.random() {
-            buffer.buf[partial_encode.start] ^= FIXED_BIT;
+            buffer.as_mut_slice()[partial_encode.start] ^= FIXED_BIT;
         }
 
         let (sample_size, tag_len) = if let Some(ref crypto) = space.crypto {
@@ -183,7 +183,7 @@ impl PacketBuilder {
         now: Instant,
         conn: &mut Connection,
         sent: Option<SentFrames>,
-        buffer: &mut Vec<u8>,
+        buffer: &mut TransmitBuf<'_>,
     ) {
         let ack_eliciting = self.ack_eliciting;
         let exact_number = self.exact_number;
@@ -226,11 +226,15 @@ impl PacketBuilder {
     }
 
     /// Encrypt packet, returning the length of the packet and whether padding was added
-    pub(super) fn finish(self, conn: &mut Connection, buffer: &mut Vec<u8>) -> (usize, bool) {
+    pub(super) fn finish(
+        self,
+        conn: &mut Connection,
+        buffer: &mut TransmitBuf<'_>,
+    ) -> (usize, bool) {
         let pad = buffer.len() < self.min_size;
         if pad {
             trace!("PADDING * {}", self.min_size - buffer.len());
-            buffer.resize(self.min_size, 0);
+            buffer.put_bytes(0, self.min_size - buffer.len());
         }
 
         let space = &conn.spaces[self.space];
@@ -249,9 +253,9 @@ impl PacketBuilder {
             "Mismatching crypto tag len"
         );
 
-        buffer.resize(buffer.len() + packet_crypto.tag_len(), 0);
+        buffer.put_bytes(0, packet_crypto.tag_len());
         let encode_start = self.partial_encode.start;
-        let packet_buf = &mut buffer[encode_start..];
+        let packet_buf = &mut buffer.as_mut_slice()[encode_start..];
         self.partial_encode.finish(
             packet_buf,
             header_crypto,

quinn-proto/src/connection/transmit_buf.rs

@@ -27,7 +27,7 @@ use super::BufLen;
 #[derive(Debug)]
 pub(super) struct TransmitBuf<'a> {
     /// The buffer itself, packets are written to this buffer
-    pub(super) buf: &'a mut Vec<u8>,
+    buf: &'a mut Vec<u8>,
     /// Offset into the buffer at which the current datagram starts
     ///
     /// Note that when coalescing packets this might be before the start of the current
@@ -187,6 +187,11 @@ impl<'a> TransmitBuf<'a> {
         }
         self.num_datagrams += 1;
     }
+
+    /// Returns the already written bytes in the buffer
+    pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
+        self.buf.as_mut_slice()
+    }
 }
 
 unsafe impl BufMut for TransmitBuf<'_> {