diff --git a/cpp/src/arrow/util/compression.h b/cpp/src/arrow/util/compression.h index 4634f35d61bd..a8a897caf85f 100644 --- a/cpp/src/arrow/util/compression.h +++ b/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> 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> MakeCompressor() = 0; diff --git a/cpp/src/arrow/util/compression_snappy.cc b/cpp/src/arrow/util/compression_snappy.cc index 6ebb566a361b..431200ef80b6 100644 --- a/cpp/src/arrow/util/compression_snappy.cc +++ b/cpp/src/arrow/util/compression_snappy.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include @@ -44,6 +45,17 @@ namespace { class SnappyCodec : public Codec { public: + Result> DecompressedLength(int64_t input_len, + const uint8_t* input) override { + size_t decompressed_size = 0; + if (!snappy::GetUncompressedLength(reinterpret_cast(input), + static_cast(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(decompressed_size); + } + Result Decompress(int64_t input_len, const uint8_t* input, int64_t output_buffer_len, uint8_t* output_buffer) override { size_t decompressed_size; diff --git a/cpp/src/arrow/util/compression_zstd.cc b/cpp/src/arrow/util/compression_zstd.cc index cd9f0ce66fd3..17b4466fbb8b 100644 --- a/cpp/src/arrow/util/compression_zstd.cc +++ b/cpp/src/arrow/util/compression_zstd.cc @@ -19,6 +19,8 @@ #include #include +#include +#include #include #include #include @@ -158,6 +160,22 @@ class ZSTDCompressor : public Compressor { class ZSTDCodec : public Codec { public: + Result> DecompressedLength(int64_t input_len, + const uint8_t* input) override { + const unsigned long long size = ZSTD_getFrameContentSize(input, static_cast(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(std::numeric_limits::max())) { + return Status::IOError("ZSTD frame declares an unrepresentable decompressed size"); + } + return static_cast(size); + } + explicit ZSTDCodec(int compression_level, std::vector> compression_context_params, std::vector> decompression_context_params) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 060e860865ad..b8e97b4aca24 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -463,7 +464,6 @@ std::shared_ptr 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 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(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));