noq_proto/
packet.rs

1use std::{cmp::Ordering, io, ops::Range, str};
2
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4use thiserror::Error;
5
6use crate::{
7    ConnectionId, PathId,
8    coding::{self, BufExt, BufMutExt},
9    connection::{EncryptionLevel, SpaceKind},
10    crypto,
11};
12
13/// Decodes a QUIC packet's invariant header
14///
15/// Due to packet number encryption, it is impossible to fully decode a header
16/// (which includes a variable-length packet number) without crypto context.
17/// The crypto context (represented by the `Crypto` type in noq) is usually
18/// part of the `Connection`, or can be derived from the destination CID for
19/// Initial packets.
20///
21/// To cope with this, we decode the invariant header (which should be stable
22/// across QUIC versions), which gives us the destination CID and allows us
23/// to inspect the version and packet type (which depends on the version).
24/// This information allows us to fully decode and decrypt the packet.
25#[cfg_attr(test, derive(Clone))]
26#[derive(Debug)]
27pub struct PartialDecode {
28    plain_header: ProtectedHeader,
29    buf: io::Cursor<BytesMut>,
30}
31
32#[allow(clippy::len_without_is_empty)]
33impl PartialDecode {
34    /// Begins decoding a QUIC packet from `bytes`.
35    ///
36    /// Returns any trailing data not part of that packet.
37    pub fn new(
38        bytes: BytesMut,
39        cid_parser: &(impl ConnectionIdParser + ?Sized),
40        supported_versions: &[u32],
41        grease_quic_bit: bool,
42    ) -> Result<(Self, Option<BytesMut>), PacketDecodeError> {
43        let mut buf = io::Cursor::new(bytes);
44        let plain_header =
45            ProtectedHeader::decode(&mut buf, cid_parser, supported_versions, grease_quic_bit)?;
46        let dgram_len = buf.get_ref().len();
47        let packet_len = plain_header
48            .payload_len()
49            .map(|len| (buf.position() + len) as usize)
50            .unwrap_or(dgram_len);
51        match dgram_len.cmp(&packet_len) {
52            Ordering::Equal => Ok((Self { plain_header, buf }, None)),
53            Ordering::Less => Err(PacketDecodeError::InvalidHeader(
54                "packet too short to contain payload length",
55            )),
56            Ordering::Greater => {
57                let rest = Some(buf.get_mut().split_off(packet_len));
58                Ok((Self { plain_header, buf }, rest))
59            }
60        }
61    }
62
63    /// The underlying partially-decoded packet data
64    pub(crate) fn data(&self) -> &[u8] {
65        self.buf.get_ref()
66    }
67
68    pub(crate) fn initial_header(&self) -> Option<&ProtectedInitialHeader> {
69        self.plain_header.as_initial()
70    }
71
72    pub(crate) fn has_long_header(&self) -> bool {
73        !matches!(self.plain_header, ProtectedHeader::Short { .. })
74    }
75
76    pub(crate) fn is_initial(&self) -> bool {
77        self.encryption_level() == Some(EncryptionLevel::Initial)
78    }
79
80    pub(crate) fn encryption_level(&self) -> Option<EncryptionLevel> {
81        use ProtectedHeader::*;
82        match self.plain_header {
83            Initial { .. } => Some(EncryptionLevel::Initial),
84            Long { ty, .. } => Some(match ty {
85                LongType::Handshake => EncryptionLevel::Handshake,
86                LongType::ZeroRtt => EncryptionLevel::ZeroRtt,
87            }),
88            Short { .. } => Some(EncryptionLevel::OneRtt),
89            _ => None,
90        }
91    }
92
93    pub(crate) fn is_0rtt(&self) -> bool {
94        match self.plain_header {
95            ProtectedHeader::Long { ty, .. } => ty == LongType::ZeroRtt,
96            _ => false,
97        }
98    }
99
100    /// The destination connection ID of the packet
101    pub fn dst_cid(&self) -> ConnectionId {
102        self.plain_header.dst_cid()
103    }
104
105    /// Length of QUIC packet being decoded
106    #[allow(unreachable_pub)] // fuzzing only
107    pub fn len(&self) -> usize {
108        self.buf.get_ref().len()
109    }
110
111    pub(crate) fn finish(
112        self,
113        header_crypto: Option<&dyn crypto::HeaderKey>,
114    ) -> Result<Packet, PacketDecodeError> {
115        use ProtectedHeader::*;
116        let Self {
117            plain_header,
118            mut buf,
119        } = self;
120
121        if let Initial(ProtectedInitialHeader {
122            dst_cid,
123            src_cid,
124            token_pos,
125            version,
126            ..
127        }) = plain_header
128        {
129            let number = Self::decrypt_header(&mut buf, header_crypto.unwrap())?;
130            let header_len = buf.position() as usize;
131            let mut bytes = buf.into_inner();
132
133            let header_data = bytes.split_to(header_len).freeze();
134            let token = header_data.slice(token_pos.start..token_pos.end);
135            return Ok(Packet {
136                header: Header::Initial(InitialHeader {
137                    dst_cid,
138                    src_cid,
139                    token,
140                    number,
141                    version,
142                }),
143                header_data,
144                payload: bytes,
145            });
146        }
147
148        let header = match plain_header {
149            Long {
150                ty,
151                dst_cid,
152                src_cid,
153                version,
154                ..
155            } => Header::Long {
156                ty,
157                dst_cid,
158                src_cid,
159                number: Self::decrypt_header(&mut buf, header_crypto.unwrap())?,
160                version,
161            },
162            Retry {
163                dst_cid,
164                src_cid,
165                version,
166            } => Header::Retry {
167                dst_cid,
168                src_cid,
169                version,
170            },
171            Short { spin, dst_cid, .. } => {
172                let number = Self::decrypt_header(&mut buf, header_crypto.unwrap())?;
173                let key_phase = buf.get_ref()[0] & KEY_PHASE_BIT != 0;
174                Header::Short {
175                    spin,
176                    key_phase,
177                    dst_cid,
178                    number,
179                }
180            }
181            VersionNegotiate {
182                random,
183                dst_cid,
184                src_cid,
185            } => Header::VersionNegotiate {
186                random,
187                dst_cid,
188                src_cid,
189            },
190            Initial { .. } => unreachable!(),
191        };
192
193        let header_len = buf.position() as usize;
194        let mut bytes = buf.into_inner();
195        Ok(Packet {
196            header,
197            header_data: bytes.split_to(header_len).freeze(),
198            payload: bytes,
199        })
200    }
201
202    fn decrypt_header(
203        buf: &mut io::Cursor<BytesMut>,
204        header_crypto: &dyn crypto::HeaderKey,
205    ) -> Result<PacketNumber, PacketDecodeError> {
206        let packet_length = buf.get_ref().len();
207        let pn_offset = buf.position() as usize;
208        if packet_length < pn_offset + 4 + header_crypto.sample_size() {
209            return Err(PacketDecodeError::InvalidHeader(
210                "packet too short to extract header protection sample",
211            ));
212        }
213
214        header_crypto.decrypt(pn_offset, buf.get_mut());
215
216        let len = PacketNumber::decode_len(buf.get_ref()[0]);
217        PacketNumber::decode(len, buf)
218    }
219}
220
221/// A buffer that can tell how much has been written to it already
222///
223/// This is commonly used for when a buffer is passed and the user may not write past a
224/// given size. It allows the user of such a buffer to know the current cursor position in
225/// the buffer. The maximum write size is usually passed in the same unit as
226/// [`BufLen::len`]: bytes since the buffer start.
227pub(crate) trait BufLen {
228    /// Returns the number of bytes written into the buffer so far
229    fn len(&self) -> usize;
230}
231
232impl BufLen for Vec<u8> {
233    fn len(&self) -> usize {
234        self.len()
235    }
236}
237
238/// A received packet with header protection removed.
239// TODO(flub): Would be grand to make this typesafe by adding a generic that indicates
240//    whether the payload is encrypted or decrypted.
241pub(crate) struct Packet {
242    /// The decoded header.
243    pub(crate) header: Header,
244    /// The encoded header data, with header protection removed (i.e. decrypted).
245    pub(crate) header_data: Bytes,
246    /// Packet payload, still encrypted when created, later decrypted in-place.
247    pub(crate) payload: BytesMut,
248}
249
250impl Packet {
251    pub(crate) fn reserved_bits_valid(&self) -> bool {
252        let mask = match self.header {
253            Header::Short { .. } => SHORT_RESERVED_BITS,
254            _ => LONG_RESERVED_BITS,
255        };
256        self.header_data[0] & mask == 0
257    }
258}
259
260pub(crate) struct InitialPacket {
261    pub(crate) header: InitialHeader,
262    pub(crate) header_data: Bytes,
263    pub(crate) payload: BytesMut,
264}
265
266impl From<InitialPacket> for Packet {
267    fn from(x: InitialPacket) -> Self {
268        Self {
269            header: Header::Initial(x.header),
270            header_data: x.header_data,
271            payload: x.payload,
272        }
273    }
274}
275
276#[cfg_attr(test, derive(Clone))]
277#[derive(Debug)]
278pub(crate) enum Header {
279    Initial(InitialHeader),
280    Long {
281        ty: LongType,
282        dst_cid: ConnectionId,
283        src_cid: ConnectionId,
284        number: PacketNumber,
285        version: u32,
286    },
287    Retry {
288        dst_cid: ConnectionId,
289        src_cid: ConnectionId,
290        version: u32,
291    },
292    Short {
293        spin: bool,
294        key_phase: bool,
295        dst_cid: ConnectionId,
296        number: PacketNumber,
297    },
298    VersionNegotiate {
299        random: u8,
300        src_cid: ConnectionId,
301        dst_cid: ConnectionId,
302    },
303}
304
305impl Header {
306    pub(crate) fn encode(&self, w: &mut (impl BufMut + BufLen)) -> PartialEncode {
307        use Header::*;
308        let start = w.len();
309        match *self {
310            Initial(InitialHeader {
311                ref dst_cid,
312                ref src_cid,
313                ref token,
314                number,
315                version,
316            }) => {
317                w.write(u8::from(LongHeaderType::Initial) | number.tag());
318                w.write(version);
319                dst_cid.encode_long(w);
320                src_cid.encode_long(w);
321                w.write_var(token.len() as u64);
322                w.put_slice(token);
323                w.write::<u16>(0); // Placeholder for payload length; see `set_payload_length`
324                number.encode(w);
325                PartialEncode {
326                    start,
327                    header_len: w.len() - start,
328                    pn: Some((number.len(), true)),
329                }
330            }
331            Long {
332                ty,
333                ref dst_cid,
334                ref src_cid,
335                number,
336                version,
337            } => {
338                w.write(u8::from(LongHeaderType::Standard(ty)) | number.tag());
339                w.write(version);
340                dst_cid.encode_long(w);
341                src_cid.encode_long(w);
342                w.write::<u16>(0); // Placeholder for payload length; see `set_payload_length`
343                number.encode(w);
344                PartialEncode {
345                    start,
346                    header_len: w.len() - start,
347                    pn: Some((number.len(), true)),
348                }
349            }
350            Retry {
351                ref dst_cid,
352                ref src_cid,
353                version,
354            } => {
355                w.write(u8::from(LongHeaderType::Retry));
356                w.write(version);
357                dst_cid.encode_long(w);
358                src_cid.encode_long(w);
359                PartialEncode {
360                    start,
361                    header_len: w.len() - start,
362                    pn: None,
363                }
364            }
365            Short {
366                spin,
367                key_phase,
368                ref dst_cid,
369                number,
370            } => {
371                w.write(
372                    FIXED_BIT
373                        | if key_phase { KEY_PHASE_BIT } else { 0 }
374                        | if spin { SPIN_BIT } else { 0 }
375                        | number.tag(),
376                );
377                w.put_slice(dst_cid);
378                number.encode(w);
379                PartialEncode {
380                    start,
381                    header_len: w.len() - start,
382                    pn: Some((number.len(), false)),
383                }
384            }
385            VersionNegotiate {
386                ref random,
387                ref dst_cid,
388                ref src_cid,
389            } => {
390                w.write(0x80u8 | random);
391                w.write::<u32>(0);
392                dst_cid.encode_long(w);
393                src_cid.encode_long(w);
394                PartialEncode {
395                    start,
396                    header_len: w.len() - start,
397                    pn: None,
398                }
399            }
400        }
401    }
402
403    /// Whether the packet is encrypted on the wire
404    pub(crate) fn is_protected(&self) -> bool {
405        !matches!(*self, Self::Retry { .. } | Self::VersionNegotiate { .. })
406    }
407
408    pub(crate) fn number(&self) -> Option<PacketNumber> {
409        use Header::*;
410        Some(match *self {
411            Initial(InitialHeader { number, .. }) => number,
412            Long { number, .. } => number,
413            Short { number, .. } => number,
414            _ => {
415                return None;
416            }
417        })
418    }
419
420    pub(crate) fn space(&self) -> SpaceKind {
421        use Header::*;
422        match *self {
423            Short { .. } => SpaceKind::Data,
424            Long {
425                ty: LongType::ZeroRtt,
426                ..
427            } => SpaceKind::Data,
428            Long {
429                ty: LongType::Handshake,
430                ..
431            } => SpaceKind::Handshake,
432            _ => SpaceKind::Initial,
433        }
434    }
435
436    pub(crate) fn key_phase(&self) -> bool {
437        match *self {
438            Self::Short { key_phase, .. } => key_phase,
439            _ => false,
440        }
441    }
442
443    pub(crate) fn is_short(&self) -> bool {
444        matches!(*self, Self::Short { .. })
445    }
446
447    pub(crate) fn is_1rtt(&self) -> bool {
448        self.is_short()
449    }
450
451    pub(crate) fn is_0rtt(&self) -> bool {
452        matches!(
453            *self,
454            Self::Long {
455                ty: LongType::ZeroRtt,
456                ..
457            }
458        )
459    }
460
461    pub(crate) fn dst_cid(&self) -> ConnectionId {
462        use Header::*;
463        match *self {
464            Initial(InitialHeader { dst_cid, .. }) => dst_cid,
465            Long { dst_cid, .. } => dst_cid,
466            Retry { dst_cid, .. } => dst_cid,
467            Short { dst_cid, .. } => dst_cid,
468            VersionNegotiate { dst_cid, .. } => dst_cid,
469        }
470    }
471
472    /// Is this packet allowed to be coalesced with others?
473    ///
474    /// Ref <https://www.rfc-editor.org/rfc/rfc9000.html#name-coalescing-packets>
475    pub(crate) fn can_coalesce(&self) -> bool {
476        use Header::*;
477        match *self {
478            Initial(_) => true,
479            Long { .. } => true,
480            Retry { .. } => false,
481            Short { .. } => false,
482            VersionNegotiate { .. } => false,
483        }
484    }
485
486    /// Whether the payload of this packet contains QUIC frames
487    pub(crate) fn has_frames(&self) -> bool {
488        use Header::*;
489        match *self {
490            Initial(_) => true,
491            Long { .. } => true,
492            Retry { .. } => false,
493            Short { .. } => true,
494            VersionNegotiate { .. } => false,
495        }
496    }
497
498    #[cfg(feature = "qlog")]
499    pub(crate) fn src_cid(&self) -> Option<ConnectionId> {
500        match self {
501            Self::Initial(initial_header) => Some(initial_header.src_cid),
502            Self::Long { src_cid, .. } => Some(*src_cid),
503            Self::Retry { src_cid, .. } => Some(*src_cid),
504            Self::Short { .. } => None,
505            Self::VersionNegotiate { src_cid, .. } => Some(*src_cid),
506        }
507    }
508}
509
510pub(crate) struct PartialEncode {
511    pub(crate) start: usize,
512    pub(crate) header_len: usize,
513    // Packet number length, payload length needed
514    pn: Option<(usize, bool)>,
515}
516
517impl PartialEncode {
518    pub(crate) fn finish(
519        self,
520        buf: &mut [u8],
521        header_crypto: &dyn crypto::HeaderKey,
522        crypto: Option<(u64, PathId, &dyn crypto::PacketKey)>,
523    ) {
524        let Self { header_len, pn, .. } = self;
525        let Some((pn_len, write_len)) = pn else {
526            return;
527        };
528
529        let pn_pos = header_len - pn_len;
530        if write_len {
531            let len = buf.len() - header_len + pn_len;
532            assert!(len < 2usize.pow(14)); // Fits in reserved space
533            let mut slice = &mut buf[pn_pos - 2..pn_pos];
534            slice.put_u16(len as u16 | (0b01 << 14));
535        }
536
537        if let Some((packet_number, path_id, crypto)) = crypto {
538            crypto.encrypt(path_id, packet_number, buf, header_len);
539        }
540
541        debug_assert!(
542            pn_pos + 4 + header_crypto.sample_size() <= buf.len(),
543            "packet must be padded to at least {} bytes for header protection sampling",
544            pn_pos + 4 + header_crypto.sample_size()
545        );
546        header_crypto.encrypt(pn_pos, buf);
547    }
548
549    /// Creates a [`PartialEncode`] that has not encoded a header into the buffer.
550    ///
551    /// This is used exclusively for testing as such a type is otherwise invalid.
552    #[cfg(test)]
553    pub(crate) fn no_header() -> Self {
554        Self {
555            start: 0,
556            header_len: 0,
557            pn: None,
558        }
559    }
560}
561
562/// Plain packet header
563#[derive(Clone, Debug)]
564pub enum ProtectedHeader {
565    /// An Initial packet header
566    Initial(ProtectedInitialHeader),
567    /// A Long packet header, as used during the handshake
568    Long {
569        /// Type of the Long header packet
570        ty: LongType,
571        /// Destination Connection ID
572        dst_cid: ConnectionId,
573        /// Source Connection ID
574        src_cid: ConnectionId,
575        /// Length of the packet payload
576        len: u64,
577        /// QUIC version
578        version: u32,
579    },
580    /// A Retry packet header
581    Retry {
582        /// Destination Connection ID
583        dst_cid: ConnectionId,
584        /// Source Connection ID
585        src_cid: ConnectionId,
586        /// QUIC version
587        version: u32,
588    },
589    /// A short packet header, as used during the data phase
590    Short {
591        /// Spin bit
592        spin: bool,
593        /// Destination Connection ID
594        dst_cid: ConnectionId,
595    },
596    /// A Version Negotiation packet header
597    VersionNegotiate {
598        /// Random value
599        random: u8,
600        /// Destination Connection ID
601        dst_cid: ConnectionId,
602        /// Source Connection ID
603        src_cid: ConnectionId,
604    },
605}
606
607impl ProtectedHeader {
608    fn as_initial(&self) -> Option<&ProtectedInitialHeader> {
609        match self {
610            Self::Initial(x) => Some(x),
611            _ => None,
612        }
613    }
614
615    /// The destination Connection ID of the packet
616    pub fn dst_cid(&self) -> ConnectionId {
617        use ProtectedHeader::*;
618        match self {
619            Initial(header) => header.dst_cid,
620            &Long { dst_cid, .. } => dst_cid,
621            &Retry { dst_cid, .. } => dst_cid,
622            &Short { dst_cid, .. } => dst_cid,
623            &VersionNegotiate { dst_cid, .. } => dst_cid,
624        }
625    }
626
627    fn payload_len(&self) -> Option<u64> {
628        use ProtectedHeader::*;
629        match self {
630            Initial(ProtectedInitialHeader { len, .. }) | Long { len, .. } => Some(*len),
631            _ => None,
632        }
633    }
634
635    /// Decode a plain header from given buffer, with given [`ConnectionIdParser`].
636    pub fn decode(
637        buf: &mut io::Cursor<impl AsRef<[u8]>>,
638        cid_parser: &(impl ConnectionIdParser + ?Sized),
639        supported_versions: &[u32],
640        grease_quic_bit: bool,
641    ) -> Result<Self, PacketDecodeError> {
642        let first = buf.get::<u8>()?;
643        if !grease_quic_bit && first & FIXED_BIT == 0 {
644            return Err(PacketDecodeError::InvalidHeader("fixed bit unset"));
645        }
646        if first & LONG_HEADER_FORM == 0 {
647            let spin = first & SPIN_BIT != 0;
648
649            Ok(Self::Short {
650                spin,
651                dst_cid: cid_parser.parse(buf)?,
652            })
653        } else {
654            let version = buf.get::<u32>()?;
655
656            let dst_cid = ConnectionId::decode_long(buf)
657                .ok_or(PacketDecodeError::InvalidHeader("malformed cid"))?;
658            let src_cid = ConnectionId::decode_long(buf)
659                .ok_or(PacketDecodeError::InvalidHeader("malformed cid"))?;
660
661            // TODO: Support long CIDs for compatibility with future QUIC versions
662            if version == 0 {
663                let random = first & !LONG_HEADER_FORM;
664                return Ok(Self::VersionNegotiate {
665                    random,
666                    dst_cid,
667                    src_cid,
668                });
669            }
670
671            if !supported_versions.contains(&version) {
672                return Err(PacketDecodeError::UnsupportedVersion {
673                    src_cid,
674                    dst_cid,
675                    version,
676                });
677            }
678
679            match LongHeaderType::from_byte(first)? {
680                LongHeaderType::Initial => {
681                    let token_len = buf.get_var()? as usize;
682                    let token_start = buf.position() as usize;
683                    if token_len > buf.remaining() {
684                        return Err(PacketDecodeError::InvalidHeader("token out of bounds"));
685                    }
686                    buf.advance(token_len);
687
688                    let len = buf.get_var()?;
689                    Ok(Self::Initial(ProtectedInitialHeader {
690                        dst_cid,
691                        src_cid,
692                        token_pos: token_start..token_start + token_len,
693                        len,
694                        version,
695                    }))
696                }
697                LongHeaderType::Retry => Ok(Self::Retry {
698                    dst_cid,
699                    src_cid,
700                    version,
701                }),
702                LongHeaderType::Standard(ty) => Ok(Self::Long {
703                    ty,
704                    dst_cid,
705                    src_cid,
706                    len: buf.get_var()?,
707                    version,
708                }),
709            }
710        }
711    }
712}
713
714/// Header of an Initial packet, before decryption
715#[derive(Clone, Debug)]
716pub struct ProtectedInitialHeader {
717    /// Destination Connection ID
718    pub dst_cid: ConnectionId,
719    /// Source Connection ID
720    pub src_cid: ConnectionId,
721    /// The position of a token in the packet buffer
722    pub token_pos: Range<usize>,
723    /// Length of the packet payload
724    pub len: u64,
725    /// QUIC version
726    pub version: u32,
727}
728
729#[derive(Clone, Debug)]
730pub(crate) struct InitialHeader {
731    pub(crate) dst_cid: ConnectionId,
732    pub(crate) src_cid: ConnectionId,
733    pub(crate) token: Bytes,
734    pub(crate) number: PacketNumber,
735    pub(crate) version: u32,
736}
737
738// An encoded packet number
739#[derive(Debug, Copy, Clone, Eq, PartialEq)]
740pub(crate) enum PacketNumber {
741    U8(u8),
742    U16(u16),
743    U24(u32),
744    U32(u32),
745}
746
747impl PacketNumber {
748    pub(crate) fn new(n: u64, largest_acked: u64) -> Self {
749        let range = (n - largest_acked) * 2;
750        if range < 1 << 8 {
751            Self::U8(n as u8)
752        } else if range < 1 << 16 {
753            Self::U16(n as u16)
754        } else if range < 1 << 24 {
755            Self::U24(n as u32)
756        } else if range < 1 << 32 {
757            Self::U32(n as u32)
758        } else {
759            panic!("packet number too large to encode")
760        }
761    }
762
763    pub(crate) fn len(self) -> usize {
764        use PacketNumber::*;
765        match self {
766            U8(_) => 1,
767            U16(_) => 2,
768            U24(_) => 3,
769            U32(_) => 4,
770        }
771    }
772
773    pub(crate) fn encode<W: BufMut>(self, w: &mut W) {
774        use PacketNumber::*;
775        match self {
776            U8(x) => w.write(x),
777            U16(x) => w.write(x),
778            U24(x) => w.put_uint(u64::from(x), 3),
779            U32(x) => w.write(x),
780        }
781    }
782
783    pub(crate) fn decode<R: Buf>(len: usize, r: &mut R) -> Result<Self, PacketDecodeError> {
784        use PacketNumber::*;
785        let pn = match len {
786            1 => U8(r.get()?),
787            2 => U16(r.get()?),
788            3 => U24(r.get_uint(3) as u32),
789            4 => U32(r.get()?),
790            _ => unreachable!(),
791        };
792        Ok(pn)
793    }
794
795    pub(crate) fn decode_len(tag: u8) -> usize {
796        1 + (tag & 0x03) as usize
797    }
798
799    fn tag(self) -> u8 {
800        use PacketNumber::*;
801        match self {
802            U8(_) => 0b00,
803            U16(_) => 0b01,
804            U24(_) => 0b10,
805            U32(_) => 0b11,
806        }
807    }
808
809    pub(crate) fn expand(self, expected: u64) -> u64 {
810        // From Appendix A
811        use PacketNumber::*;
812        let truncated = match self {
813            U8(x) => u64::from(x),
814            U16(x) => u64::from(x),
815            U24(x) => u64::from(x),
816            U32(x) => u64::from(x),
817        };
818        let nbits = self.len() * 8;
819        let win = 1 << nbits;
820        let hwin = win / 2;
821        let mask = win - 1;
822        // The incoming packet number should be greater than expected - hwin and less than or equal
823        // to expected + hwin
824        //
825        // This means we can't just strip the trailing bits from expected and add the truncated
826        // because that might yield a value outside the window.
827        //
828        // The following code calculates a candidate value and makes sure it's within the packet
829        // number window.
830        let candidate = (expected & !mask) | truncated;
831        if expected.checked_sub(hwin).is_some_and(|x| candidate <= x) {
832            candidate + win
833        } else if candidate > expected + hwin && candidate > win {
834            candidate - win
835        } else {
836            candidate
837        }
838    }
839}
840
841/// A [`ConnectionIdParser`] implementation that assumes the connection ID is of fixed length
842pub struct FixedLengthConnectionIdParser {
843    expected_len: usize,
844}
845
846impl FixedLengthConnectionIdParser {
847    /// Create a new instance of `FixedLengthConnectionIdParser`
848    pub fn new(expected_len: usize) -> Self {
849        Self { expected_len }
850    }
851}
852
853impl ConnectionIdParser for FixedLengthConnectionIdParser {
854    fn parse(&self, buffer: &mut dyn Buf) -> Result<ConnectionId, PacketDecodeError> {
855        (buffer.remaining() >= self.expected_len)
856            .then(|| ConnectionId::from_buf(buffer, self.expected_len))
857            .ok_or(PacketDecodeError::InvalidHeader("packet too small"))
858    }
859}
860
861/// Parse connection id in short header packet
862pub trait ConnectionIdParser {
863    /// Parse a connection id from given buffer
864    fn parse(&self, buf: &mut dyn Buf) -> Result<ConnectionId, PacketDecodeError>;
865}
866
867/// Long packet type including non-uniform cases
868#[derive(Clone, Copy, Debug, Eq, PartialEq)]
869pub(crate) enum LongHeaderType {
870    Initial,
871    Retry,
872    Standard(LongType),
873}
874
875impl LongHeaderType {
876    fn from_byte(b: u8) -> Result<Self, PacketDecodeError> {
877        use {LongHeaderType::*, LongType::*};
878        debug_assert!(b & LONG_HEADER_FORM != 0, "not a long packet");
879        Ok(match (b & 0x30) >> 4 {
880            0x0 => Initial,
881            0x1 => Standard(ZeroRtt),
882            0x2 => Standard(Handshake),
883            0x3 => Retry,
884            _ => unreachable!(),
885        })
886    }
887}
888
889impl From<LongHeaderType> for u8 {
890    fn from(ty: LongHeaderType) -> Self {
891        use {LongHeaderType::*, LongType::*};
892        match ty {
893            Initial => LONG_HEADER_FORM | FIXED_BIT,
894            Standard(ZeroRtt) => LONG_HEADER_FORM | FIXED_BIT | (0x1 << 4),
895            Standard(Handshake) => LONG_HEADER_FORM | FIXED_BIT | (0x2 << 4),
896            Retry => LONG_HEADER_FORM | FIXED_BIT | (0x3 << 4),
897        }
898    }
899}
900
901/// Long packet types with uniform header structure
902#[derive(Clone, Copy, Debug, Eq, PartialEq)]
903pub enum LongType {
904    /// Handshake packet
905    Handshake,
906    /// 0-RTT packet
907    ZeroRtt,
908}
909
910/// Packet decode error
911#[derive(Debug, Error, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
912pub enum PacketDecodeError {
913    /// Packet uses a QUIC version that is not supported
914    #[error("unsupported version {version:x}")]
915    UnsupportedVersion {
916        /// Source Connection ID
917        src_cid: ConnectionId,
918        /// Destination Connection ID
919        dst_cid: ConnectionId,
920        /// The version that was unsupported
921        version: u32,
922    },
923    /// The packet header is invalid
924    #[error("invalid header: {0}")]
925    InvalidHeader(&'static str),
926}
927
928impl From<coding::UnexpectedEnd> for PacketDecodeError {
929    fn from(_: coding::UnexpectedEnd) -> Self {
930        Self::InvalidHeader("unexpected end of packet")
931    }
932}
933
934pub(crate) const LONG_HEADER_FORM: u8 = 0x80;
935pub(crate) const FIXED_BIT: u8 = 0x40;
936pub(crate) const SPIN_BIT: u8 = 0x20;
937const SHORT_RESERVED_BITS: u8 = 0x18;
938const LONG_RESERVED_BITS: u8 = 0x0c;
939const KEY_PHASE_BIT: u8 = 0x04;
940
941/// Packet number space identifiers
942#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
943pub(crate) enum SpaceId {
944    /// Unprotected packets, used to bootstrap the handshake
945    Initial = 0,
946    Handshake = 1,
947    /// Application data space, used for 0-RTT and post-handshake/1-RTT packets
948    Data = 2,
949}
950
951impl SpaceId {
952    pub(crate) fn iter() -> impl Iterator<Item = Self> {
953        [Self::Initial, Self::Handshake, Self::Data].iter().cloned()
954    }
955
956    /// Returns the next higher packet space.
957    ///
958    /// Returns `None` if at  [`SpaceId::Data`].
959    pub(crate) fn next(&self) -> Option<Self> {
960        match self {
961            Self::Initial => Some(Self::Handshake),
962            Self::Handshake => Some(Self::Data),
963            Self::Data => None,
964        }
965    }
966
967    /// Returns the encryption level for this packet space.
968    pub(crate) fn encryption_level(self) -> EncryptionLevel {
969        match self {
970            Self::Initial => EncryptionLevel::Initial,
971            Self::Handshake => EncryptionLevel::Handshake,
972            Self::Data => EncryptionLevel::OneRtt,
973        }
974    }
975
976    /// Returns the [`SpaceKind`] for this packet space.
977    pub(crate) fn kind(self) -> SpaceKind {
978        match self {
979            Self::Initial => SpaceKind::Initial,
980            Self::Handshake => SpaceKind::Handshake,
981            Self::Data => SpaceKind::Data,
982        }
983    }
984}
985
986#[cfg(test)]
987mod tests {
988    use super::*;
989    use hex_literal::hex;
990    use std::io;
991
992    fn check_pn(typed: PacketNumber, encoded: &[u8]) {
993        let mut buf = Vec::new();
994        typed.encode(&mut buf);
995        assert_eq!(&buf[..], encoded);
996        let decoded = PacketNumber::decode(typed.len(), &mut io::Cursor::new(&buf)).unwrap();
997        assert_eq!(typed, decoded);
998    }
999
1000    #[test]
1001    fn roundtrip_packet_numbers() {
1002        check_pn(PacketNumber::U8(0x7f), &hex!("7f"));
1003        check_pn(PacketNumber::U16(0x80), &hex!("0080"));
1004        check_pn(PacketNumber::U16(0x3fff), &hex!("3fff"));
1005        check_pn(PacketNumber::U32(0x0000_4000), &hex!("0000 4000"));
1006        check_pn(PacketNumber::U32(0xffff_ffff), &hex!("ffff ffff"));
1007    }
1008
1009    #[test]
1010    fn pn_encode() {
1011        check_pn(PacketNumber::new(0x10, 0), &hex!("10"));
1012        check_pn(PacketNumber::new(0x100, 0), &hex!("0100"));
1013        check_pn(PacketNumber::new(0x10000, 0), &hex!("010000"));
1014    }
1015
1016    #[test]
1017    fn pn_expand_roundtrip() {
1018        for expected in 0..1024 {
1019            for actual in expected..1024 {
1020                assert_eq!(actual, PacketNumber::new(actual, expected).expand(expected));
1021            }
1022        }
1023    }
1024
1025    #[cfg(all(feature = "rustls", any(feature = "aws-lc-rs", feature = "ring")))]
1026    #[test]
1027    fn header_encoding() {
1028        use crate::Side;
1029        use crate::crypto::rustls::{configured_provider, initial_keys, initial_suite_from_provider};
1030        use rustls::quic::Version;
1031
1032        let dcid = ConnectionId::new(&hex!("06b858ec6f80452b"));
1033        let provider = configured_provider();
1034
1035        let suite = initial_suite_from_provider(&provider).unwrap();
1036        let client = initial_keys(Version::V1, dcid, Side::Client, &suite);
1037        let mut buf = Vec::new();
1038        let header = Header::Initial(InitialHeader {
1039            number: PacketNumber::U8(0),
1040            src_cid: ConnectionId::new(&[]),
1041            dst_cid: dcid,
1042            token: Bytes::new(),
1043            version: crate::DEFAULT_SUPPORTED_VERSIONS[0],
1044        });
1045        let encode = header.encode(&mut buf);
1046        let header_len = buf.len();
1047        buf.resize(header_len + 16 + client.packet.local.tag_len(), 0);
1048        encode.finish(
1049            &mut buf,
1050            &*client.header.local,
1051            Some((0, PathId::ZERO, &*client.packet.local)),
1052        );
1053
1054        for byte in &buf {
1055            print!("{byte:02x}");
1056        }
1057        println!();
1058        assert_eq!(
1059            buf[..],
1060            hex!(
1061                "c8000000010806b858ec6f80452b00004021be
1062                 3ef50807b84191a196f760a6dad1e9d1c430c48952cba0148250c21c0a6a70e1"
1063            )[..]
1064        );
1065
1066        let server = initial_keys(Version::V1, dcid, Side::Server, &suite);
1067        let supported_versions = crate::DEFAULT_SUPPORTED_VERSIONS.to_vec();
1068        let decode = PartialDecode::new(
1069            buf.as_slice().into(),
1070            &FixedLengthConnectionIdParser::new(0),
1071            &supported_versions,
1072            false,
1073        )
1074        .unwrap()
1075        .0;
1076        let mut packet = decode.finish(Some(&*server.header.remote)).unwrap();
1077        assert_eq!(
1078            packet.header_data[..],
1079            hex!("c0000000010806b858ec6f80452b0000402100")[..]
1080        );
1081        server
1082            .packet
1083            .remote
1084            .decrypt(PathId::ZERO, 0, &packet.header_data, &mut packet.payload)
1085            .unwrap();
1086        assert_eq!(packet.payload[..], [0; 16]);
1087        match packet.header {
1088            Header::Initial(InitialHeader {
1089                number: PacketNumber::U8(0),
1090                ..
1091            }) => {}
1092            _ => {
1093                panic!("unexpected header {:?}", packet.header);
1094            }
1095        }
1096    }
1097}