From e360832b7a7bd7a2fc400858cf80603aa078ea37 Mon Sep 17 00:00:00 2001 From: Floris Bruynooghe Date: Thu, 10 Apr 2025 17:31:46 +0200 Subject: [PATCH] Switch PacketBuilder to use TransmitBuf This allows making TransmitBuf::buf private at last. Now all the logic it handles is fully encapsulated. --- quinn-proto/src/connection/mod.rs | 12 ++++++------ quinn-proto/src/connection/packet_builder.rs | 20 ++++++++++++-------- quinn-proto/src/connection/transmit_buf.rs | 7 ++++++- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index bf9843d59..a90e77877 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/quinn-proto/src/connection/mod.rs @@ -717,7 +717,7 @@ impl Connection { builder.pad_to(buf.segment_size() as u16); } - builder.finish_and_track(now, self, path_id, sent_frames.take(), buf.buf); + builder.finish_and_track(now, self, path_id, sent_frames.take(), &mut buf); if buf.num_datagrams() == 1 { buf.clip_datagram_size(); @@ -762,7 +762,7 @@ impl Connection { // datagram. // Finish current packet without adding extra padding if let Some(builder) = builder_storage.take() { - builder.finish_and_track(now, self, path_id, sent_frames.take(), buf.buf); + builder.finish_and_track(now, self, path_id, sent_frames.take(), &mut buf); } } @@ -898,7 +898,7 @@ impl Connection { non_retransmits: true, ..SentFrames::default() }), - buf.buf, + &mut buf, ); self.stats.udp_tx.on_sent(1, buf.len()); return Some(Transmit { @@ -955,7 +955,7 @@ impl Connection { builder.pad_to(MIN_INITIAL_SIZE); } let last_packet_number = builder.exact_number; - builder.finish_and_track(now, self, path_id, sent_frames, buf.buf); + builder.finish_and_track(now, self, path_id, sent_frames, &mut buf); self.path_data_mut(path_id).congestion.on_sent( now, buf.len() as u64, @@ -1006,7 +1006,7 @@ impl Connection { non_retransmits: true, ..Default::default() }; - builder.finish_and_track(now, self, path_id, Some(sent_frames), buf.buf); + builder.finish_and_track(now, self, path_id, Some(sent_frames), &mut buf); self.stats.path.sent_plpmtud_probes += 1; @@ -1086,7 +1086,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 { diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs index 4134f8774..b1057d565 100644 --- a/quinn-proto/src/connection/packet_builder.rs +++ b/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}; @@ -121,9 +121,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 { @@ -186,7 +186,7 @@ impl PacketBuilder { conn: &mut Connection, path_id: PathId, sent: Option, - buffer: &mut Vec, + buffer: &mut TransmitBuf<'_>, ) { let ack_eliciting = self.ack_eliciting; let exact_number = self.exact_number; @@ -235,11 +235,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) -> (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]; @@ -258,9 +262,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..]; // for packet protection, PathId(0) and no path are equivalent. self.partial_encode.finish( packet_buf, diff --git a/quinn-proto/src/connection/transmit_buf.rs b/quinn-proto/src/connection/transmit_buf.rs index a8188e156..61d3217c3 100644 --- a/quinn-proto/src/connection/transmit_buf.rs +++ b/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, + buf: &'a mut Vec, /// 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> { pub(super) fn len(&self) -> usize { self.buf.len() } + + /// 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<'_> {