Reject DATAGRAMs larger than the send buffer

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 {

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();