feat: message type validation in ZeroCopyParser

Author: sundanc

src/zerocopy.rs

@@ -361,54 +361,73 @@ impl<'a> ZeroCopyParser<'a> {
 
     #[inline(always)]
     pub fn parse_next(&mut self) -> Option<ZeroCopyMessage<'a>> {
-        let data_len = self.data.len();
-        let pos = self.position;
+        loop {
+            let data_len = self.data.len();
+            let pos = self.position;
 
-        if pos + 3 > data_len {
-            return None;
-        }
+            if pos + 3 > data_len {
+                return None;
+            }
 
-        let length = unsafe {
-            let ptr = self.data.as_ptr().add(pos);
-            u16::from_be_bytes(std::ptr::read_unaligned(ptr as *const [u8; 2]))
-        } as usize;
+            let length = unsafe {
+                let ptr = self.data.as_ptr().add(pos);
+                u16::from_be_bytes(std::ptr::read_unaligned(ptr as *const [u8; 2]))
+            } as usize;
 
-        let total_size = length + 2;
-        let msg_end = pos + total_size;
+            let total_size = length + 2;
+            let msg_end = pos + total_size;
 
-        if msg_end > data_len {
-            return None;
-        }
+            if msg_end <= pos {
+                self.position = pos + 1;
+                continue;
+            }
 
-        #[cfg(all(feature = "simd", target_arch = "x86_64"))]
-        {
-            if msg_end + 64 <= data_len {
-                unsafe {
-                    use std::arch::x86_64::*;
-                    _mm_prefetch(self.data.as_ptr().add(msg_end) as *const i8, _MM_HINT_T0);
+            if msg_end > data_len {
+                return None;
+            }
+
+            #[cfg(all(feature = "simd", target_arch = "x86_64"))]
+            {
+                if msg_end + 64 <= data_len {
+                    unsafe {
+                        use std::arch::x86_64::*;
+                        _mm_prefetch(self.data.as_ptr().add(msg_end) as *const i8, _MM_HINT_T0);
+                    }
                 }
             }
-        }
 
-        let msg_type = self.data[pos + 2];
-        let header_start = pos + 3;
-        let header_end = header_start + 10;
+            let msg_type = self.data[pos + 2];
 
-        if header_end > data_len {
-            return None;
-        }
+            if !crate::simd::is_valid_message_type(msg_type) {
+                self.position = msg_end;
+                continue;
+            }
 
-        let payload_start = header_end;
+            let header_start = pos + 3;
+            let header_end = header_start + 10;
+
+            if header_end > msg_end {
+                self.position = msg_end;
+                continue;
+            }
+
+            if header_end > data_len {
+                return None;
+            }
+
+            let payload_start = header_end;
+
+            if let Ok((hdr_ref, _)) =
+                Ref::<&[u8], MessageHeaderRaw>::from_prefix(&self.data[header_start..header_end])
+            {
+                let payload = &self.data[payload_start..msg_end];
+                self.position = msg_end;
+                return Some(ZeroCopyMessage::new(msg_type, hdr_ref, payload));
+            }
 
-        if let Ok((hdr_ref, _)) =
-            Ref::<&[u8], MessageHeaderRaw>::from_prefix(&self.data[header_start..header_end])
-        {
-            let payload = &self.data[payload_start..msg_end];
             self.position = msg_end;
-            return Some(ZeroCopyMessage::new(msg_type, hdr_ref, payload));
+            continue;
         }
-
-        None
     }
 
     #[inline]
@@ -741,14 +760,14 @@ mod tests {
         let mut data = Vec::new();
         for _ in 0..2 {
             data.extend(&[0u8, 11u8]);
-            data.push(1u8);
+            data.push(b'S');
             data.extend(&[0u8; 10]);
         }
 
         let mut parser = ZeroCopyParser::new(&data);
         let arcs = parser.parse_all_arc();
         assert_eq!(arcs.len(), 2);
-        assert_eq!(arcs[0].msg_type, 1u8);
+        assert_eq!(arcs[0].msg_type, b'S');
         assert!(Arc::strong_count(&arcs[0].payload) >= 1);
     }
 }

tests/fixtures.rs

@@ -16,6 +16,7 @@ pub fn create_uniform_buffer(count: usize, msg_type: u8, payload_len: usize) ->
     create_test_buffer(&messages)
 }
 
+#[allow(dead_code)]
 pub fn standard_fixture() -> Vec<u8> {
     let mut messages = Vec::new();
     for i in 0..1000u16 {

tests/integration.rs

@@ -197,3 +197,35 @@ fn test_large_buffer_no_data_loss() {
         "All messages should be parsed without data loss"
     );
 }
+
+#[test]
+fn test_skip_invalid_message_type() {
+    let messages = vec![
+        (0xFF, &[][..]),      // invalid
+        (b'S', &[0; 10][..]), // valid
+    ];
+    let buf = fixtures::create_test_buffer(&messages);
+    let mut parser = ZeroCopyParser::new(&buf);
+    let msg = parser.parse_next().unwrap();
+    assert_eq!(msg.msg_type(), b'S');
+    assert_eq!(parser.position(), buf.len());
+    assert!(parser.parse_next().is_none());
+}
+
+#[test]
+fn test_skip_invalid_with_while_loop() {
+    let messages = vec![(0xFF, &[][..]), (b'S', &[0; 10][..]), (b'R', &[1; 5][..])];
+    let buf = fixtures::create_test_buffer(&messages);
+    let mut parser = ZeroCopyParser::new(&buf);
+    let mut count = 0;
+    while let Some(msg) = parser.parse_next() {
+        count += 1;
+        if count == 1 {
+            assert_eq!(msg.msg_type(), b'S');
+        } else {
+            panic!("Unexpected extra message with type {}", msg.msg_type());
+        }
+    }
+    assert_eq!(count, 1);
+    assert_eq!(parser.position(), buf.len());
+}

tests/mmap_shared.rs

@@ -9,7 +9,7 @@ fn test_mmap_shared_and_owned_messages() {
     let mut buf = Vec::new();
     for _ in 0..2 {
         buf.extend(&[0, 11]);
-        buf.push(1u8);
+        buf.push(b'S');
         buf.extend(&[0u8; 10]);
     }
 
@@ -23,5 +23,5 @@ fn test_mmap_shared_and_owned_messages() {
     assert_eq!(owned.len(), 2);
 
     drop(shared);
-    assert_eq!(owned[0].msg_type, 1u8);
+    assert_eq!(owned[0].msg_type, b'S');
 }

tests/zerocopy_validation.rs

@@ -0,0 +1,37 @@
+mod fixtures;
+
+use lunary::ZeroCopyParser;
+
+use fixtures::create_test_buffer;
+
+#[test]
+fn test_skip_invalid_message_type() {
+    let messages = vec![
+        (0xFF, &[][..]),      // invalid
+        (b'S', &[0; 10][..]), // valid
+    ];
+    let buf = create_test_buffer(&messages);
+    let mut parser = ZeroCopyParser::new(&buf);
+    let msg = parser.parse_next().unwrap();
+    assert_eq!(msg.msg_type(), b'S');
+    assert_eq!(parser.position(), buf.len());
+    assert!(parser.parse_next().is_none());
+}
+
+#[test]
+fn test_skip_invalid_with_while_loop() {
+    let messages = vec![(0xFF, &[][..]), (b'S', &[0; 10][..]), (b'R', &[1; 5][..])];
+    let buf = create_test_buffer(&messages);
+    let mut parser = ZeroCopyParser::new(&buf);
+    let mut count = 0;
+    while let Some(msg) = parser.parse_next() {
+        count += 1;
+        if count == 1 {
+            assert_eq!(msg.msg_type(), b'S');
+        } else {
+            panic!("Unexpected extra message with type {}", msg.msg_type());
+        }
+    }
+    assert_eq!(count, 1);
+    assert_eq!(parser.position(), buf.len());
+}