wtransport_proto/
datagram.rs

1use crate::bytes::BufferReader;
2use crate::bytes::BufferWriter;
3use crate::bytes::BytesReader;
4use crate::bytes::BytesWriter;
5use crate::bytes::EndOfBuffer;
6use crate::error::ErrorCode;
7use crate::ids::InvalidQStreamId;
8use crate::ids::QStreamId;
9
10/// An HTTP3 datagram.
11pub struct Datagram<'a> {
12    qstream_id: QStreamId,
13    payload: &'a [u8],
14}
15
16impl<'a> Datagram<'a> {
17    /// Creates a new [`Datagram`] with a given payload.
18    #[inline(always)]
19    pub fn new(qstream_id: QStreamId, payload: &'a [u8]) -> Self {
20        Self {
21            qstream_id,
22            payload,
23        }
24    }
25
26    /// Reads [`Datagram`] from a QUIC datagram.
27    pub fn read(quic_datagram: &'a [u8]) -> Result<Self, ErrorCode> {
28        let mut buffer_reader = BufferReader::new(quic_datagram);
29
30        let varint = buffer_reader.get_varint().ok_or(ErrorCode::Datagram)?;
31
32        let qstream_id =
33            QStreamId::try_from_varint(varint).map_err(|InvalidQStreamId| ErrorCode::Datagram)?;
34
35        let payload = buffer_reader.buffer_remaining();
36
37        Ok(Self {
38            qstream_id,
39            payload,
40        })
41    }
42
43    /// Writes a [`Datagram`] as QUIC datagram into `buffer`.
44    ///
45    /// It returns the number of bytes written.
46    /// It returns [`Err`] if the `buffer` does not have enough capacity.
47    /// See [`Self::write_size`].
48    ///
49    /// In case of [`Err`], `buffer` is not written.
50    pub fn write(&self, buffer: &mut [u8]) -> Result<usize, EndOfBuffer> {
51        if buffer.len() < self.write_size() {
52            return Err(EndOfBuffer);
53        }
54
55        let mut buffer_writer = BufferWriter::new(buffer);
56
57        buffer_writer
58            .put_varint(self.qstream_id.into_varint())
59            .expect("Buffer has capacity");
60
61        buffer_writer
62            .put_bytes(self.payload)
63            .expect("Buffer has capacity");
64
65        Ok(buffer_writer.offset())
66    }
67
68    /// Returns the needed capacity to write this datagram into a buffer.
69    #[inline(always)]
70    pub fn write_size(&self) -> usize {
71        Self::header_size(self.qstream_id) + self.payload.len()
72    }
73
74    /// Returns the HTTP3 header.
75    ///
76    /// Computes the space overhead (HTTP3 header) due to the `qstream_id`
77    /// encoding into an HTTP3 datagram.
78    #[inline(always)]
79    pub fn header_size(qstream_id: QStreamId) -> usize {
80        qstream_id.into_varint().size()
81    }
82
83    /// Returns the associated [`QStreamId`].
84    #[inline(always)]
85    pub fn qstream_id(&self) -> QStreamId {
86        self.qstream_id
87    }
88
89    /// Returns the payload.
90    #[inline(always)]
91    pub fn payload(&self) -> &[u8] {
92        self.payload
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::varint::VarInt;
100    use utils::build_datagram;
101    use utils::QStreamIdType;
102    use utils::PAYLOAD;
103
104    #[test]
105    fn read_ok() {
106        let dgram = build_datagram(QStreamIdType::Valid, PAYLOAD);
107        let qstream_id = dgram.qstream_id();
108
109        let mut buffer = vec![0; dgram.write_size() + 42];
110        let written = dgram.write(&mut buffer).unwrap();
111
112        let dgram = Datagram::read(&buffer[..written]).unwrap();
113        assert_eq!(dgram.qstream_id(), qstream_id);
114        assert_eq!(dgram.payload(), PAYLOAD);
115    }
116
117    #[test]
118    fn read_too_short() {
119        let dgram = build_datagram(QStreamIdType::Valid, PAYLOAD);
120
121        let mut buffer = vec![0; dgram.write_size() + 42];
122        dgram.write(&mut buffer).unwrap();
123
124        assert!(matches!(
125            Datagram::read(&buffer[..1]),
126            Err(ErrorCode::Datagram)
127        ));
128    }
129
130    #[test]
131    fn read_invalid_qstream_id() {
132        let dgram = build_datagram(QStreamIdType::Invalid, PAYLOAD);
133
134        let mut buffer = vec![0; dgram.write_size() + 42];
135        let written = dgram.write(&mut buffer).unwrap();
136
137        assert!(matches!(
138            Datagram::read(&buffer[..written]),
139            Err(ErrorCode::Datagram)
140        ));
141    }
142
143    #[test]
144    fn write_ok() {
145        let dgram = build_datagram(QStreamIdType::Valid, PAYLOAD);
146        let dgram_write_size = dgram.write_size();
147
148        let mut buffer = vec![0; dgram_write_size];
149        let written = dgram.write(&mut buffer).unwrap();
150        assert_eq!(written, dgram_write_size);
151    }
152
153    #[test]
154    fn write_out() {
155        let dgram = build_datagram(QStreamIdType::Valid, PAYLOAD);
156        let dgram_write_size = dgram.write_size();
157
158        let mut buffer = vec![0; dgram_write_size - 1];
159        assert!(dgram.write(&mut buffer).is_err());
160    }
161
162    mod utils {
163        use super::*;
164
165        pub const PAYLOAD: &[u8] = b"This is a testing payload";
166
167        pub enum QStreamIdType {
168            Valid,
169            Invalid,
170        }
171
172        impl QStreamIdType {
173            /// This function is for **testing purpose only**; it might produce an invalid `QStreamId`!
174            fn into_session_id(self) -> QStreamId {
175                match self {
176                    QStreamIdType::Valid => QStreamId::MAX,
177                    QStreamIdType::Invalid => {
178                        let varint = VarInt::try_from_u64(QStreamId::MAX.into_u64() + 1).unwrap();
179                        QStreamId::maybe_invalid(varint)
180                    }
181                }
182            }
183        }
184
185        /// This function is for **testing purpose only**; it might produce an invalid `Datagram`!
186        pub fn build_datagram(qstream_id_type: QStreamIdType, payload: &[u8]) -> Datagram {
187            Datagram::new(qstream_id_type.into_session_id(), payload)
188        }
189    }
190}