diff --git a/ssz/src/decode/try_from_iter.rs b/ssz/src/decode/try_from_iter.rs index 1ff89a1..9a70426 100644 --- a/ssz/src/decode/try_from_iter.rs +++ b/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(len_hint: Option) -> usize { + len_hint + .unwrap_or(0) + .min(MAX_BYTES_TO_PRE_ALLOCATE / std::mem::size_of::().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 TryFromIter for Vec { where I: IntoIterator, { - // 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::(opt_max_len)); vec.extend(iter); Ok(vec) } @@ -51,7 +63,12 @@ impl TryFromIter for SmallVec<[T; N]> { where I: IntoIterator, { - 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::(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 { + hint: usize, + iter: I, + } + + impl Iterator for Wonky { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.iter.next() + } + + fn size_hint(&self) -> (usize, Option) { + (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 = wonky.try_collect().unwrap(); + assert_eq!(vec, vec![1u64; 5]); + assert!(vec.capacity() <= MAX_BYTES_TO_PRE_ALLOCATE / std::mem::size_of::()); + } + + #[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::()); + } + + #[test] + fn small_lists_still_reserve_exactly() { + let exact: Vec = (0..3u64) + .collect::>() + .into_iter() + .try_collect() + .unwrap(); + assert_eq!(exact, vec![0, 1, 2]); + assert_eq!(exact.capacity(), 3); + } +}