diff --git a/crates/polars-parquet/src/parquet/handwritten_thrift/file_metadata_thrift.rs b/crates/polars-parquet/src/parquet/handwritten_thrift/file_metadata_thrift.rs index 01cd5ce31d51..41bd7a0fd33e 100644 --- a/crates/polars-parquet/src/parquet/handwritten_thrift/file_metadata_thrift.rs +++ b/crates/polars-parquet/src/parquet/handwritten_thrift/file_metadata_thrift.rs @@ -41,8 +41,19 @@ fn read_list( prot: &mut ThriftSliceInputProtocol<'_>, mut read_one: impl FnMut(&mut ThriftSliceInputProtocol<'_>) -> ParquetResult, ) -> ParquetResult> { + // `read_list_begin` rejects sizes above `i32::MAX`, but `i32::MAX` of + // any non-trivial element would still be many GB. Bound the up-front + // capacity by the bytes still available: every list element occupies + // at least one byte on the wire. let list = prot.read_list_begin()?; - let mut v = Vec::with_capacity(list.size.max(0) as usize); + let declared = list.size.max(0) as usize; + let bound = declared.min(prot.bytes_remaining()); + let mut v: Vec = Vec::new(); + if v.try_reserve_exact(bound).is_err() { + return Err(crate::parquet::error::ParquetError::oos(format!( + "cannot allocate Thrift list of {declared} elements", + ))); + } for _ in 0..list.size { v.push(read_one(prot)?); } diff --git a/crates/polars-parquet/src/parquet/handwritten_thrift/parquet_thrift.rs b/crates/polars-parquet/src/parquet/handwritten_thrift/parquet_thrift.rs index b84d6ea4ea3a..2f3e16cccfc3 100644 --- a/crates/polars-parquet/src/parquet/handwritten_thrift/parquet_thrift.rs +++ b/crates/polars-parquet/src/parquet/handwritten_thrift/parquet_thrift.rs @@ -65,6 +65,7 @@ pub(crate) enum ThriftProtocolError { InvalidElementType(u8), FieldDeltaOverflow { field_delta: u8, last_field_id: i16 }, InvalidBoolean(u8), + IntegerOverflow, Utf8Error, SkipDepth(FieldType), SkipUnsupportedType(FieldType), @@ -89,6 +90,9 @@ impl From for ParquetError { ThriftProtocolError::InvalidBoolean(value) => { general_err!("cannot convert {} into bool", value) }, + ThriftProtocolError::IntegerOverflow => { + general_err!("integer overflow decoding thrift value") + }, ThriftProtocolError::Utf8Error => general_err!("invalid utf8"), ThriftProtocolError::SkipDepth(field_type) => { general_err!("cannot parse past {:?}", field_type) @@ -113,6 +117,13 @@ impl From for ThriftProtocolError { } } +impl From for ThriftProtocolError { + fn from(_: std::num::TryFromIntError) -> Self { + // ignore error payload to reduce the size of ThriftProtocolError + Self::IntegerOverflow + } +} + // Arrow writes `Result` here, relying on its two-arg // `std::result::Result` alias. Our local `Result = ParquetResult` takes // one arg, so spell `std::result::Result` explicitly for this one alias. @@ -296,6 +307,23 @@ pub(crate) trait ThriftCompactInputProtocol<'a> { /// Skip the next `n` bytes of input. fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()>; + /// Upper bound on the number of bytes still available to read from this + /// input. Implementations backed by an in-memory slice can return the + /// exact remaining length; implementations wrapping a streaming reader + /// that doesn't know its bound should fall back to `usize::MAX`, which + /// is what the default impl returns. + /// + /// Used by [`read_thrift_vec`] (and the local `read_list` helper in + /// `file_metadata_thrift`) to cap up-front list allocations: a Thrift + /// list element occupies at least one byte on the wire, so the list + /// size declared in the header is always at most `bytes_remaining()`. + /// Without this cap an attacker-controlled list-size varint of + /// ~`i32::MAX` drives a multi-GiB up-front allocation through + /// `Vec::with_capacity`. + fn bytes_remaining(&self) -> usize { + usize::MAX + } + /// Read a ULEB128 encoded unsigned varint from the input. /// /// Fast path for the 1-byte case (~80-90% of Parquet metadata varints). @@ -343,7 +371,13 @@ pub(crate) trait ThriftCompactInputProtocol<'a> { // high bits set high if count and type encoded separately possible_element_count as i32 } else { - self.read_vlq()? as _ + // The list size on the wire is an unsigned varint, but the + // Parquet Thrift schema (and Java's `int`) represent it as + // `i32`. A varint that decodes above `i32::MAX` is malformed + // input — reject it here at the protocol layer rather than + // letting the cast wrap into a negative size that downstream + // allocation code has to re-validate. + i32::try_from(self.read_vlq()?)? }; Ok(ListIdentifier { @@ -591,6 +625,11 @@ impl<'a> ThriftSliceInputProtocol<'a> { } impl<'b, 'a: 'b> ThriftCompactInputProtocol<'b> for ThriftSliceInputProtocol<'a> { + #[inline] + fn bytes_remaining(&self) -> usize { + self.buf.len() + } + #[inline] fn read_byte(&mut self) -> ThriftProtocolResult { let ret = *self.buf.first().ok_or(ThriftProtocolError::Eof)?; @@ -745,8 +784,25 @@ where R: ThriftCompactInputProtocol<'a>, T: ReadThrift<'a, R>, { + // `read_list_begin` already rejects sizes above `i32::MAX`, but on a + // 64-bit target that is still ~2 GiB of `T`s. Cap the up-front + // capacity by the bytes still in the input: each Thrift list element + // occupies at least one byte on the wire, so a declared size larger + // than what the reader has left is invariably invalid input. The + // subsequent `try_reserve_exact` then serves as belt-and-suspenders. let list_ident = prot.read_list_begin()?; - let mut res = Vec::with_capacity(list_ident.size as usize); + let declared = list_ident.size as usize; + let bound = declared.min(prot.bytes_remaining()); + let mut res: Vec = Vec::new(); + if res.try_reserve_exact(bound).is_err() { + return Err(general_err!( + "cannot allocate Thrift list of {} elements", + declared + )); + } + // Iterate up to the declared length: `T::read_thrift` will surface a + // protocol-level EOF the moment the input is exhausted, which is the + // correct user-visible error for a header that overstates its list size. for _ in 0..list_ident.size { let val = T::read_thrift(prot)?; res.push(val);