Added test cases, added a helper function specialized for containers

bincode-org · Dec 10, 2021 · bc30d0a · bc30d0a
1 parent 08c1784
commit bc30d0a
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 12 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -266,7 +266,7 @@ mod internal {
         const ENDIAN: Endian;
     }
 
-    impl<E: InternalEndianConfig, I, A> InternalEndianConfig for Configuration<E, I, A> {
+    impl<E: InternalEndianConfig, I, A, L> InternalEndianConfig for Configuration<E, I, A, L> {
         const ENDIAN: Endian = E::ENDIAN;
     }
 
@@ -280,7 +280,9 @@ mod internal {
         const INT_ENCODING: IntEncoding;
     }
 
-    impl<E, I: InternalIntEncodingConfig, A> InternalIntEncodingConfig for Configuration<E, I, A> {
+    impl<E, I: InternalIntEncodingConfig, A, L> InternalIntEncodingConfig
+        for Configuration<E, I, A, L>
+    {
         const INT_ENCODING: IntEncoding = I::INT_ENCODING;
     }
 
@@ -294,7 +296,9 @@ mod internal {
         const SKIP_FIXED_ARRAY_LENGTH: bool;
     }
 
-    impl<E, I, A: InternalArrayLengthConfig> InternalArrayLengthConfig for Configuration<E, I, A> {
+    impl<E, I, A: InternalArrayLengthConfig, L> InternalArrayLengthConfig
+        for Configuration<E, I, A, L>
+    {
         const SKIP_FIXED_ARRAY_LENGTH: bool = A::SKIP_FIXED_ARRAY_LENGTH;
     }
 

diff --git a/src/de/mod.rs b/src/de/mod.rs
@@ -6,7 +6,11 @@ mod impl_tuples;
 mod impls;
 
 use self::read::{BorrowReader, Reader};
-use crate::{config::Config, error::DecodeError, utils::Sealed};
+use crate::{
+    config::{Config, InternalLimitConfig},
+    error::DecodeError,
+    utils::Sealed,
+};
 
 pub mod read;
 
@@ -58,6 +62,19 @@ pub trait Decoder: Sealed {
     /// This can be used to validate `Configuration::Limit<N>()`.
     fn claim_bytes_read(&mut self, n: usize) -> Result<(), DecodeError>;
 
+    /// Claim that we're going to read a container which contains `len` entries of `T`.
+    /// This will correctly handle overflowing if `len * size_of::<T>() > usize::max_value`
+    fn claim_container_read<T>(&mut self, len: usize) -> Result<(), DecodeError> {
+        if <Self::C as InternalLimitConfig>::LIMIT.is_some() {
+            match len.checked_mul(core::mem::size_of::<T>()) {
+                Some(val) => self.claim_bytes_read(val),
+                None => Err(DecodeError::LimitExceeded),
+            }
+        } else {
+            Ok(())
+        }
+    }
+
     /// Notify the decoder that `n` bytes are being reclaimed.
     ///
     /// When decoding container types, a typical implementation would claim to read `len * size_of::<T>()` bytes.

diff --git a/src/features/impl_alloc.rs b/src/features/impl_alloc.rs
@@ -52,7 +52,7 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len * core::mem::size_of::<T>())?;
+        decoder.claim_container_read::<T>(len)?;
 
         let mut map = BinaryHeap::with_capacity(len);
         for _ in 0..len {
@@ -86,12 +86,12 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len * (core::mem::size_of::<K>() + core::mem::size_of::<V>()))?;
+        decoder.claim_container_read::<(K, V)>(len)?;
 
         let mut map = BTreeMap::new();
         for _ in 0..len {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
-            decoder.unclaim_bytes_read(core::mem::size_of::<K>() + core::mem::size_of::<V>());
+            decoder.unclaim_bytes_read(core::mem::size_of::<(K, V)>());
 
             let key = K::decode(&mut decoder)?;
             let value = V::decode(&mut decoder)?;
@@ -122,7 +122,7 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len)?;
+        decoder.claim_container_read::<T>(len)?;
 
         let mut map = BTreeSet::new();
         for _ in 0..len {
@@ -155,7 +155,7 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len)?;
+        decoder.claim_container_read::<T>(len)?;
 
         let mut map = VecDeque::with_capacity(len);
         for _ in 0..len {
@@ -188,7 +188,7 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len)?;
+        decoder.claim_container_read::<T>(len)?;
 
         let mut vec = Vec::with_capacity(len);
         for _ in 0..len {

diff --git a/src/features/impl_std.rs b/src/features/impl_std.rs
@@ -369,12 +369,12 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        decoder.claim_bytes_read(len * (core::mem::size_of::<K>() + core::mem::size_of::<V>()))?;
+        decoder.claim_container_read::<(K, V)>(len)?;
 
         let mut map = HashMap::with_capacity(len);
         for _ in 0..len {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
-            decoder.unclaim_bytes_read(core::mem::size_of::<K>() + core::mem::size_of::<V>());
+            decoder.unclaim_bytes_read(core::mem::size_of::<(K, V)>());
 
             let k = K::decode(&mut decoder)?;
             let v = V::decode(&mut decoder)?;

diff --git a/tests/alloc.rs b/tests/alloc.rs
@@ -89,3 +89,41 @@ fn test_alloc_commons() {
         set
     });
 }
+
+#[test]
+fn test_container_limits() {
+    use bincode::{error::DecodeError, Decode};
+
+    const DECODE_LIMIT: usize = 100_000;
+
+    // for this test we'll create a malformed package of a lot of bytes
+    let test_cases = &[
+        // u64::max_value(), should overflow
+        bincode::encode_to_vec(u64::max_value(), Configuration::standard()).unwrap(),
+        // A high value which doesn't overflow, but exceeds the decode limit
+        bincode::encode_to_vec(DECODE_LIMIT as u64, Configuration::standard()).unwrap(),
+    ];
+
+    fn validate_fail<T: Decode + core::fmt::Debug>(slice: &[u8]) {
+        let result = bincode::decode_from_slice::<T, _>(
+            slice,
+            Configuration::standard().with_limit::<DECODE_LIMIT>(),
+        );
+
+        assert_eq!(result.unwrap_err(), DecodeError::LimitExceeded);
+    }
+
+    for slice in test_cases {
+        validate_fail::<BinaryHeap<i32>>(slice);
+        validate_fail::<BTreeMap<i32, i32>>(slice);
+        validate_fail::<BTreeSet<i32>>(slice);
+        validate_fail::<VecDeque<i32>>(slice);
+        validate_fail::<Vec<i32>>(slice);
+        validate_fail::<String>(slice);
+        validate_fail::<Box<[u8]>>(slice);
+        #[cfg(feature = "std")]
+        {
+            validate_fail::<std::collections::HashMap<i32, i32>>(slice);
+        }
+    }
+}