noq_proto/connection/
packet_builder.rs

1use bytes::{BufMut, Bytes};
2use rand::RngExt;
3use tracing::{debug, trace, trace_span};
4
5use super::{Connection, PathId, SentFrames, TransmitBuf, spaces::SentPacket};
6use crate::{
7    ConnectionId, FrameStats, Instant, MIN_INITIAL_SIZE, TransportError,
8    coding::Encodable,
9    connection::{ConnectionSide, EncryptionLevel, qlog::QlogSentPacket, spaces::Retransmits},
10    frame::EncodableFrame,
11    packet::{FIXED_BIT, Header, InitialHeader, LongType, PacketNumber, PartialEncode, SpaceId},
12};
13
14/// QUIC packet builder
15///
16/// This allows building QUIC packets: it takes care of writing the header, allows writing
17/// frames and on [`PacketBuilder::finish`] (or [`PacketBuilder::finish_and_track`]) it
18/// encrypts the packet so it is ready to be sent on the wire.
19///
20/// The builder manages the write buffer into which the packet is written, and directly
21/// implements [`BufMut`] to write frames into the packet.
22pub(super) struct PacketBuilder<'a, 'b> {
23    pub(super) buf: &'a mut TransmitBuf<'b>,
24    pub(super) space: SpaceId,
25    path: PathId,
26    pub(super) partial_encode: PartialEncode,
27    pub(super) ack_eliciting: bool,
28    pub(super) packet_number: u64,
29    /// Is this packet allowed to be coalesced?
30    pub(super) can_coalesce: bool,
31    /// Smallest absolute position in the associated buffer that must be occupied by this packet's
32    /// frames
33    pub(super) min_size: usize,
34    pub(super) tag_len: usize,
35    level: EncryptionLevel,
36    pub(super) _span: tracing::span::EnteredSpan,
37    qlog: QlogSentPacket,
38    sent_frames: SentFrames,
39}
40
41impl<'a, 'b> PacketBuilder<'a, 'b> {
42    /// Write a new packet header to `buffer` and determine the packet's properties
43    ///
44    /// Marks the connection drained and returns `None` if the confidentiality limit would be
45    /// violated.
46    pub(super) fn new(
47        now: Instant,
48        space_id: SpaceId,
49        path_id: PathId,
50        dst_cid: ConnectionId,
51        buffer: &'a mut TransmitBuf<'b>,
52        conn: &mut Connection,
53    ) -> Option<Self>
54    where
55        'b: 'a,
56    {
57        let mut qlog = QlogSentPacket::default();
58
59        let version = conn.version;
60        // Initiate key update if we're approaching the confidentiality limit
61        let (_header_crypto, _packet_crypto, level) = conn
62            .crypto_state
63            .encryption_keys(space_id.kind(), conn.side.side())
64            .expect("tried to build packet without encryption keys");
65
66        let remaining_packet_budget = conn
67            .crypto_state
68            .remaining_packet_budget(level)
69            .expect("keys are installed");
70
71        if level == EncryptionLevel::OneRtt {
72            // 1-RTT keys can be rotated, so exhaustion just triggers a key update
73            if remaining_packet_budget == 0 {
74                debug!("routine key update due to phase exhaustion");
75                conn.force_key_update();
76            }
77        } else {
78            // Other encryption levels have a fixed key budget
79            if remaining_packet_budget == 0 {
80                let close = TransportError::AEAD_LIMIT_REACHED("confidentiality limit reached");
81                conn.kill(close.into());
82                return None;
83            } else if remaining_packet_budget == 1 {
84                let close = TransportError::AEAD_LIMIT_REACHED("confidentiality limit reached");
85                conn.close_inner(now, close.into());
86            }
87        }
88
89        let (header_crypto, packet_crypto, level) = conn
90            .crypto_state
91            .encryption_keys(space_id.kind(), conn.side.side())
92            .expect("verified");
93
94        let space = &mut conn.spaces[space_id];
95        let packet_number = space.for_path(path_id).get_tx_number(&mut conn.rng);
96        let span = trace_span!("send", space = ?space_id, pn = packet_number, %path_id).entered();
97
98        let number = PacketNumber::new(
99            packet_number,
100            space.for_path(path_id).largest_acked_packet_pn.unwrap_or(0),
101        );
102        let header = match level {
103            EncryptionLevel::OneRtt => Header::Short {
104                dst_cid,
105                number,
106                spin: if conn.spin_enabled {
107                    conn.spin
108                } else {
109                    conn.rng.random()
110                },
111                key_phase: conn.crypto_state.key_phase,
112            },
113            EncryptionLevel::ZeroRtt => Header::Long {
114                ty: LongType::ZeroRtt,
115                src_cid: conn.handshake_cid,
116                dst_cid,
117                number,
118                version,
119            },
120            EncryptionLevel::Handshake => Header::Long {
121                ty: LongType::Handshake,
122                src_cid: conn.handshake_cid,
123                dst_cid,
124                number,
125                version,
126            },
127            EncryptionLevel::Initial => Header::Initial(InitialHeader {
128                src_cid: conn.handshake_cid,
129                dst_cid,
130                token: match &conn.side {
131                    ConnectionSide::Client { token, .. } => token.clone(),
132                    ConnectionSide::Server { .. } => Bytes::new(),
133                },
134                number,
135                version,
136            }),
137        };
138
139        let partial_encode = header.encode(buffer);
140        if conn.peer_params.grease_quic_bit && conn.rng.random() {
141            buffer.as_mut_slice()[partial_encode.start] ^= FIXED_BIT;
142        }
143
144        let (sample_size, tag_len) = (header_crypto.sample_size(), packet_crypto.tag_len());
145
146        // Each packet must be large enough for header protection sampling, i.e. the combined
147        // lengths of the encoded packet number and protected payload must be at least 4 bytes
148        // longer than the sample required for header protection. Further, each packet should be at
149        // least tag_len + 6 bytes larger than the destination CID on incoming packets so that the
150        // peer may send stateless resets that are indistinguishable from regular traffic.
151
152        // pn_len + payload_len + tag_len >= sample_size + 4
153        // payload_len >= sample_size + 4 - pn_len - tag_len
154        let min_size = Ord::max(
155            buffer.len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
156            partial_encode.start + dst_cid.len() + 6,
157        );
158        let max_size = buffer.datagram_max_offset() - tag_len;
159        debug_assert!(max_size >= min_size);
160
161        qlog.header(&header, Some(packet_number), level, path_id);
162
163        Some(Self {
164            buf: buffer,
165            space: space_id,
166            path: path_id,
167            partial_encode,
168            packet_number,
169            can_coalesce: header.can_coalesce(),
170            min_size,
171            tag_len,
172            level,
173            ack_eliciting: false,
174            qlog,
175            sent_frames: SentFrames::default(),
176            _span: span,
177        })
178    }
179
180    #[cfg(test)]
181    pub(crate) fn simple_data_buf(buf: &'a mut TransmitBuf<'b>) -> Self {
182        Self {
183            buf,
184            space: SpaceId::Data,
185            path: PathId::ZERO,
186            partial_encode: PartialEncode::no_header(),
187            ack_eliciting: true,
188            packet_number: 0,
189            can_coalesce: true,
190            min_size: 0,
191            tag_len: 0,
192            _span: trace_span!("test").entered(),
193            qlog: QlogSentPacket::default(),
194            sent_frames: SentFrames::default(),
195            level: EncryptionLevel::Initial,
196        }
197    }
198
199    /// Append the minimum amount of padding to the packet such that, after encryption, the
200    /// enclosing datagram will occupy at least `min_size` bytes
201    pub(super) fn pad_to(&mut self, min_size: u16) {
202        // The datagram might already have a larger minimum size than the caller is requesting, if
203        // e.g. we're coalescing packets and have populated more than `min_size` bytes with packets
204        // already.
205        self.min_size = Ord::max(
206            self.min_size,
207            self.buf.datagram_start_offset() + (min_size as usize) - self.tag_len,
208        );
209    }
210
211    /// Writes a frame into the underlying buffer.
212    ///
213    /// It will also:
214    /// - Track the frame so that it's registered with the path once [`Self::finish_and_track`] is
215    ///   called.
216    /// - Register the sent frame with the given [`FrameStats`].
217    /// - If the qlog feature is enabled, register the frame.
218    /// - Log the frame.
219    pub(super) fn write_frame<'c>(
220        &mut self,
221        frame: impl Into<EncodableFrame<'c>>,
222        stats: &mut FrameStats,
223    ) {
224        self.write_frame_with_log_msg(frame, stats, None);
225    }
226
227    /// Writes a frame into the underlying buffer.
228    ///
229    /// It will also:
230    /// - Track the frame so that it's registered with the path once [`Self::finish_and_track`] is
231    ///   called.
232    /// - Register the sent frame with the given [`FrameStats`].
233    /// - If the qlog feature is enabled, register the frame.
234    /// - Log the frame. If a `msg` is given, this will be added to the log.
235    pub(super) fn write_frame_with_log_msg<'c>(
236        &mut self,
237        frame: impl Into<EncodableFrame<'c>>,
238        stats: &mut FrameStats,
239        msg: Option<&'static str>,
240    ) {
241        let frame = frame.into();
242        frame.encode(&mut self.frame_space_mut());
243        self.ack_eliciting |= frame.is_ack_eliciting();
244        stats.record(frame.get_type());
245        self.qlog.record(&frame);
246        match msg {
247            Some(msg) => trace!(%frame, msg),
248            None => trace!(%frame),
249        }
250        self.sent_frames.record_sent_frame(frame);
251    }
252
253    /// Returns a writable buffer limited to the remaining frame space
254    ///
255    /// The [`BufMut::remaining_mut`] call on the returned buffer indicates the amount of
256    /// space available to write QUIC frames into.
257    // In rust 1.82 we can use `-> impl BufMut + use<'_, 'a, 'b>`
258    fn frame_space_mut(&mut self) -> bytes::buf::Limit<&mut TransmitBuf<'b>> {
259        self.buf.limit(self.frame_space_remaining())
260    }
261
262    pub(super) fn sent_frames(&self) -> &SentFrames {
263        &self.sent_frames
264    }
265
266    pub(super) fn finish_and_track(
267        mut self,
268        now: Instant,
269        conn: &mut Connection,
270        path_id: PathId,
271        pad_datagram: PadDatagram,
272    ) {
273        match pad_datagram {
274            PadDatagram::No => (),
275            PadDatagram::ToSize(size) => self.pad_to(size),
276            PadDatagram::ToSegmentSize => self.pad_to(self.buf.segment_size() as u16),
277            PadDatagram::ToMinMtu => self.pad_to(MIN_INITIAL_SIZE),
278        }
279        let ack_eliciting = self.ack_eliciting;
280        let packet_number = self.packet_number;
281        let space_id = self.space;
282        let (size, padded, sent) = self.finish(conn, now);
283
284        let size = match padded || ack_eliciting {
285            true => size as u16,
286            false => 0,
287        };
288
289        let packet = SentPacket {
290            path_generation: conn.paths.get_mut(&path_id).unwrap().data.generation(),
291            largest_acked: sent.largest_acked,
292            time_sent: now,
293            size,
294            ack_eliciting,
295            retransmits: sent.retransmits,
296            path_retransmits: sent.path_retransmits,
297            stream_frames: sent.stream_frames,
298        };
299
300        conn.paths.get_mut(&path_id).unwrap().data.sent(
301            packet_number,
302            packet,
303            conn.spaces[space_id].for_path(path_id),
304        );
305        conn.reset_keep_alive(path_id, now);
306        if size != 0 {
307            if ack_eliciting {
308                conn.spaces[space_id]
309                    .for_path(path_id)
310                    .time_of_last_ack_eliciting_packet = Some(now);
311                if conn.path_data(path_id).permit_idle_reset {
312                    conn.reset_idle_timeout(now, space_id.kind(), path_id);
313                }
314                conn.path_data_mut(path_id).permit_idle_reset = false;
315                conn.path_data_mut(path_id)
316                    .congestion
317                    .on_packet_sent(now, size, packet_number)
318            }
319            conn.set_loss_detection_timer(now, path_id);
320            conn.path_data_mut(path_id).pacing.on_transmit(size);
321        }
322    }
323
324    /// Encrypt packet, returning the length of the packet and whether padding was added
325    pub(super) fn finish(
326        mut self,
327        conn: &mut Connection,
328        now: Instant,
329    ) -> (usize, bool, SentFrames) {
330        debug_assert!(
331            self.buf.len() <= self.buf.datagram_max_offset() - self.tag_len,
332            "packet exceeds maximum size"
333        );
334        let pad = self.buf.len() < self.min_size;
335        if pad {
336            let padding = self.min_size - self.buf.len();
337            trace!("PADDING * {}", padding);
338            self.buf.put_bytes(0, padding);
339            self.qlog.frame_padding(padding);
340        }
341
342        let (header_crypto, packet_crypto) = conn
343            .crypto_state
344            .local_crypto(self.level)
345            .expect("tried to send packet without keys");
346
347        debug_assert_eq!(
348            packet_crypto.tag_len(),
349            self.tag_len,
350            "Mismatching crypto tag len"
351        );
352
353        self.buf.put_bytes(0, packet_crypto.tag_len());
354        let encode_start = self.partial_encode.start;
355        let packet_buf = &mut self.buf.as_mut_slice()[encode_start..];
356        // for packet protection, PathId::ZERO and no path are equivalent.
357        self.partial_encode.finish(
358            packet_buf,
359            header_crypto,
360            Some((self.packet_number, self.path, packet_crypto)),
361        );
362
363        conn.crypto_state.inc_sent_with_keys(self.level);
364
365        let packet_len = self.buf.len() - encode_start;
366        trace!(size = %packet_len, "wrote packet");
367        self.qlog.finalize(packet_len);
368        conn.qlog.emit_packet_sent(self.qlog, now);
369        (packet_len, pad, self.sent_frames)
370    }
371
372    /// The number of additional bytes the current packet would take up if it was finished now
373    ///
374    /// This will include any padding which is required to make the size large enough to be
375    /// encrypted correctly.
376    pub(super) fn predict_packet_end(&self) -> usize {
377        self.buf.len().max(self.min_size) + self.tag_len - self.buf.len()
378    }
379
380    /// Returns the remaining space in the packet that can be taken up by QUIC frames
381    ///
382    /// This leaves space in the datagram for the cryptographic tag that needs to be written
383    /// when the packet is finished.
384    pub(super) fn frame_space_remaining(&self) -> usize {
385        let max_offset = self.buf.datagram_max_offset() - self.tag_len;
386        max_offset.saturating_sub(self.buf.len())
387    }
388
389    pub(crate) fn require_padding(&mut self) {
390        self.sent_frames.requires_padding = true;
391    }
392
393    pub(crate) fn retransmits_mut(&mut self) -> &mut Retransmits {
394        self.sent_frames.retransmits_mut()
395    }
396}
397
398#[derive(Debug, Copy, Clone)]
399pub(super) enum PadDatagram {
400    /// Do not pad the datagram
401    No,
402    /// To a specific size
403    ToSize(u16),
404    /// Pad to the current MTU/segment size
405    ///
406    /// For the first datagram in a transmit the MTU is the same as the
407    /// [`TransmitBuf::segment_size`].
408    ToSegmentSize,
409    /// Pad to [`MIN_INITIAL_SIZE`], the minimal QUIC MTU of 1200 bytes
410    ToMinMtu,
411}
412
413impl std::ops::BitOrAssign for PadDatagram {
414    fn bitor_assign(&mut self, rhs: Self) {
415        *self = *self | rhs;
416    }
417}
418
419impl std::ops::BitOr for PadDatagram {
420    type Output = Self;
421
422    fn bitor(self, rhs: Self) -> Self::Output {
423        match (self, rhs) {
424            (Self::No, rhs) => rhs,
425            (Self::ToSize(size), Self::No) => Self::ToSize(size),
426            (Self::ToSize(a), Self::ToSize(b)) => Self::ToSize(a.max(b)),
427            (Self::ToSize(_), Self::ToSegmentSize) => Self::ToSegmentSize,
428            (Self::ToSize(_), Self::ToMinMtu) => Self::ToMinMtu,
429            (Self::ToSegmentSize, Self::No) => Self::ToSegmentSize,
430            (Self::ToSegmentSize, Self::ToSize(_)) => Self::ToSegmentSize,
431            (Self::ToSegmentSize, Self::ToSegmentSize) => Self::ToSegmentSize,
432            (Self::ToSegmentSize, Self::ToMinMtu) => Self::ToMinMtu,
433            (Self::ToMinMtu, _) => Self::ToMinMtu,
434        }
435    }
436}
noq_proto/connection/packet_builder.rs

noq_proto/connection/
packet_builder.rs
