fix(parquet): return error instead of panic on corrupt PLAIN BYTE_ARRAY pages

parquet/src/arrow/array_reader/byte_array.rs

@@ -366,12 +366,11 @@ impl ByteArrayDecoderPlain {
         let initial_values_length = output.values.len();
 
         let to_read = len.min(self.max_remaining_values);
-        output.offsets.reserve(to_read);
-
         let remaining_bytes = self.buf.len() - self.offset;
-        if remaining_bytes == 0 {
+        if to_read == 0 || remaining_bytes == 0 {
             return Ok(0);
         }
+        output.offsets.reserve(to_read);
 
         let estimated_bytes = remaining_bytes
             .checked_mul(to_read)
@@ -420,8 +419,11 @@ impl ByteArrayDecoderPlain {
             }
             let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap();
             let len = u32::from_le_bytes(len_bytes) as usize;
+            self.offset = (self.offset + 4)
+                .checked_add(len)
+                .filter(|&end| end <= buf.len())
+                .ok_or_else(|| ParquetError::EOF("eof decoding byte array".into()))?;
             skip += 1;
-            self.offset = self.offset + 4 + len;
         }
         self.max_remaining_values -= skip;
         Ok(skip)
@@ -767,6 +769,57 @@ mod tests {
         assert_eq!(decoder.max_remaining_values, 1);
     }
 
+    #[test]
+    fn test_plain_decoder_skip_past_end_returns_error() {
+        let buffer = Bytes::from_static(&[0xFF, 0xFF, 0, 0, b'a']);
+        let mut decoder = ByteArrayDecoderPlain::new(buffer, 2, Some(2), false);
+
+        let error = decoder.skip(1).unwrap_err();
+        assert!(matches!(error, ParquetError::EOF(_)), "{error}");
+        assert_eq!(decoder.offset, 0);
+        assert_eq!(decoder.max_remaining_values, 2);
+    }
+
+    #[test]
+    fn test_plain_decoder_no_remaining_values() {
+        let buffer = Bytes::from_static(&[1, 0, 0, 0, b'a']);
+        let mut decoder = ByteArrayDecoderPlain::new(buffer, 0, Some(0), false);
+        let mut output = OffsetBuffer::<i32>::with_capacity(0);
+
+        assert_eq!(decoder.read(&mut output, 1).unwrap(), 0);
+        assert!(output.values.is_empty());
+        assert_eq!(output.offsets.as_slice(), &[0]);
+    }
+
+    #[test]
+    fn test_plain_decoder_truncated_value() {
+        let bad_values: [&[u8]; 5] = [
+            &[10, 0, 0, 0, b'x'],
+            &[0xFF, 0xFF, 0xFF, 0xFF, b'x'],
+            &[3],
+            &[3, 0],
+            &[3, 0, 0],
+        ];
+        for bad in bad_values {
+            let mut page = vec![2, 0, 0, 0, b'a', b'b'];
+            page.extend_from_slice(bad);
+            let page = Bytes::from(page);
+
+            let mut decoder = ByteArrayDecoderPlain::new(page.clone(), 3, Some(3), false);
+            let mut output = OffsetBuffer::<i32>::with_capacity(0);
+            let error = decoder.read(&mut output, 3).unwrap_err();
+            assert!(matches!(error, ParquetError::EOF(_)), "{bad:?}: {error}");
+            assert_eq!(output.values.as_slice(), b"ab");
+            assert_eq!(output.offsets.as_slice(), &[0, 2]);
+            assert_eq!(decoder.offset, 6);
+
+            let mut decoder = ByteArrayDecoderPlain::new(page, 3, Some(3), false);
+            let error = decoder.skip(3).unwrap_err();
+            assert!(matches!(error, ParquetError::EOF(_)), "{bad:?}: {error}");
+            assert_eq!(decoder.offset, 6);
+        }
+    }
+
     #[test]
     fn test_delta_length_checked_offset() {
         let max = i32::MAX as usize;

parquet/src/arrow/array_reader/byte_view_array.rs

@@ -440,8 +440,11 @@ impl ByteViewArrayDecoderPlain {
             }
             let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap();
             let len = u32::from_le_bytes(len_bytes) as usize;
+            self.offset = (self.offset + 4)
+                .checked_add(len)
+                .filter(|&end| end <= buf.len())
+                .ok_or_else(|| ParquetError::EOF("eof decoding byte array".into()))?;
             skip += 1;
-            self.offset = self.offset + 4 + len;
         }
         self.max_remaining_values -= skip;
         Ok(skip)
@@ -831,4 +834,15 @@ mod tests {
         decoder.read(&mut view_buffer, 1).unwrap();
         assert_eq!(view_buffer.buffers.len(), 1);
     }
+
+    #[test]
+    fn test_byte_view_array_plain_decoder_skip_past_end_returns_error() {
+        let buffer = Bytes::from_static(&[0xFF, 0xFF, 0, 0, b'a']);
+        let mut decoder = ByteViewArrayDecoderPlain::new(buffer, 2, Some(2), false);
+
+        let error = decoder.skip(1).unwrap_err();
+        assert!(matches!(error, ParquetError::EOF(_)), "{error}");
+        assert_eq!(decoder.offset, 0);
+        assert_eq!(decoder.max_remaining_values, 2);
+    }
 }

parquet/src/arrow/arrow_reader/tests/column_reader.rs

@@ -46,6 +46,88 @@ fn test_truncated_definition_level_bitpacked_run_returns_error() {
     );
 }
 
+fn write_plain_utf8(array: StringArray, writer_version: WriterVersion) -> Vec<u8> {
+    let nullable = array.null_count() > 0;
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "c",
+        ArrowDataType::Utf8,
+        nullable,
+    )]));
+    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+    let props = WriterProperties::builder()
+        .set_dictionary_enabled(false)
+        .set_encoding(Encoding::PLAIN)
+        .set_writer_version(writer_version)
+        .build();
+
+    let mut buf = Vec::new();
+    let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
+    writer.write(&batch).unwrap();
+    writer.close().unwrap();
+    buf
+}
+
+#[test]
+fn test_plain_byte_array_skip_corrupt_length_returns_error() {
+    let mut buf = write_plain_utf8(
+        StringArray::from(vec!["first", "second", "third"]),
+        WriterVersion::PARQUET_1_0,
+    );
+
+    // Make the first value's length prefix point past the end of the page.
+    let pattern = [5, 0, 0, 0, b'f', b'i', b'r', b's', b't'];
+    let pos = buf
+        .windows(pattern.len())
+        .position(|w| w == pattern)
+        .unwrap();
+    buf[pos + 1] = 0xFF;
+
+    let selection = RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(2)]);
+    let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buf))
+        .unwrap()
+        .with_row_selection(selection)
+        .with_row_selection_policy(crate::arrow::arrow_reader::RowSelectionPolicy::Selectors)
+        .build()
+        .unwrap();
+    let error = reader.next().unwrap().unwrap_err();
+
+    assert!(
+        error.to_string().contains("eof decoding byte array"),
+        "{error}"
+    );
+}
+
+#[test]
+fn test_plain_byte_array_all_null_page_with_values_returns_error() {
+    let mut buf = write_plain_utf8(
+        StringArray::from(vec![Some("a"), None]),
+        WriterVersion::PARQUET_2_0,
+    );
+
+    // Patch the v2 page header from num_nulls = 1 to num_nulls = 2, so the
+    // header says the page has no non-null values while the levels and value
+    // buffer still contain one.
+    let pattern = [0x15, 0x04, 0x15, 0x02, 0x15, 0x04];
+    let pos = buf
+        .windows(pattern.len())
+        .position(|w| w == pattern)
+        .unwrap();
+    buf[pos + 3] = 0x04;
+
+    let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buf))
+        .unwrap()
+        .build()
+        .unwrap();
+    let error = reader.next().unwrap().unwrap_err();
+
+    assert!(
+        error
+            .to_string()
+            .contains("insufficient values read from column - expected: 1, got: 0"),
+        "{error}"
+    );
+}
+
 #[test]
 fn test_null_column_reader_test() {
     let mut file = tempfile::tempfile().unwrap();