GH-48146: [C++][Parquet] Fix undefined behavior with invalid column/offset index (#48147)

### Rationale for this change

This should fix the following issues found by OSS-Fuzz:
* https://issues.oss-fuzz.com/issues/461058054
* https://issues.oss-fuzz.com/issues/461058060

Also fixes a signed integer overflow issue in `CappedMemoryPool`:
* https://issues.oss-fuzz.com/issues/461314335

### Are these changes tested?

Yes, by existing tests and additional fuzz regression files.

### Are there any user-facing changes?

No.

**This PR contains a "Critical Fix".** (If the changes fix either (a) a security vulnerability, (b) a bug that caused incorrect or invalid data to be produced, or (c) a bug that causes a crash (even when the API contract is upheld), please provide explanation. If not, you can remove this.)

* GitHub Issue: #48146

Authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>

Author: pitrou

cpp/src/arrow/memory_pool.cc

@@ -857,6 +857,45 @@ std::vector<std::string> SupportedMemoryBackendNames() {
   return supported;
 }
 
+///////////////////////////////////////////////////////////////////////
+// CappedMemoryPool implementation
+
+Status CappedMemoryPool::Allocate(int64_t size, int64_t alignment, uint8_t** out) {
+  // XXX Another thread may allocate memory between the limit check and
+  // the `Allocate` call. It is possible for the two allocations to be successful
+  // while going above the limit.
+  // Solving this issue would require refactoring the `MemoryPool` implementation
+  // to delegate the limit check to `MemoryPoolStats`.
+  const int64_t allocated = wrapped_->bytes_allocated();
+  if (ARROW_PREDICT_FALSE(bytes_allocated_limit_ - allocated < size)) {
+    return OutOfMemory(allocated, size);
+  }
+  return wrapped_->Allocate(size, alignment, out);
+}
+
+Status CappedMemoryPool::Reallocate(int64_t old_size, int64_t new_size, int64_t alignment,
+                                    uint8_t** ptr) {
+  if (new_size > old_size) {
+    const int64_t allocated = wrapped_->bytes_allocated();
+    if (ARROW_PREDICT_FALSE(bytes_allocated_limit_ - allocated < new_size - old_size)) {
+      return OutOfMemory(allocated, new_size - old_size);
+    }
+  }
+  return wrapped_->Reallocate(old_size, new_size, alignment, ptr);
+}
+
+void CappedMemoryPool::Free(uint8_t* buffer, int64_t size, int64_t alignment) {
+  return wrapped_->Free(buffer, size, alignment);
+}
+
+Status CappedMemoryPool::OutOfMemory(int64_t current_allocated, int64_t requested) const {
+  return Status::OutOfMemory(
+      "MemoryPool bytes_allocated cap exceeded: "
+      "limit=",
+      bytes_allocated_limit_, ", current allocation=", current_allocated,
+      ", requested=", requested);
+}
+
 // -----------------------------------------------------------------------
 // Pool buffer and allocation
 

cpp/src/arrow/memory_pool.h

@@ -259,31 +259,10 @@ class ARROW_EXPORT CappedMemoryPool : public MemoryPool {
   using MemoryPool::Allocate;
   using MemoryPool::Reallocate;
 
-  Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override {
-    // XXX Another thread may allocate memory between the limit check and
-    // the `Allocate` call. It is possible for the two allocations to be successful
-    // while going above the limit.
-    // Solving this issue would require refactoring the `MemoryPool` implementation
-    // to delegate the limit check to `MemoryPoolStats`.
-    const auto attempted = size + wrapped_->bytes_allocated();
-    if (ARROW_PREDICT_FALSE(attempted > bytes_allocated_limit_)) {
-      return OutOfMemory(attempted);
-    }
-    return wrapped_->Allocate(size, alignment, out);
-  }
-
+  Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override;
   Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment,
-                    uint8_t** ptr) override {
-    const auto attempted = new_size - old_size + wrapped_->bytes_allocated();
-    if (ARROW_PREDICT_FALSE(attempted > bytes_allocated_limit_)) {
-      return OutOfMemory(attempted);
-    }
-    return wrapped_->Reallocate(old_size, new_size, alignment, ptr);
-  }
-
-  void Free(uint8_t* buffer, int64_t size, int64_t alignment) override {
-    return wrapped_->Free(buffer, size, alignment);
-  }
+                    uint8_t** ptr) override;
+  void Free(uint8_t* buffer, int64_t size, int64_t alignment) override;
 
   void ReleaseUnused() override { wrapped_->ReleaseUnused(); }
 
@@ -302,12 +281,7 @@ class ARROW_EXPORT CappedMemoryPool : public MemoryPool {
   std::string backend_name() const override { return wrapped_->backend_name(); }
 
  private:
-  Status OutOfMemory(int64_t value) {
-    return Status::OutOfMemory(
-        "MemoryPool bytes_allocated cap exceeded: "
-        "limit=",
-        bytes_allocated_limit_, ", attempted=", value);
-  }
+  Status OutOfMemory(int64_t current_allocated, int64_t requested) const;
 
   MemoryPool* wrapped_;
   const int64_t bytes_allocated_limit_;

cpp/src/parquet/page_index.cc

@@ -244,9 +244,9 @@ class RowGroupPageIndexReaderImpl : public RowGroupPageIndexReader {
                           row_group_ordinal_);
 
     if (column_index_buffer_ == nullptr) {
-      PARQUET_ASSIGN_OR_THROW(column_index_buffer_,
-                              input_->ReadAt(index_read_range_.column_index->offset,
-                                             index_read_range_.column_index->length));
+      column_index_buffer_ =
+          ReadIndexBuffer(index_read_range_.column_index->offset,
+                          index_read_range_.column_index->length, "ColumnIndex");
     }
 
     int64_t buffer_offset =
@@ -285,9 +285,9 @@ class RowGroupPageIndexReaderImpl : public RowGroupPageIndexReader {
                           row_group_ordinal_);
 
     if (offset_index_buffer_ == nullptr) {
-      PARQUET_ASSIGN_OR_THROW(offset_index_buffer_,
-                              input_->ReadAt(index_read_range_.offset_index->offset,
-                                             index_read_range_.offset_index->length));
+      offset_index_buffer_ =
+          ReadIndexBuffer(index_read_range_.offset_index->offset,
+                          index_read_range_.offset_index->length, "OffsetIndex");
     }
 
     int64_t buffer_offset =
@@ -346,7 +346,16 @@ class RowGroupPageIndexReaderImpl : public RowGroupPageIndexReader {
     }
   }
 
- private:
+  std::shared_ptr<Buffer> ReadIndexBuffer(int64_t offset, int64_t length,
+                                          const char* offset_kind) {
+    PARQUET_ASSIGN_OR_THROW(auto buffer, input_->ReadAt(offset, length));
+    if (buffer->size() < length) {
+      throw ParquetException("Invalid or truncated ", offset_kind, ": attempted to read ",
+                             length, " bytes but got only ", buffer->size(), " bytes");
+    }
+    return buffer;
+  }
+
   /// The input stream that can perform random access read.
   ::arrow::io::RandomAccessFile* input_;
 
@@ -962,6 +971,11 @@ std::unique_ptr<ColumnIndex> ColumnIndex::Make(const ColumnDescriptor& descr,
   ThriftDeserializer deserializer(properties);
   deserializer.DeserializeMessage(reinterpret_cast<const uint8_t*>(serialized_index),
                                   &index_len, &column_index, decryptor);
+  if (ARROW_PREDICT_FALSE(LoadEnumSafe(&column_index.boundary_order) ==
+                          BoundaryOrder::UNDEFINED)) {
+    // Guard against UB when moving column_index
+    throw ParquetException("Invalid ColumnIndex boundary_order");
+  }
   switch (descr.physical_type()) {
     case Type::BOOLEAN:
       return std::make_unique<TypedColumnIndexImpl<BooleanType>>(descr,

cpp/src/parquet/thrift_internal.h

@@ -186,32 +186,22 @@ struct SafeLoader {
     return static_cast<ApiTypeRawEnum>(LoadEnumRaw(in));
   }
 
-  template <typename ThriftType, bool IsUnsigned = true>
-  inline static ApiTypeEnum LoadChecked(
-      const typename std::enable_if<IsUnsigned, ThriftType>::type* in) {
-    auto raw_value = LoadRaw(in);
-    if (ARROW_PREDICT_FALSE(raw_value >=
-                            static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED))) {
-      return ApiType::UNDEFINED;
-    }
-    return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
-  }
-
-  template <typename ThriftType, bool IsUnsigned = false>
-  inline static ApiTypeEnum LoadChecked(
-      const typename std::enable_if<!IsUnsigned, ThriftType>::type* in) {
-    auto raw_value = LoadRaw(in);
-    if (ARROW_PREDICT_FALSE(raw_value >=
-                                static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED) ||
-                            raw_value < 0)) {
-      return ApiType::UNDEFINED;
-    }
-    return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
-  }
-
   template <typename ThriftType>
   inline static ApiTypeEnum Load(const ThriftType* in) {
-    return LoadChecked<ThriftType, std::is_unsigned<ApiTypeRawEnum>::value>(in);
+    const auto raw_value = LoadRaw(in);
+    if constexpr (std::is_unsigned_v<ApiTypeRawEnum>) {
+      if (ARROW_PREDICT_FALSE(raw_value >=
+                              static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED))) {
+        return ApiType::UNDEFINED;
+      }
+    } else {
+      if (ARROW_PREDICT_FALSE(raw_value >=
+                                  static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED) ||
+                              raw_value < 0)) {
+        return ApiType::UNDEFINED;
+      }
+    }
+    return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
   }
 };