Defer DTLS handshake consumption until parse succeeds

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Unreleased
 
   * Fix DTLS 1.3 RFC 9147 conformance issues #147
+  * Reject malformed fragmented DTLS handshakes before consuming fragments #144
   * Represent DTLS wire-code identifiers as compact newtypes (breaking) #137
   * Make public errors structured and fatal-only (breaking) #134
 

src/dtls12/message/handshake.rs

@@ -13,13 +13,20 @@ use super::HelloVerifyRequest;
 use super::ServerHello;
 use super::ServerKeyExchange;
 use crate::buffer::Buf;
+use arrayvec::ArrayVec;
 use nom::Err;
 use nom::IResult;
 use nom::bytes::complete::take;
 use nom::error::{Error, ErrorKind};
 use nom::number::complete::be_u8;
 use nom::number::complete::{be_u16, be_u24};
 
+// Defensive stack cap over flattened handshake fragments selected for one
+// defragmentation attempt. This intentionally does not mirror the receive
+// queue's record-count cap; exceeding 50 fragments for one handshake implies
+// pathologically tiny records and is treated as invalid input.
+const MAX_DEFRAGMENT_HANDSHAKES: usize = 50;
+
 #[derive(Debug, PartialEq, Eq, Default, Clone, Copy)]
 pub struct Header {
     pub msg_type: MessageType,
@@ -154,8 +161,11 @@ impl Handshake {
         let Body::Fragment(range) = &first_handshake.body else {
             unreachable!("Non-Fragment body in defragment()")
         };
+        let mut handled = ArrayVec::<&Handshake, MAX_DEFRAGMENT_HANDSHAKES>::new();
+        handled
+            .try_push(first_handshake)
+            .map_err(|_| crate::InternalError::too_many_records())?;
         buffer.extend_from_slice(&first_buffer[range.clone()]);
-        first_handshake.set_handled();
 
         for (handshake, source_buf) in iter {
             if handshake.header.msg_type != first_handshake.header.msg_type {
@@ -166,7 +176,9 @@ impl Handshake {
                 unreachable!("Non-Fragment body in defragment()")
             };
 
-            handshake.handled.store(true, Ordering::Relaxed);
+            handled
+                .try_push(handshake)
+                .map_err(|_| crate::InternalError::too_many_records())?;
 
             buffer.extend_from_slice(&source_buf[range.clone()]);
         }
@@ -176,7 +188,35 @@ impl Handshake {
             return Err(crate::InternalError::parse_incomplete());
         }
 
-        // If transcript is provided, write the handshake header + body before parsing
+        let (rest, body) =
+            match Body::parse(buffer, 0, first_handshake.header.msg_type, cipher_suite) {
+                Ok(parsed) => parsed,
+                Err(err) => {
+                    mark_handled(handled);
+                    return Err(err.into());
+                }
+            };
+
+        if !rest.is_empty()
+            && first_handshake
+                .header
+                .msg_type
+                .rejects_trailing_body_bytes()
+        {
+            debug!("Defragmentation failed. Body::parse() did not consume the entire buffer");
+            mark_handled(handled);
+            return Err(crate::InternalError::parse_incomplete());
+        }
+
+        // Intentional boundary: Body::parse validates the handshake body shape and
+        // extension envelopes, but known extension payloads remain validated by the
+        // client/server state handlers. A transiently corrupted UDP datagram whose
+        // extension payload fails later may therefore have been consumed here; that
+        // recovery edge is accepted to keep this path parser-only and avoid the
+        // broader transaction/rollback machinery.
+        mark_handled(handled);
+
+        // If transcript is provided, write the handshake header + body after parsing succeeds.
         if let Some(transcript) = transcript {
             transcript.push(first_handshake.header.msg_type.as_u8());
             transcript.extend_from_slice(&first_handshake.header.length.to_be_bytes()[1..]);
@@ -187,13 +227,6 @@ impl Handshake {
             transcript.extend_from_slice(&buffer[..first_handshake.header.length as usize]);
         }
 
-        let (rest, body) = Body::parse(buffer, 0, first_handshake.header.msg_type, cipher_suite)?;
-
-        if !rest.is_empty() && first_handshake.header.msg_type == MessageType::Finished {
-            debug!("Defragmentation failed. Body::parse() did not consume the entire buffer");
-            return Err(crate::InternalError::parse_incomplete());
-        }
-
         let handshake = Handshake {
             header: Header {
                 msg_type: first_handshake.header.msg_type,
@@ -287,6 +320,12 @@ impl Handshake {
     }
 }
 
+fn mark_handled(handled: ArrayVec<&Handshake, MAX_DEFRAGMENT_HANDSHAKES>) {
+    for handshake in handled {
+        handshake.set_handled();
+    }
+}
+
 #[repr(transparent)]
 #[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
 pub struct MessageType(u8);
@@ -318,6 +357,10 @@ impl MessageType {
         !matches!(*self, Self(0..=4 | 11..=16 | 20))
     }
 
+    const fn rejects_trailing_body_bytes(&self) -> bool {
+        !self.is_unknown()
+    }
+
     pub fn parse(input: &[u8]) -> IResult<&[u8], MessageType> {
         let (input, byte) = be_u8(input)?;
         Ok((input, Self::from_u8(byte)))
@@ -695,4 +738,59 @@ mod tests {
 
         assert!(rest.is_empty());
     }
+
+    #[test]
+    fn failed_defragment_parse_discards_candidate_without_writing_transcript() {
+        let mut body = MESSAGE[12..].to_vec();
+        body[38] = 0;
+        body[39] = 3;
+
+        let handshake = Handshake::new(
+            MessageType::ClientHello,
+            body.len() as u32,
+            0,
+            0,
+            body.len() as u32,
+            Body::Fragment(0..body.len()),
+        );
+
+        let mut defragmented_buffer = Buf::new();
+        let mut transcript = Buf::new();
+        let result = Handshake::defragment(
+            std::iter::once((&handshake, body.as_slice())),
+            &mut defragmented_buffer,
+            None,
+            Some(&mut transcript),
+        );
+
+        assert!(result.is_err());
+        assert!(handshake.is_handled());
+        assert!(transcript.is_empty());
+    }
+
+    #[test]
+    fn known_body_rejects_trailing_bytes() {
+        let body = [0];
+        let handshake = Handshake::new(
+            MessageType::HelloRequest,
+            body.len() as u32,
+            0,
+            0,
+            body.len() as u32,
+            Body::Fragment(0..body.len()),
+        );
+
+        let mut defragmented_buffer = Buf::new();
+        let mut transcript = Buf::new();
+        let result = Handshake::defragment(
+            std::iter::once((&handshake, body.as_slice())),
+            &mut defragmented_buffer,
+            None,
+            Some(&mut transcript),
+        );
+
+        assert!(result.is_err());
+        assert!(handshake.is_handled());
+        assert!(transcript.is_empty());
+    }
 }

src/dtls13/message/handshake.rs

@@ -5,13 +5,20 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use super::{Certificate, CertificateVerify, ClientHello, Dtls13CipherSuite};
 use super::{EncryptedExtensions, Finished, ServerHello};
 use crate::buffer::Buf;
+use arrayvec::ArrayVec;
 use nom::Err;
 use nom::IResult;
 use nom::bytes::complete::take;
 use nom::error::{Error, ErrorKind};
 use nom::number::complete::be_u8;
 use nom::number::complete::{be_u16, be_u24};
 
+// Defensive stack cap over flattened handshake fragments selected for one
+// defragmentation attempt. This intentionally does not mirror the receive
+// queue's record-count cap; exceeding 50 fragments for one handshake implies
+// pathologically tiny records and is treated as invalid input.
+const MAX_DEFRAGMENT_HANDSHAKES: usize = 50;
+
 #[derive(Debug, PartialEq, Eq, Default, Clone, Copy)]
 pub struct Header {
     pub msg_type: MessageType,
@@ -165,8 +172,11 @@ impl Handshake {
         let Body::Fragment(range) = &first_handshake.body else {
             unreachable!("Non-Fragment body in defragment()")
         };
+        let mut handled = ArrayVec::<&Handshake, MAX_DEFRAGMENT_HANDSHAKES>::new();
+        handled
+            .try_push(first_handshake)
+            .map_err(|_| crate::InternalError::too_many_records())?;
         buffer.extend_from_slice(&first_buffer[range.clone()]);
-        first_handshake.set_handled();
 
         let mut assembled_end =
             first_handshake.header.fragment_offset + first_handshake.header.fragment_length;
@@ -180,7 +190,9 @@ impl Handshake {
                 unreachable!("Non-Fragment body in defragment()")
             };
 
-            handshake.handled.store(true, Ordering::Relaxed);
+            handled
+                .try_push(handshake)
+                .map_err(|_| crate::InternalError::too_many_records())?;
 
             // Handle overlapping fragment data: skip bytes already assembled
             let frag_start = handshake.header.fragment_offset as usize;
@@ -200,31 +212,57 @@ impl Handshake {
             return Err(crate::InternalError::parse_incomplete());
         }
 
-        // If transcript is provided, write the TLS 1.3-style header + body.
-        // Per RFC 9147 Section 5.2, the transcript uses msg_type(1) + length(3)
-        // WITHOUT the DTLS-specific message_seq, fragment_offset, fragment_length.
-        if let Some(transcript) = transcript {
-            transcript.push(first_handshake.header.msg_type.as_u8());
-            transcript.extend_from_slice(&first_handshake.header.length.to_be_bytes()[1..]);
-            transcript.extend_from_slice(&buffer[..first_handshake.header.length as usize]);
-        }
-
         let (rest, body) = if allow_unknown_client_hello_suites {
-            Body::parse_allow_unknown_client_hello_suites(
+            match Body::parse_allow_unknown_client_hello_suites(
                 buffer,
                 0,
                 first_handshake.header.msg_type,
                 cipher_suite,
-            )?
+            ) {
+                Ok(parsed) => parsed,
+                Err(err) => {
+                    mark_handled(handled);
+                    return Err(err.into());
+                }
+            }
         } else {
-            Body::parse(buffer, 0, first_handshake.header.msg_type, cipher_suite)?
+            match Body::parse(buffer, 0, first_handshake.header.msg_type, cipher_suite) {
+                Ok(parsed) => parsed,
+                Err(err) => {
+                    mark_handled(handled);
+                    return Err(err.into());
+                }
+            }
         };
 
-        if !rest.is_empty() && first_handshake.header.msg_type == MessageType::Finished {
+        if !rest.is_empty()
+            && first_handshake
+                .header
+                .msg_type
+                .rejects_trailing_body_bytes()
+        {
             debug!("Defragmentation failed. Body::parse() did not consume the entire buffer");
+            mark_handled(handled);
             return Err(crate::InternalError::parse_incomplete());
         }
 
+        // Intentional boundary: Body::parse validates the handshake body shape and
+        // extension envelopes, but known extension payloads remain validated by the
+        // client/server state handlers. A transiently corrupted UDP datagram whose
+        // extension payload fails later may therefore have been consumed here; that
+        // recovery edge is accepted to keep this path parser-only and avoid the
+        // broader transaction/rollback machinery.
+        mark_handled(handled);
+
+        // If transcript is provided, write the TLS 1.3-style header + body after parsing succeeds.
+        // Per RFC 9147 Section 5.2, the transcript uses msg_type(1) + length(3)
+        // WITHOUT the DTLS-specific message_seq, fragment_offset, fragment_length.
+        if let Some(transcript) = transcript {
+            transcript.push(first_handshake.header.msg_type.as_u8());
+            transcript.extend_from_slice(&first_handshake.header.length.to_be_bytes()[1..]);
+            transcript.extend_from_slice(&buffer[..first_handshake.header.length as usize]);
+        }
+
         let handshake = Handshake {
             header: Header {
                 msg_type: first_handshake.header.msg_type,
@@ -268,6 +306,12 @@ impl Handshake {
     }
 }
 
+fn mark_handled(handled: ArrayVec<&Handshake, MAX_DEFRAGMENT_HANDSHAKES>) {
+    for handshake in handled {
+        handshake.set_handled();
+    }
+}
+
 #[repr(transparent)]
 #[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
 pub struct MessageType(u8);
@@ -295,6 +339,10 @@ impl MessageType {
         !matches!(*self, Self(1..=2 | 8 | 11 | 13 | 15 | 20 | 24))
     }
 
+    const fn rejects_trailing_body_bytes(&self) -> bool {
+        !self.is_unknown()
+    }
+
     pub fn parse(input: &[u8]) -> IResult<&[u8], MessageType> {
         let (input, byte) = be_u8(input)?;
         Ok((input, Self::from_u8(byte)))
@@ -620,16 +668,45 @@ mod tests {
         );
 
         let mut buffer = Buf::new();
+        let mut transcript = Buf::new();
         let result = Handshake::defragment(
             std::iter::once((&handshake, source.as_slice())),
             &mut buffer,
             None,
-            None,
+            Some(&mut transcript),
         );
 
         assert!(
             result.is_err(),
             "KeyUpdate bodies with trailing bytes must be rejected"
         );
+        assert!(handshake.is_handled());
+        assert!(transcript.is_empty());
+    }
+
+    #[test]
+    fn known_body_rejects_trailing_bytes() {
+        let source = [0, 0, 0];
+        let handshake = Handshake::new(
+            MessageType::EncryptedExtensions,
+            source.len() as u32,
+            0,
+            0,
+            source.len() as u32,
+            Body::Fragment(0..source.len()),
+        );
+
+        let mut buffer = Buf::new();
+        let mut transcript = Buf::new();
+        let result = Handshake::defragment(
+            std::iter::once((&handshake, source.as_slice())),
+            &mut buffer,
+            None,
+            Some(&mut transcript),
+        );
+
+        assert!(result.is_err());
+        assert!(handshake.is_handled());
+        assert!(transcript.is_empty());
     }
 }