Bound SSZ list pre-allocation

ssz/src/decode/try_from_iter.rs

@@ -3,6 +3,19 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::convert::Infallible;
 use std::fmt::Debug;
 
+/// The largest buffer we pre-allocate from an untrusted length, in bytes.
+///
+/// This protects against a hostile input claiming a huge length. The decoded list
+/// can still grow past this limit, so a valid input is never rejected.
+const MAX_BYTES_TO_PRE_ALLOCATE: usize = 1_048_576;
+
+/// Returns the capacity to reserve for a length, bounded by `MAX_BYTES_TO_PRE_ALLOCATE`.
+fn capped_capacity<T>(len_hint: Option<usize>) -> usize {
+    len_hint
+        .unwrap_or(0)
+        .min(MAX_BYTES_TO_PRE_ALLOCATE / std::mem::size_of::<T>().max(1))
+}
+
 /// Partial variant of `std::iter::FromIterator`.
 ///
 /// This trait is implemented for types which can be constructed from an iterator of decoded SSZ
@@ -33,12 +46,11 @@ impl<T> TryFromIter<T> for Vec<T> {
     where
         I: IntoIterator<Item = T>,
     {
-        // Pre-allocate the expected size of the Vec, which is parsed from the SSZ input bytes as
-        // `num_items`. This length has already been checked to be less than or equal to the type's
-        // maximum length in `decode_list_of_variable_length_items`.
+        // The expected size comes from the SSZ input and is not yet checked, so bound the
+        // pre-allocation to a fixed amount. The Vec still grows to fit every element.
         let iter = values.into_iter();
         let (_, opt_max_len) = iter.size_hint();
-        let mut vec = Vec::with_capacity(opt_max_len.unwrap_or(0));
+        let mut vec = Vec::with_capacity(capped_capacity::<T>(opt_max_len));
         vec.extend(iter);
         Ok(vec)
     }
@@ -51,7 +63,12 @@ impl<T, const N: usize> TryFromIter<T> for SmallVec<[T; N]> {
     where
         I: IntoIterator<Item = T>,
     {
-        Ok(Self::from_iter(iter))
+        // Bound the pre-allocation like we do in the `Vec` impl.
+        let iter = iter.into_iter();
+        let (_, opt_max_len) = iter.size_hint();
+        let mut out = SmallVec::with_capacity(capped_capacity::<T>(opt_max_len));
+        out.extend(iter);
+        Ok(out)
     }
 }
 
@@ -101,3 +118,60 @@ where
         C::try_from_iter(self)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// An iterator that overstates its length through `size_hint`.
+    struct Wonky<I> {
+        hint: usize,
+        iter: I,
+    }
+
+    impl<I: Iterator> Iterator for Wonky<I> {
+        type Item = I::Item;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.iter.next()
+        }
+
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (0, Some(self.hint))
+        }
+    }
+
+    #[test]
+    fn vec_try_from_iter_caps_preallocation() {
+        // A raw `Vec::with_capacity(hint)` for this hint aborts with a capacity overflow.
+        let wonky = Wonky {
+            hint: usize::MAX / 8,
+            iter: std::iter::repeat(1u64).take(5),
+        };
+        let vec: Vec<u64> = wonky.try_collect().unwrap();
+        assert_eq!(vec, vec![1u64; 5]);
+        assert!(vec.capacity() <= MAX_BYTES_TO_PRE_ALLOCATE / std::mem::size_of::<u64>());
+    }
+
+    #[test]
+    fn smallvec_try_from_iter_caps_preallocation() {
+        let wonky = Wonky {
+            hint: usize::MAX / 8,
+            iter: std::iter::repeat(1u64).take(5),
+        };
+        let sv: SmallVec<[u64; 4]> = wonky.try_collect().unwrap();
+        assert_eq!(sv.as_slice(), &[1u64; 5]);
+        assert!(sv.capacity() <= MAX_BYTES_TO_PRE_ALLOCATE / std::mem::size_of::<u64>());
+    }
+
+    #[test]
+    fn small_lists_still_reserve_exactly() {
+        let exact: Vec<u64> = (0..3u64)
+            .collect::<Vec<_>>()
+            .into_iter()
+            .try_collect()
+            .unwrap();
+        assert_eq!(exact, vec![0, 1, 2]);
+        assert_eq!(exact.capacity(), 3);
+    }
+}