diff --git a/quinn-proto/src/connection/datagrams.rs b/quinn-proto/src/connection/datagrams.rs index 88cbcce9a..0e0a2b785 100644 --- a/quinn-proto/src/connection/datagrams.rs +++ b/quinn-proto/src/connection/datagrams.rs @@ -4,7 +4,7 @@ use bytes::{BufMut, Bytes}; use thiserror::Error; use tracing::{debug, trace}; -use super::{BufLen, Connection}; +use super::Connection; use crate::{ TransportError, frame::{Datagram, FrameStruct}, @@ -164,13 +164,13 @@ impl DatagramState { /// /// Returns whether a frame was written. At most `max_size` bytes will be written, including /// framing. - pub(super) fn write(&mut self, buf: &mut (impl BufMut + BufLen), max_size: usize) -> bool { + pub(super) fn write(&mut self, buf: &mut impl BufMut) -> bool { let datagram = match self.outgoing.pop_front() { Some(x) => x, None => return false, }; - if buf.len() + datagram.size(true) > max_size { + if buf.remaining_mut() < datagram.size(true) { // Future work: we could be more clever about cramming small datagrams into // mostly-full packets when a larger one is queued first self.outgoing.push_front(datagram); diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index cf134f292..2074c0bf8 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/quinn-proto/src/connection/mod.rs @@ -711,11 +711,11 @@ impl Connection { // to encode the ConnectionClose frame too. However we still have the // check here to prevent crashes if something changes. debug_assert!( - builder.buf.len() + frame::ConnectionClose::SIZE_BOUND < builder.max_size, + builder.frame_space_remaining() > frame::ConnectionClose::SIZE_BOUND, "ACKs should leave space for ConnectionClose" ); - if builder.buf.len() + frame::ConnectionClose::SIZE_BOUND < builder.max_size { - let max_frame_size = builder.max_size - builder.buf.len(); + if frame::ConnectionClose::SIZE_BOUND < builder.frame_space_remaining() { + let max_frame_size = builder.frame_space_remaining(); match self.state { State::Closed(state::Closed { ref reason }) => { if space_id == SpaceId::Data || reason.is_transport_layer() { @@ -789,9 +789,8 @@ impl Connection { } let sent_frames = { - let max_size = builder.max_size; let pn = builder.exact_number; - self.populate_packet(now, space_id, path_id, &mut builder, max_size, pn) + self.populate_packet(now, space_id, path_id, &mut builder.frame_space_mut(), pn) }; // ACK-only packets should only be sent when explicitly allowed. If we write them due to @@ -3478,8 +3477,7 @@ impl Connection { now: Instant, space_id: SpaceId, path_id: PathId, - buf: &mut (impl BufMut + BufLen), - max_size: usize, + buf: &mut impl BufMut, pn: u64, ) -> SentFrames { let mut sent = SentFrames::default(); @@ -3508,7 +3506,7 @@ impl Connection { && (!path.observed_addr_sent || space.pending.observed_addr) { let frame = frame::ObservedAddr::new(path.remote, self.next_observed_addr_seq_no); - if buf.len() + frame.size() < max_size { + if buf.remaining_mut() > frame.size() { frame.write(buf); self.next_observed_addr_seq_no = self.next_observed_addr_seq_no.saturating_add(1u8); @@ -3563,7 +3561,7 @@ impl Connection { // Ensure the delay is within bounds to avoid a PROTOCOL_VIOLATION error let max_ack_delay = self.ack_frequency.candidate_max_ack_delay( - self.paths.get(&path_id).expect("known path").path.rtt.get(), + path.rtt.get(), config, &self.peer_params, ); @@ -3585,9 +3583,8 @@ impl Connection { } // PATH_CHALLENGE - if (buf.len() + 9 < max_size) && space_id == SpaceId::Data { + if buf.remaining_mut() > 9 && space_id == SpaceId::Data { // Transmit challenges with every outgoing frame on an unvalidated path - let path = &mut self.paths.get_mut(&path_id).expect("known path").path; if let Some(token) = path.challenge { // But only send a packet solely for that purpose at most once path.challenge_pending = false; @@ -3607,7 +3604,7 @@ impl Connection { { let frame = frame::ObservedAddr::new(path.remote, self.next_observed_addr_seq_no); - if buf.len() + frame.size() < max_size { + if buf.remaining_mut() > frame.size() { frame.write(buf); self.next_observed_addr_seq_no = @@ -3623,8 +3620,7 @@ impl Connection { } // PATH_RESPONSE - if buf.len() + 9 < max_size && space_id == SpaceId::Data { - let path = &mut self.paths.get_mut(&path_id).expect("known path").path; + if buf.remaining_mut() > 9 && space_id == SpaceId::Data { if let Some(token) = self.path_responses.pop_on_path(path.remote) { sent.non_retransmits = true; sent.requires_padding = true; @@ -3644,7 +3640,7 @@ impl Connection { { let frame = frame::ObservedAddr::new(path.remote, self.next_observed_addr_seq_no); - if buf.len() + frame.size() < max_size { + if buf.remaining_mut() > frame.size() { frame.write(buf); self.next_observed_addr_seq_no = @@ -3660,7 +3656,7 @@ impl Connection { } // CRYPTO - while (buf.len() + frame::Crypto::SIZE_BOUND < max_size) && !is_0rtt { + while buf.remaining_mut() > frame::Crypto::SIZE_BOUND && !is_0rtt { let mut frame = match space.pending.crypto.pop_front() { Some(x) => x, None => break, @@ -3670,8 +3666,7 @@ impl Connection { // Since the offset is known, we can reserve the exact size required to encode it. // For length we reserve 2bytes which allows to encode up to 2^14, // which is more than what fits into normally sized QUIC frames. - let max_crypto_data_size = max_size - - buf.len() + let max_crypto_data_size = buf.remaining_mut() - 1 // Frame Type - VarInt::size(unsafe { VarInt::from_u64_unchecked(frame.offset) }) - 2; // Maximum encoded length for frame size, given we send less than 2^14 bytes @@ -3707,7 +3702,6 @@ impl Connection { &mut space.pending, &mut sent.retransmits, &mut self.stats.frame_tx, - max_size, ); } @@ -3719,7 +3713,7 @@ impl Connection { .max() .expect("some local CID state must exist"); let new_cid_size_bound = frame::NewConnectionId::size_bound(is_multipath_enabled, cid_len); - while buf.len() + new_cid_size_bound < max_size { + while buf.remaining_mut() > new_cid_size_bound { let issued = match space.pending.new_cids.pop() { Some(x) => x, None => break, @@ -3766,7 +3760,7 @@ impl Connection { // RETIRE_CONNECTION_ID let retire_cid_bound = frame::RetireConnectionId::size_bound(is_multipath_enabled); - while buf.len() + retire_cid_bound < max_size { + while buf.remaining_mut() > retire_cid_bound { let (path_id, sequence) = match space.pending.retire_cids.pop() { Some((PathId(0), seq)) if !is_multipath_enabled => (None, seq), Some((path_id, seq)) => (Some(path_id), seq), @@ -3783,8 +3777,8 @@ impl Connection { // DATAGRAM let mut sent_datagrams = false; - while (buf.len() + Datagram::SIZE_BOUND < max_size) && space_id == SpaceId::Data { - match self.datagrams.write(buf, max_size) { + while buf.remaining_mut() > Datagram::SIZE_BOUND && space_id == SpaceId::Data { + match self.datagrams.write(buf) { true => { sent_datagrams = true; sent.non_retransmits = true; @@ -3824,7 +3818,7 @@ impl Connection { token: token.encode(&*server_config.token_key).into(), }; - if buf.len() + new_token.size() >= max_size { + if buf.remaining_mut() < new_token.size() { space.pending.new_tokens.push(remote_addr); break; } @@ -3839,9 +3833,9 @@ impl Connection { // STREAM if space_id == SpaceId::Data { - sent.stream_frames = - self.streams - .write_stream_frames(buf, max_size, self.config.send_fairness); + sent.stream_frames = self + .streams + .write_stream_frames(buf, self.config.send_fairness); self.stats.frame_tx.stream += sent.stream_frames.len() as u64; } diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs index 4a8e7e56e..27e21fa9b 100644 --- a/quinn-proto/src/connection/packet_builder.rs +++ b/quinn-proto/src/connection/packet_builder.rs @@ -30,8 +30,9 @@ pub(super) struct PacketBuilder<'a, 'b> { /// Smallest absolute position in the associated buffer that must be occupied by this packet's /// frames pub(super) min_size: usize, - /// Largest absolute position in the associated buffer that may be occupied by this packet's - /// frames + /// Largest absolute position in the buffer that may be occupied by this packet's frames + /// + /// This takes the size of the cryptographic tag into account. pub(super) max_size: usize, pub(super) tag_len: usize, pub(super) _span: tracing::span::EnteredSpan, @@ -194,6 +195,15 @@ impl<'a, 'b> PacketBuilder<'a, 'b> { ); } + /// Returns a writable buffer limited to the remaining frame space + /// + /// The [`BufMut::remaining_mut`] call on the returned buffer indicates the amount of + /// space available to write QUIC frames into. + // In rust 1.82 we can use `-> impl BufMut + use<'_, 'a, 'b>` + pub(super) fn frame_space_mut(&mut self) -> bytes::buf::Limit<&mut Self> { + self.limit(self.frame_space_remaining()) + } + pub(super) fn finish_and_track( self, now: Instant, @@ -245,6 +255,10 @@ impl<'a, 'b> PacketBuilder<'a, 'b> { /// Encrypt packet, returning the length of the packet and whether padding was added pub(super) fn finish(self, conn: &mut Connection) -> (usize, bool) { + debug_assert!( + self.buf.len() <= self.max_size, + "packet exceeds maximum size" + ); let pad = self.buf.len() < self.min_size; if pad { trace!("PADDING * {}", self.min_size - self.buf.len()); @@ -279,6 +293,15 @@ impl<'a, 'b> PacketBuilder<'a, 'b> { (self.buf.len() - encode_start, pad) } + + /// Returns the remaining space in the packet that can be taken up by QUIC frames + /// + /// This leaves space in the datagram for the cryptographic tag that needs to be written + /// when the packet is finished. + pub(super) fn frame_space_remaining(&self) -> usize { + debug_assert!(self.max_size >= self.buf.len(), "packet exceeds bounds"); + self.max_size.saturating_sub(self.buf.len()) + } } unsafe impl BufMut for PacketBuilder<'_, '_> { diff --git a/quinn-proto/src/connection/streams/state.rs b/quinn-proto/src/connection/streams/state.rs index d6be20923..fe30e1bd2 100644 --- a/quinn-proto/src/connection/streams/state.rs +++ b/quinn-proto/src/connection/streams/state.rs @@ -15,7 +15,7 @@ use super::{ use crate::{ Dir, MAX_STREAM_COUNT, Side, StreamId, TransportError, VarInt, coding::BufMutExt, - connection::{BufLen, stats::FrameStats}, + connection::stats::FrameStats, frame::{self, FrameStruct, StreamMetaVec}, transport_parameters::TransportParameters, }; @@ -411,14 +411,13 @@ impl StreamsState { pub(in crate::connection) fn write_control_frames( &mut self, - buf: &mut (impl BufMut + BufLen), + buf: &mut impl BufMut, pending: &mut Retransmits, retransmits: &mut ThinRetransmits, stats: &mut FrameStats, - max_size: usize, ) { // RESET_STREAM - while buf.len() + frame::ResetStream::SIZE_BOUND < max_size { + while buf.remaining_mut() > frame::ResetStream::SIZE_BOUND { let (id, error_code) = match pending.reset_stream.pop() { Some(x) => x, None => break, @@ -442,7 +441,7 @@ impl StreamsState { } // STOP_SENDING - while buf.len() + frame::StopSending::SIZE_BOUND < max_size { + while buf.remaining_mut() > frame::StopSending::SIZE_BOUND { let frame = match pending.stop_sending.pop() { Some(x) => x, None => break, @@ -461,7 +460,7 @@ impl StreamsState { } // MAX_DATA - if pending.max_data && buf.len() + 9 < max_size { + if pending.max_data && buf.remaining_mut() > 9 { pending.max_data = false; // `local_max_data` can grow bigger than `VarInt`. @@ -484,7 +483,7 @@ impl StreamsState { } // MAX_STREAM_DATA - while buf.len() + 17 < max_size { + while buf.remaining_mut() > 17 { let id = match pending.max_stream_data.iter().next() { Some(x) => *x, None => break, @@ -516,7 +515,7 @@ impl StreamsState { // MAX_STREAMS for dir in Dir::iter() { - if !pending.max_stream_id[dir as usize] || buf.len() + 9 >= max_size { + if !pending.max_stream_id[dir as usize] || buf.remaining_mut() <= 9 { continue; } @@ -541,21 +540,14 @@ impl StreamsState { pub(crate) fn write_stream_frames( &mut self, - buf: &mut (impl BufMut + BufLen), - max_buf_size: usize, + buf: &mut impl BufMut, fair: bool, ) -> StreamMetaVec { let mut stream_frames = StreamMetaVec::new(); - while buf.len() + frame::Stream::SIZE_BOUND < max_buf_size { - if max_buf_size - .checked_sub(buf.len() + frame::Stream::SIZE_BOUND) - .is_none() - { - break; - } - - // Pop the stream of the highest priority that currently has pending data - // If the stream still has some pending data left after writing, it will be reinserted, otherwise not + while buf.remaining_mut() > frame::Stream::SIZE_BOUND { + // Pop the stream of the highest priority that currently has pending data. If + // the stream still has some pending data left after writing, it will be + // reinserted, otherwise not let Some(stream) = self.pending.pop() else { break; }; @@ -577,7 +569,7 @@ impl StreamsState { // Now that we know the `StreamId`, we can better account for how many bytes // are required to encode it. - let max_buf_size = max_buf_size - buf.len() - 1 - VarInt::size(id.into()); + let max_buf_size = buf.remaining_mut() - 1 - VarInt::size(id.into()); let (offsets, encode_length) = stream.pending.poll_transmit(max_buf_size); let fin = offsets.end == stream.pending.offset() && matches!(stream.state, SendState::DataSent { .. }); @@ -1380,7 +1372,7 @@ mod tests { high.write(b"high").unwrap(); let mut buf = Vec::with_capacity(40); - let meta = server.write_stream_frames(&mut buf, 40, true); + let meta = server.write_stream_frames(&mut buf, true); assert_eq!(meta[0].id, id_high); assert_eq!(meta[1].id, id_mid); assert_eq!(meta[2].id, id_low); @@ -1438,16 +1430,18 @@ mod tests { }; high.set_priority(-1).unwrap(); - let mut buf = Vec::with_capacity(1000); - let meta = server.write_stream_frames(&mut buf, 40, true); + let mut buf = Vec::with_capacity(1000).limit(40); + let meta = server.write_stream_frames(&mut buf, true); assert_eq!(meta.len(), 1); assert_eq!(meta[0].id, id_high); // After requeuing we should end up with 2 priorities - not 3 assert_eq!(server.pending.len(), 2); + let mut buf = buf.into_inner(); + // Send the remaining data. The initial mid priority one should go first now - let meta = server.write_stream_frames(&mut buf, 1000, true); + let meta = server.write_stream_frames(&mut buf, true); assert_eq!(meta.len(), 2); assert_eq!(meta[0].id, id_mid); assert_eq!(meta[1].id, id_high); @@ -1507,12 +1501,13 @@ mod tests { // loop until all the streams are written loop { - let buf_len = buf.len(); - let meta = server.write_stream_frames(&mut buf, buf_len + 40, fair); + let mut chunk_buf = buf.limit(40); + let meta = server.write_stream_frames(&mut chunk_buf, fair); if meta.is_empty() { break; } metas.extend(meta); + buf = chunk_buf.into_inner(); } assert!(!server.can_send_stream_data()); @@ -1575,11 +1570,12 @@ mod tests { stream_b.write(&[b'b'; 100]).unwrap(); let mut metas = vec![]; - let mut buf = Vec::with_capacity(1024); + let buf = Vec::with_capacity(1024); // Write the first chunk of stream_a - let buf_len = buf.len(); - let meta = server.write_stream_frames(&mut buf, buf_len + 40, false); + let mut chunk_buf = buf.limit(40); + let meta = server.write_stream_frames(&mut chunk_buf, false); + let mut buf = chunk_buf.into_inner(); assert!(!meta.is_empty()); metas.extend(meta); @@ -1595,8 +1591,9 @@ mod tests { // loop until all the streams are written loop { - let buf_len = buf.len(); - let meta = server.write_stream_frames(&mut buf, buf_len + 40, false); + let mut chunk_buf = buf.limit(40); + let meta = server.write_stream_frames(&mut chunk_buf, false); + buf = chunk_buf.into_inner(); if meta.is_empty() { break; }