wtransport_proto/
datagram.rs1use 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
10pub struct Datagram<'a> {
12 qstream_id: QStreamId,
13 payload: &'a [u8],
14}
15
16impl<'a> Datagram<'a> {
17 #[inline(always)]
19 pub fn new(qstream_id: QStreamId, payload: &'a [u8]) -> Self {
20 Self {
21 qstream_id,
22 payload,
23 }
24 }
25
26 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 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 #[inline(always)]
70 pub fn write_size(&self) -> usize {
71 Self::header_size(self.qstream_id) + self.payload.len()
72 }
73
74 #[inline(always)]
79 pub fn header_size(qstream_id: QStreamId) -> usize {
80 qstream_id.into_varint().size()
81 }
82
83 #[inline(always)]
85 pub fn qstream_id(&self) -> QStreamId {
86 self.qstream_id
87 }
88
89 #[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 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 pub fn build_datagram(qstream_id_type: QStreamIdType, payload: &[u8]) -> Datagram {
187 Datagram::new(qstream_id_type.into_session_id(), payload)
188 }
189 }
190}