fix standalone footer validation

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-file/src/footer/deserializer.rs

@@ -40,8 +40,8 @@ pub struct FooterDeserializer {
 
     // Internal state that we accumulate
 
-    // The file size, possibly provided externally.
-    file_size: Option<u64>,
+    // The size of the byte stream containing the serialized footer.
+    stream_size: Option<u64>,
     // The postscript, once we've parsed it.
     postscript: Option<Postscript>,
 }
@@ -52,7 +52,7 @@ impl FooterDeserializer {
             buffer: initial_read,
             session,
             dtype: None,
-            file_size: None,
+            stream_size: None,
             postscript: None,
         }
     }
@@ -71,15 +71,19 @@ impl FooterDeserializer {
         self
     }
 
-    /// Provide the total file size.
-    pub fn with_size(mut self, file_size: u64) -> Self {
-        self.file_size = Some(file_size);
+    /// Provide the total size of the byte stream containing this serialized footer.
+    ///
+    /// For a footer read from the end of a Vortex file, this is the file size. For a standalone
+    /// blob created by [`crate::footer::FooterSerializer`] with its default offset, this is the
+    /// size of that blob, not the size of the data file described by the footer.
+    pub fn with_size(mut self, stream_size: u64) -> Self {
+        self.stream_size = Some(stream_size);
         self
     }
 
-    /// Provide or clear the total file size.
-    pub fn with_some_size(mut self, file_size: Option<u64>) -> Self {
-        self.file_size = file_size;
+    /// Provide or clear the total size of the byte stream containing this serialized footer.
+    pub fn with_some_size(mut self, stream_size: Option<u64>) -> Self {
+        self.stream_size = stream_size;
         self
     }
 
@@ -120,11 +124,18 @@ impl FooterDeserializer {
         // The other postscript segments are required, so now we figure out our the offset that
         // contains all the required segments.
 
-        // The initial offset is the file size - the size of our initial read.
-        let Some(file_size) = self.file_size else {
+        // The initial offset is the byte stream size minus the size of our initial read.
+        let Some(stream_size) = self.stream_size else {
             return Ok(DeserializeStep::NeedFileSize);
         };
-        let initial_offset = file_size - (self.buffer.len() as u64);
+        let initial_offset = stream_size
+            .checked_sub(self.buffer.len() as u64)
+            .ok_or_else(|| {
+                vortex_err!(
+                    "Footer buffer length {} exceeds declared byte stream size {stream_size}",
+                    self.buffer.len()
+                )
+            })?;
 
         let mut read_more_offset = initial_offset;
         if let Some(dtype_segment) = &dtype_segment {
@@ -269,20 +280,7 @@ impl FooterDeserializer {
             layout_segment,
         )?);
 
-        // The file size is always known by the time we parse the footer, since `deserialize`
-        // returns `NeedFileSize` before reaching this point when it is missing.
-        let file_size = self
-            .file_size
-            .vortex_expect("file size is required before parsing the footer");
-
-        Footer::from_flatbuffer(
-            footer_bytes,
-            layout_bytes,
-            dtype,
-            file_stats,
-            file_size,
-            &self.session,
-        )
+        Footer::from_flatbuffer(footer_bytes, layout_bytes, dtype, file_stats, &self.session)
     }
 }
 
@@ -386,6 +384,25 @@ mod tests {
         assert!(err.to_string().contains("out of bounds"), "{err}");
         Ok(())
     }
+
+    #[test]
+    fn deserialize_rejects_buffer_larger_than_declared_size() -> VortexResult<()> {
+        let postscript = Postscript {
+            dtype: None,
+            layout: segment(0, 1),
+            statistics: None,
+            footer: segment(1, 1),
+        };
+        let buffer = eof_buffer(&postscript)?;
+        let declared_size = buffer.len() as u64 - 1;
+
+        let mut deserializer = FooterDeserializer::new(buffer, array_session())
+            .with_dtype(DType::Primitive(PType::I32, Nullability::NonNullable))
+            .with_size(declared_size);
+        let err = deserializer.deserialize().unwrap_err();
+        assert!(err.to_string().contains("exceeds declared"), "{err}");
+        Ok(())
+    }
 }
 
 #[derive(Debug)]
@@ -398,7 +415,7 @@ pub enum DeserializeStep {
         /// Number of bytes to read and prefix into the deserializer.
         len: usize,
     },
-    /// The total file size is required before offsets can be resolved.
+    /// The total size of the byte stream is required before offsets can be resolved.
     NeedFileSize,
     /// Footer deserialization is complete.
     Done(Footer),

vortex-file/src/footer/mod.rs

@@ -73,15 +73,11 @@ impl Footer {
     }
 
     /// Read the [`Footer`] from a flatbuffer.
-    ///
-    /// `file_size` is the total size of the file in bytes and is used to validate that every
-    /// segment declared in the footer lies within the file.
     pub(crate) fn from_flatbuffer(
         footer_bytes: &[u8],
         layout_bytes: FlatBuffer,
         dtype: DType,
         statistics: Option<FileStatistics>,
-        file_size: u64,
         session: &VortexSession,
     ) -> VortexResult<Self> {
         let approx_byte_size = footer_bytes.len() + layout_bytes.len();
@@ -128,11 +124,6 @@ impl Footer {
             vortex_bail!("Segment offsets are not ordered");
         }
 
-        // A corrupt or malicious file can declare a segment whose offset or length extends past
-        // the end of the file. Reject such files here so that later slicing of the backing buffer
-        // returns a `VortexError` rather than panicking (see issue #8819).
-        validate_segments_within_file(&segments, file_size)?;
-
         Ok(Self {
             root_layout,
             segments,
@@ -181,6 +172,11 @@ impl Footer {
         self.root_layout.row_count()
     }
 
+    /// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
+    pub(crate) fn validate_file_size(&self, file_size: u64) -> VortexResult<()> {
+        validate_segments_within_file(&self.segments, file_size)
+    }
+
     /// Returns a serializer for this footer.
     pub fn into_serializer(self) -> FooterSerializer {
         FooterSerializer::new(self)

vortex-file/src/open.rs

@@ -218,6 +218,7 @@ impl VortexOpenOptions {
             Some(footer) => footer,
             None => block_on(opts.read_footer(&buffer))?,
         };
+        footer.validate_file_size(buffer.len() as u64)?;
 
         let segment_source = Arc::new(BufferSegmentSource::new(
             buffer,
@@ -246,6 +247,9 @@ impl VortexOpenOptions {
             .unwrap_or_else(|| Arc::new(DefaultMetricsRegistry::default()));
 
         let footer = if let Some(footer) = self.footer {
+            if let Some(file_size) = self.file_size {
+                footer.validate_file_size(file_size)?;
+            }
             footer
         } else {
             self.read_footer(&reader).await?
@@ -324,6 +328,11 @@ impl VortexOpenOptions {
             }
         }?;
 
+        // Segment specs describe the data file, not necessarily the byte stream used to
+        // deserialize a standalone cached footer. Validate them here, where we know this is the
+        // size of the actual data source (see issue #8819).
+        footer.validate_file_size(file_size)?;
+
         // If the initial read happened to cover any segments, then we can populate the
         // segment cache
         let initial_offset = file_size - (deserializer.buffer().len() as u64);
@@ -351,9 +360,9 @@ impl VortexOpenOptions {
                 SegmentId::from(u32::try_from(idx).vortex_expect("Invalid segment ID"));
             let offset =
                 usize::try_from(segment.offset - initial_offset).vortex_expect("Invalid offset");
-            // The segment map is validated against the file size when the footer is parsed, but we
-            // still bounds-check here so slicing never depends on that distant validation for
-            // panic safety (see issue #8819).
+            // The segment map is validated against the file size before this method is called, but
+            // still bounds-check here so slicing never depends on that distant validation for panic
+            // safety (see issue #8819).
             let end = offset
                 .checked_add(segment.length as usize)
                 .filter(|end| *end <= initial_read.len())
@@ -620,4 +629,52 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn standalone_footer_round_trip() -> VortexResult<()> {
+        let session = test_session();
+
+        let mut file_bytes = ByteBufferMut::empty();
+        let values = (0i32..65_536)
+            .map(|value| value.wrapping_mul(1_664_525).wrapping_add(1_013_904_223))
+            .collect::<Vec<_>>();
+        let array = Buffer::from(values).into_array();
+        session
+            .write_options()
+            .write(&mut file_bytes, array.to_array_stream())
+            .await?;
+        let file_bytes = ByteBuffer::from(file_bytes);
+
+        let footer = session
+            .open_options()
+            .open_buffer(file_bytes.clone())?
+            .footer()
+            .clone();
+        let last_segment_end = footer
+            .segment_map()
+            .iter()
+            .map(|segment| segment.offset + u64::from(segment.length))
+            .max()
+            .unwrap_or_default();
+        let serialized_footer = footer.into_serializer().serialize()?;
+        let serialized_footer_size = serialized_footer.iter().map(ByteBuffer::len).sum();
+        assert!(last_segment_end > serialized_footer_size as u64);
+        let mut footer_bytes = ByteBufferMut::with_capacity(serialized_footer_size);
+        for buffer in serialized_footer {
+            footer_bytes.extend_from_slice(&buffer);
+        }
+
+        let mut deserializer = Footer::deserializer(footer_bytes.freeze(), session.clone())
+            .with_size(serialized_footer_size as u64);
+        let DeserializeStep::Done(cached_footer) = deserializer.deserialize()? else {
+            vortex_bail!("standalone footer bytes must be sufficient for deserialization");
+        };
+
+        session
+            .open_options()
+            .with_footer(cached_footer)
+            .open_buffer(file_bytes)?;
+
+        Ok(())
+    }
 }