No BufMut on TransmitBuilder

BufMut is no longer directly implemented on TransmitBuilder, instead
you need to call TransmitBuilder::datagram_mut to get a BufMut for the
current datagram.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -757,7 +757,7 @@ impl Connection {
                         self.receiving_ecn,
                         &mut SentFrames::default(),
                         &mut self.spaces[space_id],
-                        &mut transmit,
+                        &mut transmit.datagram_mut(),
                         &mut self.stats,
                     );
                 }
@@ -774,22 +774,22 @@ impl Connection {
                     match self.state {
                         State::Closed(state::Closed { ref reason }) => {
                             if space_id == SpaceId::Data || reason.is_transport_layer() {
-                                reason.encode(&mut transmit, max_frame_size)
+                                reason.encode(&mut transmit.datagram_mut(), max_frame_size)
                             } else {
                                 frame::ConnectionClose {
                                     error_code: TransportErrorCode::APPLICATION_ERROR,
                                     frame_type: None,
                                     reason: Bytes::new(),
                                 }
-                                .encode(&mut transmit, max_frame_size)
+                                .encode(&mut transmit.datagram_mut(), max_frame_size)
                             }
                         }
                         State::Draining => frame::ConnectionClose {
                             error_code: TransportErrorCode::NO_ERROR,
                             frame_type: None,
                             reason: Bytes::new(),
                         }
-                        .encode(&mut transmit, max_frame_size),
+                        .encode(&mut transmit.datagram_mut(), max_frame_size),
                         _ => unreachable!(
                             "tried to make a close packet when the connection wasn't closed"
                         ),
@@ -817,8 +817,10 @@ impl Connection {
                     // above.
                     let mut builder = builder_storage.take().unwrap();
                     trace!("PATH_RESPONSE {:08x} (off-path)", token);
-                    transmit.write(frame::FrameType::PATH_RESPONSE);
-                    transmit.write(token);
+                    transmit
+                        .datagram_mut()
+                        .write(frame::FrameType::PATH_RESPONSE);
+                    transmit.datagram_mut().write(token);
                     self.stats.frame_tx.path_response += 1;
                     builder.pad_to(MIN_INITIAL_SIZE);
                     builder.finish_and_track(
@@ -844,7 +846,7 @@ impl Connection {
             let sent = self.populate_packet(
                 now,
                 space_id,
-                &mut transmit,
+                &mut transmit.datagram_mut(),
                 builder.max_size,
                 builder.exact_number,
             );
@@ -914,12 +916,14 @@ impl Connection {
             )?;
 
             // We implement MTU probes as ping packets padded up to the probe size
-            transmit.write(frame::FrameType::PING);
+            transmit.datagram_mut().write(frame::FrameType::PING);
             self.stats.frame_tx.ping += 1;
 
             // If supported by the peer, we want no delays to the probe's ACK
             if self.peer_supports_ack_frequency() {
-                transmit.write(frame::FrameType::IMMEDIATE_ACK);
+                transmit
+                    .datagram_mut()
+                    .write(frame::FrameType::IMMEDIATE_ACK);
                 self.stats.frame_tx.immediate_ack += 1;
             }
 
@@ -996,8 +1000,8 @@ impl Connection {
         debug_assert_eq!(buf.datagram_start_offset(), 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);
+        buf.datagram_mut().write(frame::FrameType::PATH_CHALLENGE);
+        buf.datagram_mut().write(token);
         self.stats.frame_tx.path_challenge += 1;
 
         // An endpoint MUST expand datagrams that contain a PATH_CHALLENGE frame
@@ -3985,6 +3989,12 @@ impl BufLen for Vec<u8> {
     }
 }
 
+impl BufLen for bytes::buf::Limit<&mut Vec<u8>> {
+    fn len(&self) -> usize {
+        self.get_ref().len()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

quinn-proto/src/connection/packet_builder.rs

@@ -120,7 +120,7 @@ impl PacketBuilder {
                 version,
             }),
         };
-        let partial_encode = header.encode(buffer);
+        let partial_encode = header.encode(&mut buffer.datagram_mut());
         if conn.peer_params.grease_quic_bit && conn.rng.random() {
             buffer.as_mut_slice()[partial_encode.start] ^= FIXED_BIT;
         }
@@ -233,8 +233,9 @@ impl PacketBuilder {
     ) -> (usize, bool) {
         let pad = buffer.len() < self.min_size;
         if pad {
-            trace!("PADDING * {}", self.min_size - buffer.len());
-            buffer.put_bytes(0, self.min_size - buffer.len());
+            let padding_bytes = self.min_size - buffer.len();
+            trace!("PADDING * {padding_bytes}");
+            buffer.datagram_mut().put_bytes(0, padding_bytes);
         }
 
         let space = &conn.spaces[self.space];
@@ -253,7 +254,7 @@ impl PacketBuilder {
             "Mismatching crypto tag len"
         );
 
-        buffer.put_bytes(0, packet_crypto.tag_len());
+        buffer.datagram_mut().put_bytes(0, packet_crypto.tag_len());
         let encode_start = self.partial_encode.start;
         let packet_buf = &mut buffer.as_mut_slice()[encode_start..];
         self.partial_encode.finish(

quinn-proto/src/connection/transmit_builder.rs

@@ -1,7 +1,5 @@
 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
@@ -147,6 +145,11 @@ impl<'a> TransmitBuilder<'a> {
         self.buf.as_mut_slice()
     }
 
+    /// 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 the GSO segment size
     ///
     /// This is also the maximum size datagrams are allowed to be. The first and last
@@ -193,23 +196,3 @@ impl<'a> TransmitBuilder<'a> {
         self.buf.len()
     }
 }
-
-unsafe impl BufMut for TransmitBuilder<'_> {
-    fn remaining_mut(&self) -> usize {
-        self.buf.remaining_mut()
-    }
-
-    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 BufLen for TransmitBuilder<'_> {
-    fn len(&self) -> usize {
-        self.len()
-    }
-}