dtls: defer oversized poll outputs

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 
+  * Defer oversized DTLS poll outputs instead of panicking #138
   * Represent DTLS wire-code identifiers as compact newtypes (breaking) #137
   * Make public errors structured and fatal-only (breaking) #134
 

README.md

@@ -69,6 +69,7 @@ The output is an [`Output`][output] enum with borrowed
 references into your provided buffer:
 - `Packet(&[u8])`: send on your UDP socket
 - `Timeout(Instant)`: schedule a timer and call `handle_timeout` at/after it
+- `BufferTooSmall { needed }`: grow the caller-owned poll buffer and retry
 - `Connected`: handshake complete
 - `PeerCert(&[u8])`: peer leaf certificate (DER) — validate in your app
 - `KeyingMaterial(KeyingMaterial, SrtpProfile)`: DTLS‑SRTP export
@@ -114,6 +115,9 @@ fn example_event_loop(mut dtls: Dtls) -> Result<(), dimpl::Error> {
             match dtls.poll_output(&mut out_buf) {
                 Output::Packet(p) => send_udp(p),
                 Output::Timeout(t) => { next_wake = Some(t); break; }
+                Output::BufferTooSmall { needed } => {
+                    out_buf.resize(needed, 0);
+                }
                 Output::Connected => {
                     // DTLS established — application may start sending
                 }

src/auto.rs

@@ -320,10 +320,7 @@ impl ClientPending {
             if buf.len() < len {
                 // Buffer too small; keep needs_send armed so the packet
                 // is emitted on the next poll with a sufficiently large buffer.
-                let next = self
-                    .retransmit_at
-                    .unwrap_or(self.last_now + Duration::from_secs(1));
-                return Output::Timeout(next);
+                return Output::BufferTooSmall { needed: len };
             }
             self.needs_send = false;
             buf[..len].copy_from_slice(&self.wire_packet);

src/dtls12/client.rs

@@ -199,7 +199,13 @@ impl Client {
 
     pub fn poll_output<'a>(&mut self, buf: &'a mut [u8]) -> Output<'a> {
         if let Some(event) = self.local_events.pop_front() {
-            return event.into_output(buf, &self.server_certificates);
+            return match event.into_output(buf, &self.server_certificates) {
+                Ok(output) => output,
+                Err((event, needed)) => {
+                    self.local_events.push_front(event);
+                    Output::BufferTooSmall { needed }
+                }
+            };
         }
         self.engine.poll_output(buf, self.last_now)
     }
@@ -1376,20 +1382,23 @@ fn handshake_create_certificate_verify(body: &mut Buf, engine: &mut Engine) -> R
 }
 
 impl LocalEvent {
-    pub fn into_output<'a>(self, buf: &'a mut [u8], peer_certs: &[Buf]) -> Output<'a> {
+    pub fn into_output<'a>(
+        self,
+        buf: &'a mut [u8],
+        peer_certs: &[Buf],
+    ) -> Result<Output<'a>, (Self, usize)> {
         match self {
             LocalEvent::PeerCert => {
                 let l = peer_certs[0].len();
-                assert!(
-                    l <= buf.len(),
-                    "Output buffer too small for peer certificate"
-                );
+                if l > buf.len() {
+                    return Err((LocalEvent::PeerCert, l));
+                }
                 buf[..l].copy_from_slice(&peer_certs[0]);
-                Output::PeerCert(&buf[..l])
+                Ok(Output::PeerCert(&buf[..l]))
             }
-            LocalEvent::Connected => Output::Connected,
+            LocalEvent::Connected => Ok(Output::Connected),
             LocalEvent::KeyingMaterial(m, profile) => {
-                Output::KeyingMaterial(KeyingMaterial::new(&m), profile)
+                Ok(Output::KeyingMaterial(KeyingMaterial::new(&m), profile))
             }
         }
     }
@@ -1466,4 +1475,31 @@ mod tests {
 
         assert!(matches!(err, Error::InvalidState(_)));
     }
+
+    #[test]
+    fn peer_cert_output_is_deferred_when_buffer_is_too_small() {
+        let mut cert = Buf::new();
+        cert.extend_from_slice(&[1, 2, 3, 4]);
+        let certs = [cert];
+
+        let mut out = [0u8; 2];
+        let result = LocalEvent::PeerCert.into_output(&mut out, &certs);
+
+        assert!(matches!(result, Err((LocalEvent::PeerCert, 4))));
+    }
+
+    #[test]
+    fn peer_cert_output_is_copied_when_buffer_fits() {
+        let mut cert = Buf::new();
+        cert.extend_from_slice(&[1, 2, 3, 4]);
+        let certs = [cert];
+
+        let mut out = [0u8; 4];
+        let result = LocalEvent::PeerCert.into_output(&mut out, &certs);
+
+        match result {
+            Ok(Output::PeerCert(bytes)) => assert_eq!(bytes, &[1, 2, 3, 4]),
+            other => panic!("expected PeerCert output, got: {other:?}"),
+        }
+    }
 }

src/dtls12/engine.rs

@@ -17,6 +17,12 @@ use crate::{Config, Error, InternalError, Output, SeededRng};
 
 const MAX_DEFRAGMENT_PACKETS: usize = 50;
 
+enum PollBuffer<'a> {
+    Ready(&'a [u8]),
+    Empty(&'a mut [u8]),
+    TooSmall { needed: usize },
+}
+
 // Using debug_ignore_primary since CryptoContext doesn't implement Debug
 pub struct Engine {
     config: Arc<Config>,
@@ -416,12 +422,15 @@ impl Engine {
 
         // First check if we have any decrypted app data.
         let buf = match self.poll_app_data(buf) {
-            Ok(p) => return Output::ApplicationData(p),
-            Err(b) => b,
+            PollBuffer::Ready(p) => return Output::ApplicationData(p),
+            PollBuffer::Empty(b) => b,
+            PollBuffer::TooSmall { needed } => return Output::BufferTooSmall { needed },
         };
 
-        if let Ok(p) = self.poll_packet_tx(buf) {
-            return Output::Packet(p);
+        match self.poll_packet_tx(buf) {
+            PollBuffer::Ready(p) => return Output::Packet(p),
+            PollBuffer::Empty(_) => {}
+            PollBuffer::TooSmall { needed } => return Output::BufferTooSmall { needed },
         }
 
         if self.close_notify_received && !self.close_notify_reported {
@@ -434,9 +443,9 @@ impl Engine {
         Output::Timeout(next_timeout)
     }
 
-    fn poll_app_data<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [u8], &'a mut [u8]> {
+    fn poll_app_data<'a>(&mut self, buf: &'a mut [u8]) -> PollBuffer<'a> {
         if !self.release_app_data {
-            return Err(buf);
+            return PollBuffer::Empty(buf);
         }
 
         let mut unhandled = self
@@ -447,24 +456,21 @@ impl Engine {
             .skip_while(|r| r.is_handled());
 
         let Some(next) = unhandled.next() else {
-            return Err(buf);
+            return PollBuffer::Empty(buf);
         };
 
         let record_buffer = next.buffer();
         let fragment = next.record().fragment(record_buffer);
         let len = fragment.len();
 
-        assert!(
-            len <= buf.len(),
-            "Output buffer too small for application data {} > {}",
-            len,
-            buf.len()
-        );
+        if len > buf.len() {
+            return PollBuffer::TooSmall { needed: len };
+        }
 
         buf[..len].copy_from_slice(fragment);
         next.set_handled();
 
-        Ok(&buf[..len])
+        PollBuffer::Ready(&buf[..len])
     }
 
     fn purge_handled_queue_rx(&mut self) {
@@ -482,22 +488,20 @@ impl Engine {
         }
     }
 
-    fn poll_packet_tx<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [u8], &'a mut [u8]> {
-        let Some(p) = self.queue_tx.pop_front() else {
-            return Err(buf);
+    fn poll_packet_tx<'a>(&mut self, buf: &'a mut [u8]) -> PollBuffer<'a> {
+        let Some(p) = self.queue_tx.front() else {
+            return PollBuffer::Empty(buf);
         };
 
-        assert!(
-            p.len() <= buf.len(),
-            "Output buffer too small for packet {} > {}",
-            p.len(),
-            buf.len()
-        );
-
         let len = p.len();
-        buf[..len].copy_from_slice(&p);
+        if len > buf.len() {
+            return PollBuffer::TooSmall { needed: len };
+        }
+
+        buf[..len].copy_from_slice(p);
+        self.queue_tx.pop_front();
 
-        Ok(&buf[..len])
+        PollBuffer::Ready(&buf[..len])
     }
 
     fn poll_timeout(&self, now: Instant) -> Instant {

src/dtls12/server.rs

@@ -201,7 +201,13 @@ impl Server {
 
     pub fn poll_output<'a>(&mut self, buf: &'a mut [u8]) -> Output<'a> {
         if let Some(event) = self.local_events.pop_front() {
-            return event.into_output(buf, &self.client_certificates);
+            return match event.into_output(buf, &self.client_certificates) {
+                Ok(output) => output,
+                Err((event, needed)) => {
+                    self.local_events.push_front(event);
+                    Output::BufferTooSmall { needed }
+                }
+            };
         }
         self.engine.poll_output(buf, self.last_now)
     }

src/dtls13/client.rs

@@ -231,7 +231,13 @@ impl Client {
 
     pub fn poll_output<'a>(&mut self, buf: &'a mut [u8]) -> Output<'a> {
         if let Some(event) = self.local_events.pop_front() {
-            return event.into_output(buf, &self.server_certificates);
+            return match event.into_output(buf, &self.server_certificates) {
+                Ok(output) => output,
+                Err((event, needed)) => {
+                    self.local_events.push_front(event);
+                    Output::BufferTooSmall { needed }
+                }
+            };
         }
         self.engine.poll_output(buf, self.last_now)
     }
@@ -1547,20 +1553,23 @@ fn parse_certificate_request(cr_data: &[u8], base_offset: usize) -> Result<Optio
 // =========================================================================
 
 impl LocalEvent {
-    pub fn into_output<'a>(self, buf: &'a mut [u8], peer_certs: &[Buf]) -> Output<'a> {
+    pub fn into_output<'a>(
+        self,
+        buf: &'a mut [u8],
+        peer_certs: &[Buf],
+    ) -> Result<Output<'a>, (Self, usize)> {
         match self {
             LocalEvent::PeerCert => {
                 let l = peer_certs[0].len();
-                assert!(
-                    l <= buf.len(),
-                    "Output buffer too small for peer certificate"
-                );
+                if l > buf.len() {
+                    return Err((LocalEvent::PeerCert, l));
+                }
                 buf[..l].copy_from_slice(&peer_certs[0]);
-                Output::PeerCert(&buf[..l])
+                Ok(Output::PeerCert(&buf[..l]))
             }
-            LocalEvent::Connected => Output::Connected,
+            LocalEvent::Connected => Ok(Output::Connected),
             LocalEvent::KeyingMaterial(m, profile) => {
-                Output::KeyingMaterial(KeyingMaterial::new(&m), profile)
+                Ok(Output::KeyingMaterial(KeyingMaterial::new(&m), profile))
             }
         }
     }
@@ -1697,4 +1706,31 @@ mod tests {
             ))
         ));
     }
+
+    #[test]
+    fn peer_cert_output_is_deferred_when_buffer_is_too_small() {
+        let mut cert = Buf::new();
+        cert.extend_from_slice(&[1, 2, 3, 4]);
+        let certs = [cert];
+
+        let mut out = [0u8; 2];
+        let result = LocalEvent::PeerCert.into_output(&mut out, &certs);
+
+        assert!(matches!(result, Err((LocalEvent::PeerCert, 4))));
+    }
+
+    #[test]
+    fn peer_cert_output_is_copied_when_buffer_fits() {
+        let mut cert = Buf::new();
+        cert.extend_from_slice(&[1, 2, 3, 4]);
+        let certs = [cert];
+
+        let mut out = [0u8; 4];
+        let result = LocalEvent::PeerCert.into_output(&mut out, &certs);
+
+        match result {
+            Ok(Output::PeerCert(bytes)) => assert_eq!(bytes, &[1, 2, 3, 4]),
+            other => panic!("expected PeerCert output, got: {other:?}"),
+        }
+    }
 }