fix(zstd): validate frame content sizes (#9164)

## Rationale for this change

Fuzzing found that zstd (arrays + buffers) serialised output sizes were
not validated against the declared uncompressed sizes in frame headers.

## What changes are included in this PR?

Validate serialized zstd output sizes against frame headers before CPU
allocation and host-backed GPU planning. Add regressions for zstd arrays
and zstd buffers.

Note that as long as the size declared in the frame header matches
what's declared in metadata, allocations are not bounded. But this is
much harder to fix, since zstd can essentially decompress to
`⌈compressed_len/3⌉ × 128 KiB` which is ~42GiB per MiB of compressed
data so not particularly useful as a bound.

## What APIs are changed? Are there any user-facing changes?

Malformed files where output sizes declared in zstd frames don't match
the header are now refused instead of allocating a buffer of
file-controlled size.

Investigated and built with Codex.

Signed-off-by: Lorenz Hübschle <lorenz@firebolt.io>

Author: lorenzhs

encodings/zstd/src/array.rs

@@ -59,6 +59,7 @@ use vortex_session::registry::CachedId;
 
 use crate::ZstdFrameMetadata;
 use crate::ZstdMetadata;
+use crate::validate_frame_content_size;
 
 // Zstd doesn't support training dictionaries on very few samples.
 const MIN_SAMPLES_FOR_DICTIONARY: usize = 8;
@@ -670,6 +671,10 @@ impl ZstdData {
             self.frames.len(),
             self.metadata.frames.len()
         );
+        for (index, (frame, metadata)) in self.frames.iter().zip(&self.metadata.frames).enumerate()
+        {
+            validate_frame_content_size(frame.as_slice(), metadata.uncompressed_size, index)?;
+        }
 
         Ok(())
     }

encodings/zstd/src/lib.rs

@@ -23,6 +23,9 @@
 
 pub use array::*;
 use vortex_array::dtype::proto::dtype as pb;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
 #[cfg(feature = "unstable_encodings")]
 pub use zstd_buffers::*;
 
@@ -36,6 +39,22 @@ mod zstd_buffers;
 #[cfg(test)]
 mod test;
 
+/// Ensure Vortex metadata agrees with the content size declared by a zstd frame.
+pub(crate) fn validate_frame_content_size(
+    frame: &[u8],
+    metadata_size: u64,
+    index: usize,
+) -> VortexResult<()> {
+    let frame_content_size = zstd::zstd_safe::get_frame_content_size(frame)
+        .map_err(|error| vortex_err!("Invalid zstd frame {index}: {error}"))?
+        .ok_or_else(|| vortex_err!("Zstd frame {index} does not declare a content size"))?;
+    vortex_ensure!(
+        metadata_size == frame_content_size,
+        "Zstd frame {index} metadata declares {metadata_size} uncompressed bytes, but its header declares {frame_content_size}"
+    );
+    Ok(())
+}
+
 #[derive(Clone, prost::Message)]
 /// Metadata for one zstd frame.
 pub struct ZstdFrameMetadata {

encodings/zstd/src/test.rs

@@ -13,6 +13,7 @@ use vortex_array::assert_nth_scalar;
 use vortex_array::builders::DynVarBinBuilder;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
 use vortex_array::validity::Validity;
 use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
@@ -289,3 +290,20 @@ fn test_zstd_frame_start_buffer_alignment() {
 
     assert!(compressed.is_ok());
 }
+
+#[test]
+fn test_zstd_rejects_mismatched_frame_content_size() {
+    let mut ctx = array_session().create_execution_ctx();
+    let compressed =
+        Zstd::from_primitive(&PrimitiveArray::from_iter([1_i32, 2, 3]), 0, 0, &mut ctx).unwrap();
+    let mut data = compressed.data().clone();
+    data.metadata.frames[0].uncompressed_size = 16 * 1024 * 1024 * 1024;
+
+    let error = Zstd::try_new(
+        DType::Primitive(PType::I32, Nullability::NonNullable),
+        data,
+        Validity::NonNullable,
+    )
+    .unwrap_err();
+    assert!(error.to_string().contains("metadata declares"));
+}

encodings/zstd/src/zstd_buffers.rs

@@ -39,6 +39,7 @@ use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
 use crate::ZstdBuffersMetadata;
+use crate::validate_frame_content_size;
 
 /// A [`ZstdBuffers`]-encoded Vortex array.
 pub type ZstdBuffersArray = Array<ZstdBuffers>;
@@ -258,6 +259,8 @@ impl ZstdBuffersData {
             let alignment = self.buffer_alignments.get(i).copied().unwrap_or(1);
 
             let aligned = Alignment::try_from(alignment)?;
+            let compressed = buf.clone().try_to_host_sync()?;
+            validate_frame_content_size(compressed.as_slice(), uncompressed_size, i)?;
             let mut output = ByteBufferMut::with_capacity_aligned(size, aligned);
             let spare = output.spare_capacity_mut();
 
@@ -274,7 +277,6 @@ impl ZstdBuffersData {
             // `set_len(size)` after zstd reports how many bytes were written.
             let dst =
                 unsafe { std::slice::from_raw_parts_mut(spare.as_mut_ptr().cast::<u8>(), size) };
-            let compressed = buf.clone().try_to_host_sync()?;
             let written = decompressor.decompress_to_buffer(compressed.as_slice(), dst)?;
             if written != size {
                 return Err(vortex_err!(
@@ -295,6 +297,18 @@ impl ZstdBuffersData {
         // If invariants are somehow broken, device decompression could have UB, so ensure
         // they still hold.
         self.validate()?;
+        // Only host-resident frames are checked: reading a device frame's header would cost a D2H
+        // copy per frame, and the device output allocation is fallible so it cannot abort.
+        for (index, (buffer, &metadata_size)) in self
+            .compressed_buffers
+            .iter()
+            .zip(&self.uncompressed_sizes)
+            .enumerate()
+        {
+            if let Some(frame) = buffer.as_host_opt() {
+                validate_frame_content_size(frame.as_slice(), metadata_size, index)?;
+            }
+        }
 
         let output_sizes = self
             .uncompressed_sizes
@@ -653,6 +667,24 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_rejects_mismatched_frame_content_size_before_output_allocation() -> VortexResult<()> {
+        let compressed = ZstdBuffers::compress(&make_primitive_array(), 3, &array_session())?;
+        let mut data = compressed.data().clone();
+        data.uncompressed_sizes[0] = 16 * 1024 * 1024 * 1024;
+
+        for error in [
+            data.decompress_buffers().unwrap_err(),
+            data.decode_plan().unwrap_err(),
+        ] {
+            assert!(
+                error.to_string().contains("metadata declares"),
+                "unexpected error: {error}"
+            );
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_validity_delegates_for_nullable_input() -> VortexResult<()> {
         let input = make_nullable_primitive_array();