diff --git a/src/tinyvec.rs b/src/tinyvec.rs index 28e4102..8fe1fc0 100644 --- a/src/tinyvec.rs +++ b/src/tinyvec.rs @@ -234,6 +234,25 @@ where } } +/// Caps an untrusted, deserialized element count before it is handed to +/// `with_capacity`, so a hostile length prefix cannot force a huge eager +/// allocation (and its allocation-abort DoS) before a single element has been +/// read. The reservation is limited to `MAX_PREALLOC_BYTES` worth of items; the +/// container still grows to the real length via `push` as elements actually +/// arrive, so well-formed input is unaffected. +#[cfg(any(feature = "borsh", feature = "bin-proto", feature = "serde"))] +fn cautious_capacity(len: usize) -> usize { + // Mirrors serde's `size_hint::cautious`: never trust a wire-provided length + // as an allocation size. + const MAX_PREALLOC_BYTES: usize = 4096; + let item_size = core::mem::size_of::(); + if item_size == 0 { + len + } else { + core::cmp::min(len, MAX_PREALLOC_BYTES / item_size) + } +} + #[cfg(feature = "borsh")] #[cfg_attr(docs_rs, doc(cfg(feature = "borsh")))] impl borsh::BorshDeserialize for TinyVec @@ -244,7 +263,8 @@ where reader: &mut R, ) -> borsh::io::Result { let len = ::deserialize_reader(reader)?; - let mut new_tinyvec = Self::with_capacity(len); + let mut new_tinyvec = + Self::with_capacity(cautious_capacity::(len)); for _ in 0..len { new_tinyvec.push( @@ -313,7 +333,8 @@ where { let item_count = tag.0.try_into().map_err(|_| bin_proto::Error::TagConvert)?; - let mut values = Self::with_capacity(item_count); + let mut values = + Self::with_capacity(cautious_capacity::(item_count)); for _ in 0..item_count { values.push(bin_proto::BitDecode::<_, _>::decode::<_, E>(read, ctx, ())?); } @@ -1951,7 +1972,9 @@ where S: SeqAccess<'de>, { let mut new_tinyvec = match seq.size_hint() { - Some(expected_size) => TinyVec::with_capacity(expected_size), + Some(expected_size) => { + TinyVec::with_capacity(cautious_capacity::(expected_size)) + } None => Default::default(), }; diff --git a/tests/tinyvec.rs b/tests/tinyvec.rs index 12311ef..2c1abd4 100644 --- a/tests/tinyvec.rs +++ b/tests/tinyvec.rs @@ -478,6 +478,20 @@ fn TinyVec_borsh_de_heap() { assert_eq!(tv, des); } +#[cfg(feature = "borsh")] +#[test] +fn TinyVec_borsh_de_hostile_length_no_abort() { + // A tiny buffer that claims an enormous element count but supplies no + // elements. Before the cautious-capacity fix this drove + // `with_capacity(huge)` and aborted the process; now the eager reservation + // is bounded and decoding simply fails on the missing element bytes. + let huge: u64 = 1u64 << 60; + let mut buffer = Vec::new(); + buffer.extend_from_slice(&huge.to_le_bytes()); + let des: Result, _> = borsh::from_slice(&buffer); + assert!(des.is_err()); +} + #[cfg(feature = "bin-proto")] #[test] fn TinyVec_bin_proto_encode_untagged() {