fix(parquet): size a page from what its data says, not what its header claims (FB-4054)

cpp/src/arrow/util/compression.h

@@ -230,6 +230,28 @@ class ARROW_EXPORT Codec {
 
   virtual int64_t MaxCompressedLen(int64_t input_len, const uint8_t* input) = 0;
 
+  /// \brief The decompressed length the compressed data states for itself
+  ///
+  /// Some formats record it -- snappy in a varint prefix, zstd in the frame header -- so a caller
+  /// sizing a buffer can use what the data says rather than a length handed to it out of band.
+  ///
+  /// Three outcomes, and they mean different things:
+  ///   - a value: the data states this length
+  ///   - nullopt: this format does not record one (lz4 raw blocks, gzip), or this writer left it
+  ///     unset. The caller learns nothing and must carry on as before.
+  ///   - an error: the format does record one, and this input is too damaged to read it. The
+  ///     input cannot decompress, so a caller sizing a buffer for it should stop rather than
+  ///     allocate on the strength of a number from elsewhere.
+  ///
+  /// A returned length is read from attacker-controllable bytes. It bounds what the caller should
+  /// be willing to allocate; it does not promise the data decompresses, which Decompress checks.
+  virtual Result<std::optional<int64_t>> DecompressedLength(int64_t input_len,
+                                                            const uint8_t* input) {
+    ARROW_UNUSED(input_len);
+    ARROW_UNUSED(input);
+    return std::nullopt;
+  }
+
   /// \brief Create a streaming compressor instance
   virtual Result<std::shared_ptr<Compressor>> MakeCompressor() = 0;
 

cpp/src/arrow/util/compression_snappy.cc

@@ -21,6 +21,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 
 #include <snappy-sinksource.h>
 #include <snappy.h>
@@ -44,6 +45,17 @@ namespace {
 
 class SnappyCodec : public Codec {
  public:
+  Result<std::optional<int64_t>> DecompressedLength(int64_t input_len,
+                                                    const uint8_t* input) override {
+    size_t decompressed_size = 0;
+    if (!snappy::GetUncompressedLength(reinterpret_cast<const char*>(input),
+                                       static_cast<size_t>(input_len), &decompressed_size)) {
+      // Snappy always records the length, so failing to read it means the input is not snappy.
+      return Status::IOError("Corrupt snappy compressed data.");
+    }
+    return static_cast<int64_t>(decompressed_size);
+  }
+
   Result<int64_t> Decompress(int64_t input_len, const uint8_t* input,
                              int64_t output_buffer_len, uint8_t* output_buffer) override {
     size_t decompressed_size;

cpp/src/arrow/util/compression_zstd.cc

@@ -19,6 +19,8 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <limits>
+#include <optional>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -158,6 +160,22 @@ class ZSTDCompressor : public Compressor {
 
 class ZSTDCodec : public Codec {
  public:
+  Result<std::optional<int64_t>> DecompressedLength(int64_t input_len,
+                                                    const uint8_t* input) override {
+    const unsigned long long size = ZSTD_getFrameContentSize(input, static_cast<size_t>(input_len));
+    if (size == ZSTD_CONTENTSIZE_ERROR) {
+      // Not a readable zstd frame header at all, so nothing can decompress from it.
+      return Status::IOError("ZSTD decompression failed: not a valid frame");
+    }
+    if (size == ZSTD_CONTENTSIZE_UNKNOWN) {
+      return std::nullopt;  // a valid frame whose writer chose not to record the size
+    }
+    if (size > static_cast<unsigned long long>(std::numeric_limits<int64_t>::max())) {
+      return Status::IOError("ZSTD frame declares an unrepresentable decompressed size");
+    }
+    return static_cast<int64_t>(size);
+  }
+
   explicit ZSTDCodec(int compression_level,
                      std::vector<std::pair<int, int>> compression_context_params,
                      std::vector<std::pair<int, int>> decompression_context_params)

cpp/src/parquet/column_reader.cc

@@ -23,6 +23,7 @@
 #include <exception>
 #include <iostream>
 #include <memory>
+#include <optional>
 #include <string>
 #include <type_traits>
 #include <unordered_map>
@@ -463,7 +464,6 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() {
     if (compressed_len < 0 || uncompressed_len < 0) {
       throw ParquetException("Invalid page header");
     }
-
     EncodedStatistics data_page_statistics;
     if (ShouldSkipPage(&data_page_statistics)) {
       PARQUET_THROW_NOT_OK(stream_->Advance(compressed_len));
@@ -641,6 +641,36 @@ std::shared_ptr<Buffer> SerializedPageReader::DecompressIfNeeded(
     throw ParquetException("Invalid page header");
   }
 
+  // The header's uncompressed size is the file's word for it, and it sizes the buffer below --
+  // a 1358-byte page was enough to ask for 2 GiB. Ask the data itself instead: snappy records its
+  // decompressed length in a varint prefix and zstd in the frame header.
+  //
+  // The error case is the one that matters. A page whose compressed bytes cannot even be parsed
+  // as the codec they claim to be will not decompress, so there is no reason to allocate for it
+  // first -- and that is exactly the shape of the attack, five bytes of noise declaring 2 GiB.
+  // Formats that record no length (lz4 raw blocks, gzip) return nullopt and are left alone, as is
+  // the corded path, whose bytes are not contiguous to hand over.
+  if (plain_buffer != nullptr) {
+    const int64_t bytes_to_decompress = compressed_len - levels_byte_len;
+    if (bytes_to_decompress > 0) {
+      auto stated = decompressor_->DecompressedLength(
+          bytes_to_decompress, plain_buffer->data() + levels_byte_len);
+      if (!stated.ok()) {
+        throw ParquetException("Page data is not readable as " +
+                               ::arrow::util::Codec::GetCodecAsString(
+                                   decompressor_->compression_type()) +
+                               ": " + stated.status().message());
+      }
+      const int64_t claimed = static_cast<int64_t>(uncompressed_len) - levels_byte_len;
+      if (stated->has_value() && **stated != claimed) {
+        std::stringstream ss;
+        ss << "Page header declares an uncompressed size of " << claimed
+           << " bytes, but the compressed data states " << **stated;
+        throw ParquetException(ss.str());
+      }
+    }
+  }
+
   // Grow the uncompressed buffer if we need to.
   PARQUET_THROW_NOT_OK(
       decompression_buffer_->Resize(uncompressed_len, /*shrink_to_fit=*/false));