Merge pull request #17 from k0nserv/fix-fragment-buffer-max-size

Backport buffer limit fix from Pion

Author: rainliu

crates/dtls/src/error.rs

@@ -144,6 +144,11 @@ pub enum Error {
     #[error("Alert is Fatal or Close Notify")]
     ErrAlertFatalOrClose,
 
+    #[error(
+        "Fragment buffer overflow. New size {new_size} is greater than specified max {max_size}"
+    )]
+    ErrFragmentBufferOverflow { new_size: usize, max_size: usize },
+
     #[error("{0}")]
     Io(#[source] IoError),
     #[error("{0}")]

crates/dtls/src/fragment_buffer/fragment_buffer_test.rs

@@ -143,3 +143,23 @@ fn test_fragment_buffer() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_fragment_buffer_overflow() -> Result<()> {
+    let mut fragment_buffer = FragmentBuffer::new();
+
+    fragment_buffer.push(&[
+        0x16, 0xfe, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03, 0x00,
+        0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00,
+    ])?;
+
+    let big_buffer = vec![0; 2_000_000];
+    let result = fragment_buffer.push(&big_buffer);
+
+    assert!(
+        result.is_err(),
+        "Pushing a buffer of size 2MB should have caused FragmentBuffer::push to return an error"
+    );
+
+    Ok(())
+}

crates/dtls/src/fragment_buffer/mod.rs

@@ -9,6 +9,9 @@ use crate::record_layer::record_layer_header::*;
 use std::collections::HashMap;
 use std::io::{BufWriter, Cursor};
 
+// 2 mb max buffer size
+const FRAGMENT_BUFFER_MAX_SIZE: usize = 2_000_000;
+
 pub(crate) struct Fragment {
     record_layer_header: RecordLayerHeader,
     handshake_header: HandshakeHeader,
@@ -34,6 +37,14 @@ impl FragmentBuffer {
     // when it returns true it means the FragmentBuffer has inserted and the buffer shouldn't be handled
     // when an error returns it is fatal, and the DTLS connection should be stopped
     pub fn push(&mut self, mut buf: &[u8]) -> Result<bool> {
+        let current_size = self.size();
+        if current_size + buf.len() >= FRAGMENT_BUFFER_MAX_SIZE {
+            return Err(Error::ErrFragmentBufferOverflow {
+                new_size: current_size + buf.len(),
+                max_size: FRAGMENT_BUFFER_MAX_SIZE,
+            });
+        }
+
         let mut reader = Cursor::new(buf);
         let record_layer_header = RecordLayerHeader::unmarshal(&mut reader)?;
 
@@ -113,6 +124,13 @@ impl FragmentBuffer {
 
         Ok((content, epoch))
     }
+
+    fn size(&self) -> usize {
+        self.cache
+            .iter()
+            .map(|(_, fragment)| fragment.iter().map(|f| f.data.len()).sum::<usize>())
+            .sum()
+    }
 }
 
 fn append_message(target_offset: u32, frags: &[Fragment], raw_message: &mut Vec<u8>) -> bool {