Support ext formats in MessageLen and report limit violations as ParseError

rmp/src/decode/est.rs

@@ -54,7 +54,7 @@ impl MessageLen {
 
     /// * `max_depth` limits nesting of arrays and maps
     ///
-    /// * `max_len` is maximum size of any string, byte string, map, or array.
+    /// * `max_len` is maximum size of any string, byte string, ext payload, map, or array.
     ///   For maps and arrays this is the number of items, not bytes.
     ///
     /// Messages can be both deep and wide, being `max_depth` * `max_len` in size.
@@ -82,7 +82,8 @@ impl MessageLen {
     /// and it would need *at least* this many bytes to parse.
     /// The `len` is always the lower bound, and never exceeds actual message length.
     ///
-    /// `Err(LenError::ParseError)` — the end of the message is unknown.
+    /// `Err(LenError::ParseError)` — the message is malformed or exceeds the limits, so the end
+    /// of the message cannot be determined.
     ///
     /// Don't call this function in a loop. Use [`MessageLen::incremental_len`] instead.
     pub fn len_of(complete_message: &[u8]) -> Result<usize, LenError> {
@@ -104,8 +105,9 @@ impl MessageLen {
     ///   The `len` is always the lower bound, and never exceeds actual message length,
     ///   so it's safe to read the additional bytes without overshooting the end of the message.
     ///
-    /// * `Err(LenError::ParseError)` — the end of the message cannot be determined, and this
-    ///   is a non-recoverable error. Any further calls to this function may return nonsense.
+    /// * `Err(LenError::ParseError)` — the message is malformed or exceeds the limits, so the end
+    ///   of the message cannot be determined. This is a non-recoverable error: any further calls
+    ///   to this function will keep returning it until [`MessageLen::reset`] is called.
     pub fn incremental_len(&mut self, mut next_message_fragment: &[u8]) -> Result<usize, LenError> {
         let data = &mut next_message_fragment;
         let Some(wip) = self.wip.take() else {
@@ -115,16 +117,16 @@ impl MessageLen {
             WIP::Data(Data { bytes_left }) => self.skip_data(data, bytes_left.get()),
             WIP::MarkerLen(wip) => self.read_marker_with_len(data, wip),
             WIP::NextMarker => self.read_one_item(data),
-            WIP::LimitExceeded => {
-                self.wip = Some(WIP::LimitExceeded); // put it back!
+            kind @ (WIP::LimitExceeded | WIP::Reserved) => {
+                self.wip = Some(kind); // put it back!
                 return Err(LenError::ParseError);
             },
-        }.ok_or(LenError::Truncated(self.max_position))?;
+        }.ok_or_else(|| self.interrupted())?;
 
         while let Some(seq) = self.sequences_wip.pop() {
             self.current_depth = seq.depth;
             debug_assert!(self.wip.is_none());
-            self.read_sequence(data, seq.items_left.get() - 1).ok_or(LenError::Truncated(self.max_position))?;
+            self.read_sequence(data, seq.items_left.get() - 1).ok_or_else(|| self.interrupted())?;
         }
         debug_assert!(self.wip.is_none());
         debug_assert!(self.max_position.get() <= self.position);
@@ -149,9 +151,10 @@ impl MessageLen {
             Marker::FixArray(len) => self.read_sequence(data, u32::from(len)),
             Marker::FixStr(len) => self.skip_data(data, len.into()),
             Marker::Null |
-            Marker::Reserved |
             Marker::False |
             Marker::True => Some(()),
+            // 0xc1 is not a valid marker
+            Marker::Reserved => self.fail(WIP::Reserved),
             Marker::Str8 |
             Marker::Str16 |
             Marker::Str32 |
@@ -161,10 +164,10 @@ impl MessageLen {
             Marker::Array16 |
             Marker::Array32 |
             Marker::Map16 |
-            Marker::Map32 => self.read_marker_with_len(data, MarkerLen { marker, buf: [0; 4], has: 0 }),
+            Marker::Map32 |
             Marker::Ext8 |
             Marker::Ext16 |
-            Marker::Ext32 => todo!(),
+            Marker::Ext32 => self.read_marker_with_len(data, MarkerLen { marker, buf: [0; 4], has: 0 }),
             Marker::F32 => self.skip_data(data, 4),
             Marker::F64 => self.skip_data(data, 8),
             Marker::U8 => self.skip_data(data, 1),
@@ -175,11 +178,11 @@ impl MessageLen {
             Marker::I16 => self.skip_data(data, 2),
             Marker::I32 => self.skip_data(data, 4),
             Marker::I64 => self.skip_data(data, 8),
-            Marker::FixExt1 |
-            Marker::FixExt2 |
-            Marker::FixExt4 |
-            Marker::FixExt8 |
-            Marker::FixExt16 => todo!(),
+            Marker::FixExt1 => self.skip_data(data, 2),
+            Marker::FixExt2 => self.skip_data(data, 3),
+            Marker::FixExt4 => self.skip_data(data, 5),
+            Marker::FixExt8 => self.skip_data(data, 9),
+            Marker::FixExt16 => self.skip_data(data, 17),
             Marker::FixNeg(_) => Some(()),
         }
     }
@@ -217,7 +220,10 @@ impl MessageLen {
             Marker::Str32 => self.skip_data(data, len),
             Marker::Ext8 |
             Marker::Ext16 |
-            Marker::Ext32 => todo!(),
+            Marker::Ext32 => {
+                // type byte + payload; `len < max_len <= u32::MAX`, so this can't overflow
+                self.skip_data(data, len + 1)
+            },
             Marker::Array16 |
             Marker::Array32 => self.read_sequence(data, len),
             Marker::Map16 |
@@ -296,19 +302,29 @@ impl MessageLen {
         let pos = match self.wip.insert(wip) {
             WIP::NextMarker => self.position + 1,
             WIP::Data(Data { bytes_left }) => self.position + bytes_left.get() as usize,
-            WIP::MarkerLen(m) => self.position + (m.size() - m.has) as usize,
-            WIP::LimitExceeded => 0,
+            WIP::MarkerLen(m) => self.position + (m.size() - m.has) as usize + m.min_trailing(),
+            WIP::LimitExceeded | WIP::Reserved => 0,
         };
         self.set_max_position(pos);
         None
     }
+
+    /// The error to report when an operation was interrupted (returned `None`)
+    fn interrupted(&self) -> LenError {
+        match self.wip {
+            Some(WIP::LimitExceeded | WIP::Reserved) => LenError::ParseError,
+            _ => LenError::Truncated(self.max_position),
+        }
+    }
 }
 
 enum WIP {
     NextMarker,
     Data(Data),
     MarkerLen(MarkerLen),
     LimitExceeded,
+    /// The message contains the reserved marker `0xc1`. Sticky, like `LimitExceeded`.
+    Reserved,
 }
 
 struct Seq { items_left: NonZeroU32, depth: u16 }
@@ -334,4 +350,13 @@ impl MarkerLen {
             _ => unimplemented!(),
         }
     }
+
+    /// Number of bytes that must follow the length, regardless of its value
+    fn min_trailing(&self) -> usize {
+        match self.marker {
+            // the type byte
+            Marker::Ext8 | Marker::Ext16 | Marker::Ext32 => 1,
+            _ => 0,
+        }
+    }
 }

rmp/tests/func/est.rs

@@ -115,3 +115,161 @@ fn nested() {
     assert!(MessageLen::with_limits(4, 1 << 16).incremental_len(out.as_slice()).is_err());
     assert!(MessageLen::with_limits(14, 1 << 16).incremental_len(out.as_slice()).is_ok());
 }
+
+#[test]
+fn extensions() {
+    let mut out = Vec::with_capacity(263);
+
+    let mut expected = Vec::with_capacity(264);
+    expected.push(1);
+
+    for len in (0i32..=17).chain([255, 256, 257, 65535]) {
+        const TOO_BIG_FOR_U8: i32 = u8::MAX as i32 + 1;
+        const TOO_BIG_FOR_U16: i32 = u16::MAX as i32 + 1;
+
+        let length_bytes = match len {
+            1 | 2 | 4 | 8 | 16 => 0,
+            ..TOO_BIG_FOR_U8 => 1,
+            TOO_BIG_FOR_U8..TOO_BIG_FOR_U16 => 2,
+            _ => 4,
+        };
+
+        let len_with_prefix = |len| 1 + length_bytes + 1 + len;
+        let msg_len = len_with_prefix(len);
+
+        expected.truncate(1);
+        // while the length is being read, the type byte is already known to follow it
+        expected.resize(1 + length_bytes as usize, 2 + length_bytes);
+        expected.resize(msg_len as usize, msg_len);
+        expected.push(-msg_len);
+
+        out.clear();
+        write_ext_meta(&mut out, len as u32, 0x67).unwrap();
+        out.resize(out.len() + len as usize, 0xab);
+
+        check_estimates(&out, &expected);
+    }
+}
+
+#[test]
+fn ext32() {
+    // Non-canonical but valid: ext32 marker with a small length, so the estimates can be
+    // checked exhaustively for every prefix.
+    let mut out = vec![Marker::Ext32.to_u8()];
+    out.extend_from_slice(&3u32.to_be_bytes());
+    out.push(7);
+    out.extend_from_slice(&[0xEE; 3]);
+    // marker + 4 len bytes + type byte + data
+    check_estimates(&out, &[1, 6, 6, 6, 6, 9, 9, 9, 9, -9]);
+
+    // Canonical ext32 as produced by the encoder.
+    let len = 70_000u32;
+    let mut out = Vec::new();
+    assert_eq!(Marker::Ext32, write_ext_meta(&mut out, len, 7).unwrap());
+    out.extend(std::iter::repeat_n(0xEE, len as usize));
+    assert_eq!(out.len(), MessageLen::len_of(&out).unwrap());
+    assert_eq!(out.len(), MessageLen::len_of(&out[..out.len() - 1]).unwrap_err().len());
+}
+
+#[test]
+fn ext_in_array() {
+    let mut out = Vec::new();
+    write_array_len(&mut out, 3).unwrap();
+    write_ext_meta(&mut out, 1, 7).unwrap(); // fixext1
+    out.push(0xEE);
+    write_ext_meta(&mut out, 3, 7).unwrap(); // ext8
+    out.extend_from_slice(&[0xEE; 3]);
+    write_nil(&mut out).unwrap();
+    assert_eq!(11, out.len());
+
+    check_estimates(&out, &[1, 4, 4, 4, 6, 7, 10, 10, 10, 10, 11, -11]);
+}
+
+#[test]
+fn ext_in_map() {
+    let mut out = Vec::new();
+    write_map_len(&mut out, 1).unwrap();
+    write_ext_meta(&mut out, 8, 1).unwrap(); // fixext8 key
+    out.extend_from_slice(&[0xEE; 8]);
+    write_ext_meta(&mut out, 300, 2).unwrap(); // ext16 value
+    out.extend_from_slice(&[0xEE; 300]);
+    // 1 + (2 + 8) + (4 + 300)
+    assert_eq!(315, out.len());
+
+    let mut expected = vec![1, 3];
+    expected.extend(std::iter::repeat_n(11, 9)); // key: type byte + 8 data bytes
+    expected.extend([12, 15, 15]); // value: marker, then 2 length bytes + type byte
+    expected.resize(315, 315);
+    expected.push(-315);
+    check_estimates(&out, &expected);
+}
+
+#[test]
+fn limit_exceeded_is_parse_error() {
+    // Nesting limit, all data available in the first call.
+    let mut out = Vec::new();
+    for _ in 0..5 {
+        write_array_len(&mut out, 1).unwrap();
+    }
+    write_nil(&mut out).unwrap();
+    let mut est = MessageLen::with_limits(4, 1 << 16);
+    assert!(matches!(est.incremental_len(&out), Err(LenError::ParseError)));
+    // The error is sticky.
+    assert!(matches!(est.incremental_len(&out), Err(LenError::ParseError)));
+    assert_eq!(6, MessageLen::with_limits(5, 1 << 16).incremental_len(&out).unwrap());
+
+    // Nesting limit hit while resuming from a truncated message.
+    let mut est = MessageLen::with_limits(2, 1 << 16);
+    assert_eq!(2, est.incremental_len(&[0x91]).unwrap_err().len());
+    assert!(matches!(est.incremental_len(&[0x91, 0x91, 0xc0]), Err(LenError::ParseError)));
+
+    // `len_of` has a fixed nesting limit of 1024.
+    let mut out = Vec::new();
+    for _ in 0..1025 {
+        write_array_len(&mut out, 1).unwrap();
+    }
+    write_nil(&mut out).unwrap();
+    assert!(matches!(MessageLen::len_of(&out), Err(LenError::ParseError)));
+
+    // Length limit is exclusive, for strings...
+    let mut out = Vec::new();
+    write_str_len(&mut out, 40).unwrap(); // str8
+    out.extend_from_slice(&[b'x'; 40]);
+    assert!(matches!(MessageLen::with_limits(1024, 40).incremental_len(&out), Err(LenError::ParseError)));
+    assert_eq!(42, MessageLen::with_limits(1024, 41).incremental_len(&out).unwrap());
+
+    // ...arrays (counted in items)...
+    let mut out = Vec::new();
+    write_array_len(&mut out, 40).unwrap(); // array16
+    out.extend(std::iter::repeat_n(0xc0, 40));
+    assert!(matches!(MessageLen::with_limits(1024, 40).incremental_len(&out), Err(LenError::ParseError)));
+    assert_eq!(43, MessageLen::with_limits(1024, 41).incremental_len(&out).unwrap());
+
+    // ...maps (counted in keys + values)...
+    let mut out = Vec::new();
+    write_map_len(&mut out, 20).unwrap(); // map16
+    out.extend(std::iter::repeat_n(0xc0, 40));
+    assert!(matches!(MessageLen::with_limits(1024, 40).incremental_len(&out), Err(LenError::ParseError)));
+    assert_eq!(43, MessageLen::with_limits(1024, 41).incremental_len(&out).unwrap());
+
+    // ...and ext payloads.
+    let mut out = Vec::new();
+    write_ext_meta(&mut out, 40, 7).unwrap(); // ext8
+    out.extend_from_slice(&[0xEE; 40]);
+    assert!(matches!(MessageLen::with_limits(1024, 40).incremental_len(&out), Err(LenError::ParseError)));
+    assert_eq!(43, MessageLen::with_limits(1024, 41).incremental_len(&out).unwrap());
+}
+
+#[test]
+fn reserved_marker_is_parse_error() {
+    assert!(matches!(MessageLen::len_of(&[0xc1]), Err(LenError::ParseError)));
+    assert!(matches!(MessageLen::len_of(&[0x92, 0xc0, 0xc1]), Err(LenError::ParseError)));
+    // Truncation before the reserved byte is still just truncation.
+    assert_eq!(3, MessageLen::len_of(&[0x92, 0xc0]).unwrap_err().len());
+
+    // The error is sticky.
+    let mut est = MessageLen::new();
+    assert_eq!(3, est.incremental_len(&[0x92, 0xc0]).unwrap_err().len());
+    assert!(matches!(est.incremental_len(&[0xc1]), Err(LenError::ParseError)));
+    assert!(matches!(est.incremental_len(&[0xc0]), Err(LenError::ParseError)));
+}