fix: Check for corrupted repeat/define lengths in Parquet headers (facebookincubator#16511)

Summary:
Pull Request resolved: facebookincubator#16511

We currently do not validate the correctness of the repeat/define lengths we read in the Parquet header,
this can lead us to access memory outside the data buffer.

Add checks to validate we do not go off the end of the pageData_ buffer.

Reviewed By: Yuhta

Differential Revision: D94270200

fbshipit-source-id: 1e78b2c09748fddc52bc01cbda58c2c6cb6a0f1e

Author: kevinwilfong

velox/dwio/parquet/reader/PageReader.cpp

@@ -232,18 +232,42 @@ void PageReader::prepareDataPageV1(const PageHeader& pageHeader, int64_t row) {
       pageHeader.compressed_page_size,
       pageHeader.uncompressed_page_size);
   auto pageEnd = pageData_ + pageHeader.uncompressed_page_size;
+  auto remainingBytes = pageHeader.uncompressed_page_size;
   if (maxRepeat_ > 0) {
+    VELOX_CHECK_GE(
+        remainingBytes,
+        sizeof(int32_t),
+        "Insufficient bytes for repetition level length (corrupt data page?)");
     uint32_t repeatLength = readField<int32_t>(pageData_);
+    remainingBytes -= sizeof(int32_t);
+    VELOX_CHECK_LE(
+        repeatLength,
+        remainingBytes,
+        "Repetition level length {} exceeds remaining page size {} (corrupt data page?)",
+        repeatLength,
+        remainingBytes);
     repeatDecoder_ = std::make_unique<RleDecoder>(
         reinterpret_cast<const uint8_t*>(pageData_),
         repeatLength,
         ::arrow::bit_util::NumRequiredBits(maxRepeat_));
 
     pageData_ += repeatLength;
+    remainingBytes -= repeatLength;
   }
 
   if (maxDefine_ > 0) {
+    VELOX_CHECK_GE(
+        remainingBytes,
+        sizeof(uint32_t),
+        "Insufficient bytes for definition level length (corrupt data page?)");
     auto defineLength = readField<uint32_t>(pageData_);
+    remainingBytes -= sizeof(uint32_t);
+    VELOX_CHECK_LE(
+        defineLength,
+        remainingBytes,
+        "Definition level length {} exceeds remaining page size {} (corrupt data page?)",
+        defineLength,
+        remainingBytes);
     if (maxDefine_ == 1) {
       defineDecoder_ = std::make_unique<RleBpDecoder>(
           pageData_,
@@ -288,6 +312,13 @@ void PageReader::prepareDataPageV2(const PageHeader& pageHeader, int64_t row) {
       pageHeader.data_page_header_v2.repetition_levels_byte_length;
 
   auto bytes = pageHeader.compressed_page_size;
+  VELOX_CHECK_LE(
+      static_cast<uint64_t>(repeatLength) + defineLength,
+      bytes,
+      "Repetition and definition level lengths ({} + {}) exceed compressed page size {} (corrupt data page?)",
+      repeatLength,
+      defineLength,
+      bytes);
   pageData_ = readBytes(bytes, pageBuffer_);
 
   if (repeatLength) {

velox/dwio/parquet/reader/PageReader.h

@@ -65,11 +65,13 @@ class PageReader {
       common::CompressionKind codec,
       int64_t chunkSize,
       dwio::common::ColumnReaderStatistics& stats,
-      const tz::TimeZone* sessionTimezone = nullptr)
+      const tz::TimeZone* sessionTimezone = nullptr,
+      int32_t maxRepeat = 0,
+      int32_t maxDefine = 1)
       : pool_(pool),
         inputStream_(std::move(stream)),
-        maxRepeat_(0),
-        maxDefine_(1),
+        maxRepeat_(maxRepeat),
+        maxDefine_(maxDefine),
         isTopLevel_(maxRepeat_ == 0 && maxDefine_ <= 1),
         codec_(codec),
         chunkSize_(chunkSize),