diff --git a/quinn-proto/src/connection/datagrams.rs b/quinn-proto/src/connection/datagrams.rs index 3b5e027f9a..957703b0b9 100644 --- a/quinn-proto/src/connection/datagrams.rs +++ b/quinn-proto/src/connection/datagrams.rs @@ -32,22 +32,18 @@ impl Datagrams<'_> { let max = self .max_size() .ok_or(SendDatagramError::UnsupportedByPeer)?; - if data.len() > max { + let send_buffer_size = self.conn.config.datagram_send_buffer_size; + if data.len() > Ord::min(max, send_buffer_size) { return Err(SendDatagramError::TooLarge); } if drop { - while self.conn.datagrams.outgoing_total > self.conn.config.datagram_send_buffer_size { - let prev = self - .conn - .datagrams - .outgoing - .pop_front() - .expect("datagrams.outgoing_total desynchronized"); - trace!(len = prev.data.len(), "dropping outgoing datagram"); - self.conn.datagrams.outgoing_total -= prev.data.len(); - } - } else if self.conn.datagrams.outgoing_total + data.len() - > self.conn.config.datagram_send_buffer_size + self.conn + .datagrams + .make_space_for(data.len(), send_buffer_size); + } else if !self + .conn + .datagrams + .has_send_buffer_space(data.len(), send_buffer_size) { self.conn.datagrams.send_blocked = true; return Err(SendDatagramError::Blocked(data)); @@ -140,6 +136,24 @@ impl DatagramState { Ok(was_empty) } + fn make_space_for(&mut self, datagram_len: usize, send_buffer_size: usize) { + while !self.has_send_buffer_space(datagram_len, send_buffer_size) { + let Some(prev) = self.outgoing.pop_front() else { + break; + }; + trace!(len = prev.data.len(), "dropping outgoing datagram"); + self.outgoing_total -= prev.data.len(); + } + } + + fn has_send_buffer_space(&self, datagram_len: usize, send_buffer_size: usize) -> bool { + let Some(total) = self.outgoing_total.checked_add(datagram_len) else { + return false; + }; + + total <= send_buffer_size + } + /// Discard outgoing datagrams with a payload larger than `max_payload` bytes /// /// Returns whether any datagrams were dropped. @@ -194,6 +208,43 @@ impl DatagramState { } } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn make_space_for_accounts_for_new_datagram() { + let mut state = DatagramState::default(); + state.outgoing.push_back(Datagram { + data: Bytes::from_static(&[0; 7]), + }); + state.outgoing.push_back(Datagram { + data: Bytes::from_static(&[0; 2]), + }); + state.outgoing_total = 9; + + state.make_space_for(4, 10); + + assert_eq!(state.outgoing.len(), 1); + assert_eq!(state.outgoing[0].data.len(), 2); + assert_eq!(state.outgoing_total, 2); + } + + #[test] + fn make_space_for_handles_overflowing_capacity_check() { + let mut state = DatagramState::default(); + state.outgoing.push_back(Datagram { + data: Bytes::from_static(&[0]), + }); + state.outgoing_total = usize::MAX - 1; + + state.make_space_for(2, usize::MAX); + + assert!(state.outgoing.is_empty()); + assert_eq!(state.outgoing_total, usize::MAX - 2); + } +} + /// Errors that can arise when sending a datagram #[derive(Debug, Error, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)] pub enum SendDatagramError { diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs index bcff9ae597..8ca1f5e6c1 100644 --- a/quinn-proto/src/tests/mod.rs +++ b/quinn-proto/src/tests/mod.rs @@ -2029,6 +2029,28 @@ fn datagram_recv_buffer_overflow() { assert_matches!(pair.server_datagrams(server_ch).recv(), None); } +#[test] +fn datagram_larger_than_send_buffer_is_too_large() { + let _guard = subscribe(); + let mut pair = Pair::default(); + let mut client_config = client_config(); + let mut transport_config = TransportConfig::default(); + transport_config.datagram_send_buffer_size(1); + client_config.transport_config(transport_config.into()); + let (client_ch, _) = pair.connect_with(client_config); + + assert_matches!( + pair.client_datagrams(client_ch) + .send(Bytes::from_static(&[0; 2]), true), + Err(SendDatagramError::TooLarge) + ); + assert_matches!( + pair.client_datagrams(client_ch) + .send(Bytes::from_static(&[0; 2]), false), + Err(SendDatagramError::TooLarge) + ); +} + #[test] fn datagram_unsupported() { let _guard = subscribe();