noq_proto/connection/
datagrams.rs1use std::collections::VecDeque;
2
3use bytes::Bytes;
4use thiserror::Error;
5use tracing::{debug, trace};
6
7use super::Connection;
8use crate::{
9 FrameStats, TransportError,
10 connection::PacketBuilder,
11 frame::{Datagram, FrameStruct},
12};
13
14pub struct Datagrams<'a> {
16 pub(super) conn: &'a mut Connection,
17}
18
19impl Datagrams<'_> {
20 pub fn send(&mut self, data: Bytes, drop: bool) -> Result<(), SendDatagramError> {
30 if self.conn.config.datagram_receive_buffer_size.is_none() {
31 return Err(SendDatagramError::Disabled);
32 }
33 let max = self
34 .max_size()
35 .ok_or(SendDatagramError::UnsupportedByPeer)?;
36 let send_buffer_size = self.conn.config.datagram_send_buffer_size;
37 if data.len() > Ord::min(max, send_buffer_size) {
38 return Err(SendDatagramError::TooLarge);
39 }
40 if drop {
41 self.conn
42 .datagrams
43 .make_space_for(data.len(), send_buffer_size);
44 } else if !self
45 .conn
46 .datagrams
47 .has_send_buffer_space(data.len(), send_buffer_size)
48 {
49 self.conn.datagrams.send_blocked = true;
50 return Err(SendDatagramError::Blocked(data));
51 }
52 self.conn.datagrams.outgoing_total += data.len();
53 self.conn.datagrams.outgoing.push_back(Datagram { data });
54 Ok(())
55 }
56
57 pub fn max_size(&self) -> Option<usize> {
70 let max_size = self.conn.current_mtu() as usize
74 - self.conn.predict_1rtt_overhead_no_pn()
75 - Datagram::SIZE_BOUND;
76 let limit = self
77 .conn
78 .peer_params
79 .max_datagram_frame_size?
80 .into_inner()
81 .saturating_sub(Datagram::SIZE_BOUND as u64);
82 Some(limit.min(max_size as u64) as usize)
83 }
84
85 pub fn recv(&mut self) -> Option<Bytes> {
87 self.conn.datagrams.recv()
88 }
89
90 pub fn send_buffer_space(&self) -> usize {
95 self.conn
96 .config
97 .datagram_send_buffer_size
98 .saturating_sub(self.conn.datagrams.outgoing_total)
99 }
100}
101
102#[derive(Default)]
103pub(super) struct DatagramState {
104 pub(super) recv_buffered: usize,
107 pub(super) incoming: VecDeque<Datagram>,
108 pub(super) outgoing: VecDeque<Datagram>,
109 pub(super) outgoing_total: usize,
110 pub(super) send_blocked: bool,
111}
112
113impl DatagramState {
114 pub(super) fn received(
115 &mut self,
116 datagram: Datagram,
117 window: &Option<usize>,
118 ) -> Result<bool, TransportError> {
119 let window = match window {
120 None => {
121 return Err(TransportError::PROTOCOL_VIOLATION(
122 "unexpected DATAGRAM frame",
123 ));
124 }
125 Some(x) => *x,
126 };
127
128 if datagram.data.len() > window {
129 return Err(TransportError::PROTOCOL_VIOLATION("oversized datagram"));
130 }
131
132 let was_empty = self.recv_buffered == 0;
133 while datagram.data.len() + self.recv_buffered > window {
134 debug!("dropping stale datagram");
135 self.recv();
136 }
137
138 self.recv_buffered += datagram.data.len();
139 self.incoming.push_back(datagram);
140 Ok(was_empty)
141 }
142
143 fn make_space_for(&mut self, datagram_len: usize, send_buffer_size: usize) {
144 while !self.has_send_buffer_space(datagram_len, send_buffer_size) {
145 let Some(prev) = self.outgoing.pop_front() else {
146 break;
147 };
148 trace!(len = prev.data.len(), "dropping outgoing datagram");
149 self.outgoing_total -= prev.data.len();
150 }
151 }
152
153 fn has_send_buffer_space(&self, datagram_len: usize, send_buffer_size: usize) -> bool {
154 let Some(total) = self.outgoing_total.checked_add(datagram_len) else {
155 return false;
156 };
157
158 total <= send_buffer_size
159 }
160
161 pub(super) fn drop_oversized(&mut self, max_payload: usize) -> bool {
168 let mut dropped_any = false;
169 self.outgoing.retain(|datagram| {
170 let result = datagram.data.len() < max_payload;
171 if !result {
172 trace!(
173 "dropping {} byte datagram violating {} byte limit",
174 datagram.data.len(),
175 max_payload
176 );
177 self.outgoing_total -= datagram.data.len();
178 dropped_any = true;
179 }
180 result
181 });
182 dropped_any
183 }
184
185 pub(super) fn write<'a, 'b>(
190 &mut self,
191 buf: &mut PacketBuilder<'a, 'b>,
192 stat: &mut FrameStats,
193 ) -> bool {
194 let Some(datagram) = self.outgoing.pop_front() else {
195 return false;
196 };
197
198 if buf.frame_space_remaining() < datagram.size(true) {
199 self.outgoing.push_front(datagram);
202 return false;
203 }
204
205 self.outgoing_total -= datagram.data.len();
206 buf.write_frame(datagram, stat);
207 true
208 }
209
210 pub(super) fn recv(&mut self) -> Option<Bytes> {
211 let x = self.incoming.pop_front()?.data;
212 self.recv_buffered -= x.len();
213 Some(x)
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220
221 #[test]
222 fn make_space_for_accounts_for_new_datagram() {
223 let mut state = DatagramState::default();
224 state.outgoing.push_back(Datagram {
225 data: Bytes::from_static(&[0; 7]),
226 });
227 state.outgoing.push_back(Datagram {
228 data: Bytes::from_static(&[0; 2]),
229 });
230 state.outgoing_total = 9;
231
232 state.make_space_for(4, 10);
233
234 assert_eq!(state.outgoing.len(), 1);
235 assert_eq!(state.outgoing[0].data.len(), 2);
236 assert_eq!(state.outgoing_total, 2);
237 }
238
239 #[test]
240 fn make_space_for_handles_overflowing_capacity_check() {
241 let mut state = DatagramState::default();
242 state.outgoing.push_back(Datagram {
243 data: Bytes::from_static(&[0]),
244 });
245 state.outgoing_total = usize::MAX - 1;
246
247 state.make_space_for(2, usize::MAX);
248
249 assert!(state.outgoing.is_empty());
250 assert_eq!(state.outgoing_total, usize::MAX - 2);
251 }
252}
253
254#[derive(Debug, Error, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
256pub enum SendDatagramError {
257 #[error("datagrams not supported by peer")]
259 UnsupportedByPeer,
260 #[error("datagram support disabled")]
262 Disabled,
263 #[error("datagram too large")]
268 TooLarge,
269 #[error("datagram send blocked")]
271 Blocked(Bytes),
272}