qbase/
frame.rs

1use std::fmt::Debug;
2
3use bytes::{Buf, BufMut, Bytes};
4use derive_more::{Deref, DerefMut};
5use enum_dispatch::enum_dispatch;
6use io::WriteFrame;
7
8use super::varint::VarInt;
9use crate::packet::r#type::Type;
10
11mod ack;
12mod connection_close;
13mod crypto;
14mod data_blocked;
15mod datagram;
16mod handshake_done;
17mod max_data;
18mod max_stream_data;
19mod max_streams;
20mod new_connection_id;
21mod new_token;
22mod padding;
23pub mod path_challenge;
24mod path_response;
25mod ping;
26mod reset_stream;
27mod retire_connection_id;
28mod stop_sending;
29mod stream;
30mod stream_data_blocked;
31mod streams_blocked;
32
33/// Error module for parsing frames
34pub mod error;
35/// IO module for frame encoding and decoding
36pub mod io;
37
38pub use ack::{AckFrame, EcnCounts};
39pub use connection_close::{AppCloseFrame, ConnectionCloseFrame, QuicCloseFrame};
40pub use crypto::CryptoFrame;
41pub use data_blocked::DataBlockedFrame;
42pub use datagram::DatagramFrame;
43#[doc(hidden)]
44pub use error::Error;
45pub use handshake_done::HandshakeDoneFrame;
46pub use max_data::MaxDataFrame;
47pub use max_stream_data::MaxStreamDataFrame;
48pub use max_streams::MaxStreamsFrame;
49pub use new_connection_id::NewConnectionIdFrame;
50pub use new_token::NewTokenFrame;
51pub use padding::PaddingFrame;
52pub use path_challenge::PathChallengeFrame;
53pub use path_response::PathResponseFrame;
54pub use ping::PingFrame;
55pub use reset_stream::{ResetStreamError, ResetStreamFrame};
56pub use retire_connection_id::RetireConnectionIdFrame;
57pub use stop_sending::StopSendingFrame;
58pub use stream::{EncodingStrategy, STREAM_FRAME_MAX_ENCODING_SIZE, StreamFrame};
59pub use stream_data_blocked::StreamDataBlockedFrame;
60pub use streams_blocked::StreamsBlockedFrame;
61
62/// Define the basic behaviors for all kinds of frames
63#[enum_dispatch]
64pub trait GetFrameType: Debug {
65    /// Return the type of frame
66    fn frame_type(&self) -> FrameType;
67}
68
69#[enum_dispatch]
70pub trait EncodeFrame {
71    /// Return the max number of bytes needed to encode this value
72    ///
73    /// Calculate the maximum size by summing up the maximum length of each field.
74    /// If a field type has a maximum length, use it, otherwise use the actual length
75    /// of the data in that field.
76    ///
77    /// When packaging data, by pre-estimating this value to effectively avoid spending
78    /// extra resources to calculate the actual encoded size.
79    fn max_encoding_size(&self) -> usize {
80        1
81    }
82
83    /// Return the exact number of bytes needed to encode this value
84    fn encoding_size(&self) -> usize {
85        1
86    }
87}
88
89/// The `Spec` summarizes any special rules governing the processing
90/// or generation of the frame type, as indicated by the following characters.
91///
92/// See [table-3](https://www.rfc-editor.org/rfc/rfc9000.html#table-3)
93/// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html) for more details.
94pub enum Spec {
95    /// Packets containing only frames with this marking are not ack-eliciting.
96    ///
97    /// See [Section 13.2](https://www.rfc-editor.org/rfc/rfc9000.html#generating-acks)
98    /// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html) for more details.
99    NonAckEliciting = 1,
100    /// Packets containing only frames with this marking do not count toward bytes
101    /// in flight for congestion control purposes.
102    /// See [section-12.4-14.4](https://www.rfc-editor.org/rfc/rfc9000.html#section-12.4-14.4)
103    /// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html).
104    ///
105    /// Similar to TCP, packets containing only ACK frames do not count toward bytes
106    /// in flight and are not congestion controlled.
107    /// See [Section 7.4](https://www.rfc-editor.org/rfc/rfc9002#section-7-4)
108    /// of [QUIC-RECOVERY](https://www.rfc-editor.org/rfc/rfc9002).
109    CongestionControlFree = 2,
110    /// Packets containing only frames with this marking can be used to probe
111    /// new network paths during connection migration.
112    ///
113    /// See [Section 9.1](https://www.rfc-editor.org/rfc/rfc9000.html#probing)
114    /// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html).
115    ProbeNewPath = 4,
116    /// The contents of frames with this marking are flow controlled.
117    ///
118    /// See [Section 4](https://www.rfc-editor.org/rfc/rfc9000.html#flow-control)
119    /// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html) for more details.
120    FlowControlled = 8,
121}
122
123pub trait ContainSpec {
124    fn contain(&self, spec: Spec) -> bool;
125}
126
127impl ContainSpec for u8 {
128    #[inline]
129    fn contain(&self, spec: Spec) -> bool {
130        *self & spec as u8 != 0
131    }
132}
133
134/// The sum type of all the core QUIC frame types.
135///
136/// See [table-3](https://www.rfc-editor.org/rfc/rfc9000.html#table-3)
137/// and [frame types and formats](https://www.rfc-editor.org/rfc/rfc9000.html#name-frame-types-and-formats)
138/// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html) for more details.
139#[derive(Debug, Copy, Clone, Eq, PartialEq)]
140pub enum FrameType {
141    /// PADDED frame, see [`PaddingFrame`].
142    Padding,
143    /// PING frame, see [`PingFrame`].
144    Ping,
145    /// ACK frame, see [`AckFrame`].
146    Ack(u8),
147    /// RESET_STREAM frame, see [`ResetStreamFrame`].
148    ResetStream,
149    /// STOP_SENDING frame, see [`StopSendingFrame`].
150    StopSending,
151    /// CRYPTO frame, see [`CryptoFrame`].
152    Crypto,
153    /// NEW_TOKEN frame, see [`NewTokenFrame`].
154    NewToken,
155    /// STREAM frame, see [`StreamFrame`].
156    Stream(u8),
157    /// MAX_DATA frame, see [`MaxDataFrame`].
158    MaxData,
159    /// MAX_STREAM_DATA frame, see [`MaxStreamDataFrame`].
160    MaxStreamData,
161    /// MAX_STREAMS frame, see [`MaxStreamsFrame`].
162    MaxStreams(u8),
163    /// DATA_BLOCKED frame, see [`DataBlockedFrame`].
164    DataBlocked,
165    /// STREAM_DATA_BLOCKED frame, see [`StreamDataBlockedFrame`].
166    StreamDataBlocked,
167    /// STREAMS_BLOCKED frame, see [`StreamsBlockedFrame`].
168    StreamsBlocked(u8),
169    /// NEW_CONNECTION_ID frame, see [`NewConnectionIdFrame`].
170    NewConnectionId,
171    /// RETIRE_CONNECTION_ID frame, see [`RetireConnectionIdFrame`].
172    RetireConnectionId,
173    /// PATH_CHALLENGE frame, see [`PathChallengeFrame`].
174    PathChallenge,
175    /// PATH_RESPONSE frame, see [`PathResponseFrame`].
176    PathResponse,
177    /// CONNECTION_CLOSE frame, see [`ConnectionCloseFrame`].
178    ConnectionClose(u8),
179    /// HANDSHAKE_DONE frame, see [`HandshakeDoneFrame`].
180    HandshakeDone,
181    /// DATAGRAM frame, see [`DatagramFrame`].
182    Datagram(u8),
183}
184
185#[enum_dispatch]
186pub trait FrameFeture {
187    /// Return whether a frame type belongs to the given packet_type
188    fn belongs_to(&self, packet_type: Type) -> bool;
189    /// Return the specs of the frame type
190    fn specs(&self) -> u8;
191}
192
193impl<T: GetFrameType> FrameFeture for T {
194    fn belongs_to(&self, packet_type: Type) -> bool {
195        self.frame_type().belongs_to(packet_type)
196    }
197
198    fn specs(&self) -> u8 {
199        self.frame_type().specs()
200    }
201}
202
203impl FrameFeture for FrameType {
204    fn belongs_to(&self, packet_type: Type) -> bool {
205        use crate::packet::r#type::{
206            long::{Type::V1, Ver1},
207            short::OneRtt,
208        };
209        // IH01
210        let i = matches!(packet_type, Type::Long(V1(Ver1::INITIAL)));
211        let h = matches!(packet_type, Type::Long(V1(Ver1::HANDSHAKE)));
212        let o = matches!(packet_type, Type::Long(V1(Ver1::ZERO_RTT)));
213        let l = matches!(packet_type, Type::Short(OneRtt(_)));
214
215        match self {
216            FrameType::Padding => i | h | o | l,
217            FrameType::Ping => i | h | o | l,
218            FrameType::Ack(_) => i | h | l,
219            FrameType::ResetStream => o | l,
220            FrameType::StopSending => o | l,
221            FrameType::Crypto => i | h | l,
222            FrameType::NewToken => l,
223            FrameType::Stream(_) => o | l,
224            FrameType::MaxData => o | l,
225            FrameType::MaxStreamData => o | l,
226            FrameType::MaxStreams(_) => o | l,
227            FrameType::DataBlocked => o | l,
228            FrameType::StreamDataBlocked => o | l,
229            FrameType::StreamsBlocked(_) => o | l,
230            FrameType::NewConnectionId => o | l,
231            FrameType::RetireConnectionId => o | l,
232            FrameType::PathChallenge => o | l,
233            FrameType::PathResponse => l,
234            // The application-specific variant of CONNECTION_CLOSE (type 0x1d) can only be
235            // sent using 0-RTT or 1-RTT packets;
236            // See [Section 12.5](https://www.rfc-editor.org/rfc/rfc9000.html#section-12.5).
237            //
238            // When an application wishes to abandon a connection during the handshake,
239            // an endpoint can send a CONNECTION_CLOSE frame (type 0x1c) with an error code
240            // of APPLICATION_ERROR in an Initial or Handshake packet.
241            FrameType::ConnectionClose(bit) => {
242                if *bit == 0 && i || h {
243                    true
244                } else {
245                    o | l
246                }
247            }
248            FrameType::HandshakeDone => l,
249            FrameType::Datagram(_) => o | l,
250        }
251    }
252
253    fn specs(&self) -> u8 {
254        let (n, c, p, f) = (
255            Spec::NonAckEliciting as u8,
256            Spec::CongestionControlFree as u8,
257            Spec::ProbeNewPath as u8,
258            Spec::FlowControlled as u8,
259        );
260        match self {
261            FrameType::Padding => n | p,
262            FrameType::Ack(_) => n | c,
263            FrameType::Stream(_) => f,
264            FrameType::NewConnectionId => p,
265            FrameType::PathChallenge => p,
266            FrameType::PathResponse => p,
267            // different from [table 3](https://www.rfc-editor.org/rfc/rfc9000.html#table-3),
268            // add the [`Spec::Con`] for the CONNECTION_CLOSE frame
269            FrameType::ConnectionClose(_) => n | c,
270            _ => 0,
271        }
272    }
273}
274
275impl TryFrom<VarInt> for FrameType {
276    type Error = Error;
277
278    fn try_from(frame_type: VarInt) -> Result<Self, Self::Error> {
279        Ok(match frame_type.into_inner() {
280            0x00 => FrameType::Padding,
281            0x01 => FrameType::Ping,
282            // The last bit is the ECN flag.
283            ty @ (0x02 | 0x03) => FrameType::Ack(ty as u8 & 0b1),
284            0x04 => FrameType::ResetStream,
285            0x05 => FrameType::StopSending,
286            0x06 => FrameType::Crypto,
287            0x07 => FrameType::NewToken,
288            // The last three bits are the offset, length, and fin flag bits respectively.
289            ty @ 0x08..=0x0f => FrameType::Stream(ty as u8 & 0b111),
290            0x10 => FrameType::MaxData,
291            0x11 => FrameType::MaxStreamData,
292            // The last bit is the direction flag bit, 0 indicates bidirectional, 1 indicates unidirectional.
293            ty @ (0x12 | 0x13) => FrameType::MaxStreams(ty as u8 & 0b1),
294            0x14 => FrameType::DataBlocked,
295            0x15 => FrameType::StreamDataBlocked,
296            // The last bit is the direction flag bit, 0 indicates bidirectional, 1 indicates unidirectional.
297            ty @ (0x16 | 0x17) => FrameType::StreamsBlocked(ty as u8 & 0b1),
298            0x18 => FrameType::NewConnectionId,
299            0x19 => FrameType::RetireConnectionId,
300            0x1a => FrameType::PathChallenge,
301            0x1b => FrameType::PathResponse,
302            // The last bit is the layer flag bit, 0 indicates application layer, 1 indicates transport layer.
303            ty @ (0x1c | 0x1d) => FrameType::ConnectionClose(ty as u8 & 0x1),
304            0x1e => FrameType::HandshakeDone,
305            // The last bit is the length flag bit, 0 the length field is absent and the Datagram Data
306            // field extends to the end of the packet, 1 the length field is present.
307            ty @ (0x30 | 0x31) => FrameType::Datagram(ty as u8 & 1),
308            // May be extension frame
309            _ => return Err(Self::Error::InvalidType(frame_type)),
310        })
311    }
312}
313
314impl From<FrameType> for VarInt {
315    fn from(frame_type: FrameType) -> Self {
316        match frame_type {
317            FrameType::Padding => VarInt::from_u32(0x00),
318            FrameType::Ping => VarInt::from_u32(0x01),
319            FrameType::Ack(ecn) => VarInt::from(0x02 | ecn),
320            FrameType::ResetStream => VarInt::from_u32(0x04),
321            FrameType::StopSending => VarInt::from_u32(0x05),
322            FrameType::Crypto => VarInt::from_u32(0x06),
323            FrameType::NewToken => VarInt::from_u32(0x07),
324            FrameType::Stream(flag) => VarInt::from(0x08 | flag),
325            FrameType::MaxData => VarInt::from_u32(0x10),
326            FrameType::MaxStreamData => VarInt::from_u32(0x11),
327            FrameType::MaxStreams(dir) => VarInt::from(0x12 | dir),
328            FrameType::DataBlocked => VarInt::from_u32(0x14),
329            FrameType::StreamDataBlocked => VarInt::from_u32(0x15),
330            FrameType::StreamsBlocked(dir) => VarInt::from(0x16 | dir),
331            FrameType::NewConnectionId => VarInt::from_u32(0x18),
332            FrameType::RetireConnectionId => VarInt::from_u32(0x19),
333            FrameType::PathChallenge => VarInt::from_u32(0x1a),
334            FrameType::PathResponse => VarInt::from_u32(0x1b),
335            FrameType::ConnectionClose(layer) => VarInt::from(0x1c | layer),
336            FrameType::HandshakeDone => VarInt::from_u32(0x1e),
337            FrameType::Datagram(with_len) => VarInt::from(0x30 | with_len),
338        }
339    }
340}
341
342/// Parse the frame type from the input buffer,
343/// [nom](https://docs.rs/nom/latest/nom/) parser style.
344pub fn be_frame_type(input: &[u8]) -> nom::IResult<&[u8], FrameType, Error> {
345    let (remain, frame_type) = crate::varint::be_varint(input).map_err(|_| {
346        nom::Err::Error(Error::IncompleteType(format!(
347            "Incomplete frame type from input: {input:?}"
348        )))
349    })?;
350    let frame_type = FrameType::try_from(frame_type).map_err(nom::Err::Error)?;
351    Ok((remain, frame_type))
352}
353
354/// Sum type of all the stream related frames except [`StreamFrame`].
355#[derive(Debug, Clone, Eq, PartialEq)]
356#[enum_dispatch(EncodeFrame, GetFrameType)]
357pub enum StreamCtlFrame {
358    /// RESET_STREAM frame, see [`ResetStreamFrame`].
359    ResetStream(ResetStreamFrame),
360    /// STOP_SENDING frame, see [`StopSendingFrame`].
361    StopSending(StopSendingFrame),
362    /// MAX_STREAM_DATA frame, see [`MaxStreamDataFrame`].
363    MaxStreamData(MaxStreamDataFrame),
364    /// MAX_STREAMS frame, see [`MaxStreamsFrame`].
365    MaxStreams(MaxStreamsFrame),
366    /// STREAM_DATA_BLOCKED frame, see [`StreamDataBlockedFrame`].
367    StreamDataBlocked(StreamDataBlockedFrame),
368    /// STREAMS_BLOCKED frame, see [`StreamsBlockedFrame`].
369    StreamsBlocked(StreamsBlockedFrame),
370}
371
372/// Sum type of all the reliable frames.
373#[derive(Debug, Clone, Eq, PartialEq)]
374#[enum_dispatch(EncodeFrame, GetFrameType)]
375pub enum ReliableFrame {
376    /// NEW_TOKEN frame, see [`NewTokenFrame`].
377    NewToken(NewTokenFrame),
378    /// MAX_DATA frame, see [`MaxDataFrame`].
379    MaxData(MaxDataFrame),
380    /// DATA_BLOCKED frame, see [`DataBlockedFrame`].
381    DataBlocked(DataBlockedFrame),
382    /// NEW_CONNECTION_ID frame, see [`NewConnectionIdFrame`].
383    NewConnectionId(NewConnectionIdFrame),
384    /// RETIRE_CONNECTION_ID frame, see [`RetireConnectionIdFrame`].
385    RetireConnectionId(RetireConnectionIdFrame),
386    /// HANDSHAKE_DONE frame, see [`HandshakeDoneFrame`].
387    HandshakeDone(HandshakeDoneFrame),
388    /// STREAM control frame, see [`StreamCtlFrame`].
389    Stream(StreamCtlFrame),
390}
391
392/// Sum type of all the frames.
393///
394/// The data frames' body are stored in the second field.
395#[derive(Debug, Clone, Eq, PartialEq)]
396pub enum Frame {
397    /// PADDING frame, see [`PaddingFrame`].
398    Padding(PaddingFrame),
399    /// PING frame, see [`PingFrame`].
400    Ping(PingFrame),
401    /// ACK frame, see [`AckFrame`].
402    Ack(AckFrame),
403    /// CONNECTION_CLOSE frame, see [`ConnectionCloseFrame`].
404    Close(ConnectionCloseFrame),
405    /// NEW_TOKEN frame, see [`NewTokenFrame`].
406    NewToken(NewTokenFrame),
407    /// MAX_DATA frame, see [`MaxDataFrame`].
408    MaxData(MaxDataFrame),
409    /// DATA_BLOCKED frame, see [`DataBlockedFrame`].
410    DataBlocked(DataBlockedFrame),
411    /// NEW_CONNECTION_ID frame, see [`NewConnectionIdFrame`].
412    NewConnectionId(NewConnectionIdFrame),
413    /// RETIRE_CONNECTION_ID frame, see [`RetireConnectionIdFrame`].
414    RetireConnectionId(RetireConnectionIdFrame),
415    /// HANDSHAKE_DONE frame, see [`HandshakeDoneFrame`].
416    HandshakeDone(HandshakeDoneFrame),
417    /// PATH_CHALLENGE frame, see [`PathChallengeFrame`].
418    Challenge(PathChallengeFrame),
419    /// PATH_RESPONSE frame, see [`PathResponseFrame`].
420    Response(PathResponseFrame),
421    /// Stream control frame, see [`StreamCtlFrame`].
422    StreamCtl(StreamCtlFrame),
423    /// STREAM frame and its data, see [`StreamFrame`].
424    Stream(StreamFrame, Bytes),
425    /// CRYPTO frame and its data, see [`CryptoFrame`].
426    Crypto(CryptoFrame, Bytes),
427    /// DATAGRAM frame and its data, see [`DatagramFrame`].
428    Datagram(DatagramFrame, Bytes),
429}
430
431impl From<ReliableFrame> for Frame {
432    fn from(frame: ReliableFrame) -> Self {
433        match frame {
434            ReliableFrame::NewToken(new_token_frame) => Frame::NewToken(new_token_frame),
435            ReliableFrame::MaxData(max_data_frame) => Frame::MaxData(max_data_frame),
436            ReliableFrame::DataBlocked(data_blocked_frame) => {
437                Frame::DataBlocked(data_blocked_frame)
438            }
439            ReliableFrame::NewConnectionId(new_connection_id_frame) => {
440                Frame::NewConnectionId(new_connection_id_frame)
441            }
442            ReliableFrame::RetireConnectionId(retire_connection_id_frame) => {
443                Frame::RetireConnectionId(retire_connection_id_frame)
444            }
445            ReliableFrame::HandshakeDone(handshake_done_frame) => {
446                Frame::HandshakeDone(handshake_done_frame)
447            }
448            ReliableFrame::Stream(stream_frame) => Frame::StreamCtl(stream_frame),
449        }
450    }
451}
452
453/// Some modules that need send specific frames can implement `SendFrame` trait directly.
454///
455/// Alternatively, a temporary buffer that stores certain frames can also implement this trait,
456/// But additional processing is required to ensure that the frames in the buffer are eventually
457/// sent to the peer.
458pub trait SendFrame<T> {
459    /// Need send the frames to the peer
460    fn send_frame<I: IntoIterator<Item = T>>(&self, iter: I);
461}
462
463/// Some modules that need receive specific frames can implement `ReceiveFrame` trait directly.
464///
465/// Alternatively, a temporary buffer that stores certain frames can also implement this trait,
466/// But additional processing is required to ensure that the frames in the buffer are eventually
467/// delivered to the corresponding modules.
468pub trait ReceiveFrame<T> {
469    type Output;
470
471    /// Receive the frames from the peer
472    fn recv_frame(&self, frame: &T) -> Result<Self::Output, crate::error::Error>;
473}
474
475/// Reads frames from a buffer until the packet buffer is empty.
476#[derive(Deref, DerefMut)]
477pub struct FrameReader {
478    #[deref]
479    #[deref_mut]
480    payload: Bytes,
481    packet_type: Type,
482}
483
484impl FrameReader {
485    /// Creates a [`FrameReader`] for a packet of type `packet_type`
486    pub fn new(payload: Bytes, packet_type: Type) -> Self {
487        Self {
488            payload,
489            packet_type,
490        }
491    }
492}
493
494impl Iterator for FrameReader {
495    type Item = Result<(Frame, FrameType), Error>;
496
497    fn next(&mut self) -> Option<Self::Item> {
498        if self.payload.is_empty() {
499            return None;
500        }
501
502        match io::be_frame(&self.payload, self.packet_type) {
503            Ok((consumed, frame, frame_type)) => {
504                self.payload.advance(consumed);
505                Some(Ok((frame, frame_type)))
506            }
507            Err(e) => Some(Err(e)),
508        }
509    }
510}
511
512impl<T: BufMut> WriteFrame<StreamCtlFrame> for T {
513    fn put_frame(&mut self, frame: &StreamCtlFrame) {
514        match frame {
515            StreamCtlFrame::ResetStream(frame) => self.put_frame(frame),
516            StreamCtlFrame::StopSending(frame) => self.put_frame(frame),
517            StreamCtlFrame::MaxStreamData(frame) => self.put_frame(frame),
518            StreamCtlFrame::MaxStreams(frame) => self.put_frame(frame),
519            StreamCtlFrame::StreamDataBlocked(frame) => self.put_frame(frame),
520            StreamCtlFrame::StreamsBlocked(frame) => self.put_frame(frame),
521        }
522    }
523}
524
525impl<T: BufMut> WriteFrame<ReliableFrame> for T {
526    fn put_frame(&mut self, frame: &ReliableFrame) {
527        match frame {
528            ReliableFrame::NewToken(frame) => self.put_frame(frame),
529            ReliableFrame::MaxData(frame) => self.put_frame(frame),
530            ReliableFrame::DataBlocked(frame) => self.put_frame(frame),
531            ReliableFrame::NewConnectionId(frame) => self.put_frame(frame),
532            ReliableFrame::RetireConnectionId(frame) => self.put_frame(frame),
533            ReliableFrame::HandshakeDone(frame) => self.put_frame(frame),
534            ReliableFrame::Stream(frame) => self.put_frame(frame),
535        }
536    }
537}
538
539#[cfg(test)]
540mod tests {
541    use nom::Parser;
542
543    use super::*;
544    use crate::{
545        packet::{
546            PacketContains,
547            r#type::{
548                Type,
549                long::{Type::V1, Ver1},
550                short::OneRtt,
551            },
552        },
553        varint::{WriteVarInt, be_varint},
554    };
555
556    #[test]
557    fn test_frame_type_conversion() {
558        let frame_types = vec![
559            FrameType::Padding,
560            FrameType::Ping,
561            FrameType::Ack(0),
562            FrameType::Stream(0),
563            FrameType::MaxData,
564            FrameType::ConnectionClose(0),
565            FrameType::HandshakeDone,
566            FrameType::Datagram(0),
567        ];
568
569        for frame_type in frame_types {
570            let byte: VarInt = frame_type.into();
571            assert_eq!(FrameType::try_from(byte).unwrap(), frame_type);
572        }
573    }
574
575    #[test]
576    fn test_frame_type_specs() {
577        assert!(FrameType::Padding.specs().contain(Spec::NonAckEliciting));
578        assert!(
579            FrameType::Ack(0)
580                .specs()
581                .contain(Spec::CongestionControlFree)
582        );
583        assert!(FrameType::Stream(0).specs().contain(Spec::FlowControlled));
584        assert!(FrameType::PathChallenge.specs().contain(Spec::ProbeNewPath));
585    }
586
587    #[test]
588    fn test_frame_type_belongs_to() {
589        let initial = Type::Long(V1(Ver1::INITIAL));
590        assert!(FrameType::Padding.belongs_to(initial));
591        assert!(FrameType::Ping.belongs_to(initial));
592        assert!(FrameType::Ack(0).belongs_to(initial));
593        assert!(!FrameType::Stream(0).belongs_to(initial));
594    }
595
596    #[test]
597    fn test_frame_reader() {
598        let mut buf = bytes::BytesMut::new();
599        buf.put_u8(0x00); // PADDING
600        buf.put_u8(0x01); // PING
601
602        let packet_type = Type::Long(V1(Ver1::INITIAL));
603        let mut reader = FrameReader::new(buf.freeze(), packet_type);
604
605        // Read PADDING frame
606        let (frame, frame_type) = reader.next().unwrap().unwrap();
607        assert!(matches!(frame, Frame::Padding(_)));
608        assert!(frame_type.specs().contain(Spec::NonAckEliciting));
609
610        // Read PING frame
611        let (frame, frame_type) = reader.next().unwrap().unwrap();
612        assert!(matches!(frame, Frame::Ping(_)));
613        assert!(!frame_type.specs().contain(Spec::NonAckEliciting));
614
615        // No more frames
616        assert!(reader.next().is_none());
617    }
618
619    #[test]
620    fn test_invalid_frame_type() {
621        assert!(FrameType::try_from(VarInt::from_u32(0xFF)).is_err());
622    }
623
624    #[test]
625    fn test_extetion_frame_type() {
626        use crate::varint::WriteVarInt;
627
628        #[derive(Debug, Clone, Eq, PartialEq)]
629        struct AddAddressFrame {
630            pub seq_num: VarInt,
631            pub tire: VarInt,
632            pub nat_type: VarInt,
633        }
634
635        fn be_add_address_frame(input: &[u8]) -> nom::IResult<&[u8], AddAddressFrame> {
636            use nom::{combinator::verify, sequence::preceded};
637            preceded(
638                verify(be_varint, |typ| typ == &VarInt::from_u32(0xff)),
639                (be_varint, be_varint, be_varint),
640            )
641            .map(|(seq_num, tire, nat_type)| AddAddressFrame {
642                seq_num,
643                tire,
644                nat_type,
645            })
646            .parse(input)
647        }
648
649        fn parse_address_frame(input: &[u8]) -> Result<(usize, AddAddressFrame), Error> {
650            let origin = input.len();
651            let (remain, frame) = be_add_address_frame(input).map_err(|_| {
652                Error::IncompleteType(format!("Incomplete frame type from input: {input:?}"))
653            })?;
654            let consumed = origin - remain.len();
655            Ok((consumed, frame))
656        }
657
658        impl<T: bytes::BufMut> super::io::WriteFrame<AddAddressFrame> for T {
659            fn put_frame(&mut self, frame: &AddAddressFrame) {
660                self.put_varint(&0xff_u32.into());
661                self.put_varint(&frame.seq_num);
662                self.put_varint(&frame.tire);
663                self.put_varint(&frame.nat_type);
664            }
665        }
666
667        let mut buf = bytes::BytesMut::new();
668        let add_address_frame = AddAddressFrame {
669            seq_num: VarInt::from_u32(0x01),
670            tire: VarInt::from_u32(0x02),
671            nat_type: VarInt::from_u32(0x03),
672        };
673        buf.put_frame(&add_address_frame);
674        buf.put_frame(&PaddingFrame);
675        buf.put_frame(&PaddingFrame);
676        buf.put_frame(&add_address_frame);
677        buf.put_varint(&0xfe_u32.into());
678        let mut padding_count = 0;
679        let mut add_address_count = 0;
680        let mut reader = FrameReader::new(buf.freeze(), Type::Short(OneRtt(0.into())));
681        loop {
682            match reader.next() {
683                Some(Ok((frame, typ))) => {
684                    assert!(matches!(frame, Frame::Padding(_)));
685                    assert_eq!(typ, FrameType::Padding);
686                    padding_count += 1;
687                }
688                Some(Err(_e)) => {
689                    // Parse Extension frame
690                    if let Ok((consum, frame)) = parse_address_frame(&reader) {
691                        reader.advance(consum);
692                        assert_eq!(frame, add_address_frame);
693                        add_address_count += 1;
694                    } else {
695                        reader.clear();
696                    }
697                }
698                None => break,
699            };
700        }
701        assert_eq!(padding_count, 2);
702        assert_eq!(add_address_count, 2);
703    }
704
705    #[test]
706    fn test_handless_extension_frame() {
707        let mut buf = bytes::BytesMut::new();
708        buf.put_frame(&PaddingFrame);
709        buf.put_frame(&PaddingFrame);
710        // error frame type
711        buf.put_varint(&0xfe_u32.into());
712        buf.put_frame(&PaddingFrame);
713
714        let mut padding_count = 0;
715        let _ = FrameReader::new(buf.freeze(), Type::Short(OneRtt(0.into()))).try_fold(
716            PacketContains::default(),
717            |packet_contains, frame| {
718                let (frame, frame_type) = frame?;
719
720                assert!(matches!(frame, Frame::Padding(_)));
721                assert_eq!(frame_type, FrameType::Padding);
722                padding_count += 1;
723                Result::<_, Error>::Ok(packet_contains)
724            },
725        );
726
727        assert_eq!(padding_count, 2);
728    }
729}