add fragment_buffer and test

rainliu · rainliu · commit f1def083fbf6 · 2020-10-31T22:49:34.000-08:00
update code
diff --git a/dtls/src/fragment_buffer.rs b/dtls/src/fragment_buffer.rs
@@ -0,0 +1,140 @@
+#[cfg(test)]
+mod fragment_buffer_test;
+
+use crate::content::*;
+use crate::handshake::handshake_header::*;
+use crate::record_layer::record_layer_header::*;
+
+use util::Error;
+
+use std::collections::HashMap;
+use std::io::{BufWriter, Cursor};
+
+pub(crate) struct Fragment {
+    record_layer_header: RecordLayerHeader,
+    handshake_header: HandshakeHeader,
+    data: Vec<u8>,
+}
+
+pub(crate) struct FragmentBuffer {
+    // map of MessageSequenceNumbers that hold slices of fragments
+    cache: HashMap<u16, Vec<Fragment>>,
+
+    current_message_sequence_number: u16,
+}
+
+impl FragmentBuffer {
+    pub fn new() -> Self {
+        FragmentBuffer {
+            cache: HashMap::new(),
+            current_message_sequence_number: 0,
+        }
+    }
+
+    // Attempts to push a DTLS packet to the 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, Error> {
+        let mut reader = Cursor::new(buf);
+        let record_layer_header = RecordLayerHeader::unmarshal(&mut reader)?;
+
+        // Fragment isn't a handshake, we don't need to handle it
+        if record_layer_header.content_type != ContentType::Handshake {
+            return Ok(false);
+        }
+
+        buf = &buf[RECORD_LAYER_HEADER_SIZE..];
+        while buf.len() != 0 {
+            let mut reader = Cursor::new(buf);
+            let handshake_header = HandshakeHeader::unmarshal(&mut reader)?;
+
+            if !self.cache.contains_key(&handshake_header.message_sequence) {
+                self.cache.insert(handshake_header.message_sequence, vec![]);
+            }
+
+            // end index should be the length of handshake header but if the handshake
+            // was fragmented, we should keep them all
+            let mut end = HANDSHAKE_HEADER_LENGTH + handshake_header.length as usize;
+            if end > buf.len() {
+                end = buf.len();
+            }
+
+            // Discard all headers, when rebuilding the packet we will re-build
+            let data = buf[HANDSHAKE_HEADER_LENGTH..end].to_vec();
+
+            if let Some(x) = self.cache.get_mut(&handshake_header.message_sequence) {
+                x.push(Fragment {
+                    record_layer_header,
+                    handshake_header,
+                    data,
+                });
+            }
+            buf = &buf[end..];
+        }
+
+        Ok(true)
+    }
+
+    pub fn pop(&mut self) -> (Vec<u8>, u16) {
+        let seq_num = self.current_message_sequence_number;
+        if !self.cache.contains_key(&seq_num) {
+            return (vec![], 0);
+        }
+
+        let (content, epoch) = if let Some(frags) = self.cache.get_mut(&seq_num) {
+            let mut raw_message = vec![];
+            // Recursively collect up
+            if !append_message(0, frags, &mut raw_message) {
+                return (vec![], 0);
+            }
+
+            let mut first_header = frags[0].handshake_header;
+            first_header.fragment_offset = 0;
+            first_header.fragment_length = first_header.length;
+
+            let mut raw_header = vec![];
+            {
+                let mut writer = BufWriter::<&mut Vec<u8>>::new(raw_header.as_mut());
+                match first_header.marshal(&mut writer) {
+                    Err(_) => return (vec![], 0),
+                    Ok(_) => {}
+                };
+            }
+
+            let message_epoch = frags[0].record_layer_header.epoch;
+
+            raw_header.extend_from_slice(&raw_message);
+
+            (raw_header, message_epoch)
+        } else {
+            (vec![], 0)
+        };
+
+        self.cache.remove(&seq_num);
+        self.current_message_sequence_number += 1;
+
+        (content, epoch)
+    }
+}
+
+fn append_message(target_offset: u32, frags: &[Fragment], raw_message: &mut Vec<u8>) -> bool {
+    for f in frags {
+        if f.handshake_header.fragment_offset == target_offset {
+            let fragment_end =
+                f.handshake_header.fragment_offset + f.handshake_header.fragment_length;
+            if fragment_end != f.handshake_header.length {
+                if !append_message(fragment_end, frags, raw_message) {
+                    return false;
+                }
+            }
+
+            let mut message = vec![];
+            message.extend_from_slice(&f.data);
+            message.extend_from_slice(raw_message);
+            *raw_message = message;
+            return true;
+        }
+    }
+
+    false
+}
diff --git a/dtls/src/fragment_buffer/fragment_buffer_test.rs b/dtls/src/fragment_buffer/fragment_buffer_test.rs
@@ -0,0 +1,148 @@
+use super::*;
+
+use util::Error;
+
+#[test]
+fn test_fragment_buffer() -> Result<(), Error> {
+    let tests = vec![
+        (
+            "Single Fragment",
+            vec![vec![
+                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,
+            ]],
+            vec![vec![
+                0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff,
+                0x00,
+            ]],
+            0,
+        ),
+        (
+            "Single Fragment Epoch 3",
+            vec![vec![
+                0x16, 0xfe, 0xff, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03,
+                0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00,
+            ]],
+            vec![vec![
+                0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff,
+                0x00,
+            ]],
+            3,
+        ),
+        (
+            "Multiple Fragments",
+            vec![
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00,
+                    0x01, 0x02, 0x03, 0x04,
+                ],
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x05, 0x05,
+                    0x06, 0x07, 0x08, 0x09,
+                ],
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x05, 0x0A,
+                    0x0B, 0x0C, 0x0D, 0x0E,
+                ],
+            ],
+            vec![vec![
+                0x0b, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x01,
+                0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
+            ]],
+            0,
+        ),
+        (
+            "Multiple Unordered Fragments",
+            vec![
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00,
+                    0x01, 0x02, 0x03, 0x04,
+                ],
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x05, 0x0A,
+                    0x0B, 0x0C, 0x0D, 0x0E,
+                ],
+                vec![
+                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x81,
+                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x05, 0x05,
+                    0x06, 0x07, 0x08, 0x09,
+                ],
+            ],
+            vec![vec![
+                0x0b, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x01,
+                0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
+            ]],
+            0,
+        ),
+        (
+            "Multiple Handshakes in Signle Fragment",
+            vec![vec![
+                0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
+                0x30, /* record header */
+                0x03, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
+                0x01, 0x01, /*handshake msg 1*/
+                0x03, 0x00, 0x00, 0x04, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
+                0x01, 0x01, /*handshake msg 2*/
+                0x03, 0x00, 0x00, 0x04, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
+                0x01, 0x01, /*handshake msg 3*/
+            ]],
+            vec![
+                vec![
+                    0x03, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
+                    0xff, 0x01, 0x01,
+                ],
+                vec![
+                    0x03, 0x00, 0x00, 0x04, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
+                    0xff, 0x01, 0x01,
+                ],
+                vec![
+                    0x03, 0x00, 0x00, 0x04, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
+                    0xff, 0x01, 0x01,
+                ],
+            ],
+            0,
+        ),
+    ];
+
+    for (name, inputs, expects, expected_epoch) in tests {
+        let mut fragment_buffer = FragmentBuffer::new();
+        for frag in inputs {
+            let status = fragment_buffer.push(&frag)?;
+            assert_eq!(
+                true, status,
+                "fragment_buffer didn't accept fragments for '{}'",
+                name
+            );
+        }
+
+        for expected in expects {
+            let (out, epoch) = fragment_buffer.pop();
+            assert_eq!(
+                out, expected,
+                "fragment_buffer '{}' push/pop: got {:?}, want {:?}",
+                name, out, expected
+            );
+
+            assert_eq!(
+                epoch, expected_epoch,
+                "fragment_buffer returned wrong epoch: got {}, want {}",
+                epoch, expected_epoch
+            );
+        }
+
+        let (frag, _) = fragment_buffer.pop();
+        assert_eq!(
+            frag.len(),
+            0,
+            "fragment_buffer popped single buffer multiple times for '{}'",
+            name
+        );
+    }
+
+    Ok(())
+}
diff --git a/dtls/src/handshake.rs b/dtls/src/handshake.rs
@@ -1,15 +1,15 @@
-mod handshake_header;
-mod handshake_message_certificate;
-mod handshake_message_certificate_request;
-mod handshake_message_certificate_verify;
-mod handshake_message_client_hello;
-mod handshake_message_client_key_exchange;
-mod handshake_message_finished;
-mod handshake_message_hello_verify_request;
-mod handshake_message_server_hello;
-mod handshake_message_server_hello_done;
-mod handshake_message_server_key_exchange;
-mod handshake_random;
+pub mod handshake_header;
+pub mod handshake_message_certificate;
+pub mod handshake_message_certificate_request;
+pub mod handshake_message_certificate_verify;
+pub mod handshake_message_client_hello;
+pub mod handshake_message_client_key_exchange;
+pub mod handshake_message_finished;
+pub mod handshake_message_hello_verify_request;
+pub mod handshake_message_server_hello;
+pub mod handshake_message_server_hello_done;
+pub mod handshake_message_server_key_exchange;
+pub mod handshake_random;
 
 use std::fmt;
 use std::io::{Read, Write};
diff --git a/dtls/src/handshake/handshake_header.rs b/dtls/src/handshake/handshake_header.rs
@@ -7,16 +7,16 @@ use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use util::Error;
 
 // msg_len for Handshake messages assumes an extra 12 bytes for
-// sequence, fragment and version information
-const HANDSHAKE_HEADER_LENGTH: usize = 12;
+// sequence, Fragment and version information
+pub(crate) const HANDSHAKE_HEADER_LENGTH: usize = 12;
 
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Copy, Clone, PartialEq, Debug)]
 pub struct HandshakeHeader {
     pub handshake_type: HandshakeType,
-    length: u32, // uint24 in spec
-    message_sequence: u16,
-    fragment_offset: u32, // uint24 in spec
-    fragment_length: u32, // uint24 in spec
+    pub(crate) length: u32, // uint24 in spec
+    pub message_sequence: u16,
+    pub(crate) fragment_offset: u32, // uint24 in spec
+    pub(crate) fragment_length: u32, // uint24 in spec
 }
 
 impl HandshakeHeader {
diff --git a/dtls/src/lib.rs b/dtls/src/lib.rs
@@ -16,7 +16,9 @@ pub mod crypto;
 pub mod curve;
 pub mod errors;
 pub mod extension;
+pub mod fragment_buffer;
 pub mod handshake;
 pub mod prf;
 pub mod record_layer;
 pub mod signature_hash_algorithm;
+pub mod state;
diff --git a/dtls/src/state.rs b/dtls/src/state.rs
