diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd65306..4037dedb066c 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -928,12 +928,16 @@ auto RleBitPackedDecoder::ProcessValues(Callable&& func, // Remaining from a previous call that would have left some unread data from a run. if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { + const auto run_available = run_remaining(); const auto processed = std::visit([&](auto& decoder) { return func(decoder, batch_size); }, decoder_); values_processed += processed; - // Either we fulfilled all the batch to be read or we finished remaining run. - if (ARROW_PREDICT_FALSE(values_processed == batch_size)) { + // Either we fulfilled all the batch to be read, or the run stopped short + // because of invalid data (e.g. an out-of-bounds dictionary index), in which + // case we must not resume with the next run. + if (ARROW_PREDICT_FALSE(values_processed == batch_size || + processed < run_available)) { return values_processed; } ARROW_DCHECK(run_remaining() == 0); @@ -946,11 +950,14 @@ auto RleBitPackedDecoder::ProcessValues(Callable&& func, ARROW_DCHECK_LT(values_processed, batch_size); // Local decoder to keep its type transparent to the compiler. RunDecoder decoder(run, value_bit_width_); + const auto run_available = decoder.remaining(); const auto read = func(decoder, batch_size - values_processed); ARROW_DCHECK_LE(read, batch_size - values_processed); values_processed += read; - if (ARROW_PREDICT_FALSE(values_processed == batch_size || read == 0)) { + // Stop if the batch is full, or if the run was not entirely consumed although + // the batch is not full, which means invalid data. + if (ARROW_PREDICT_FALSE(values_processed == batch_size || read < run_available)) { // Stop reading and store remaining decoder decoder_ = std::move(decoder); return ControlFlow::Break; @@ -1277,6 +1284,7 @@ auto RleBitPackedDecoder::GetSpaced(Converter converter, // Remaining from a previous call that would have left some unread data from a run. if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { + const auto run_available = run_remaining(); const auto read = internal::RunGetSpaced(&converter, out, batch.total_remaining(), batch.null_remaining(), value_bit_width_, &validity_reader, &validity_run, &decoder_); @@ -1292,6 +1300,12 @@ auto RleBitPackedDecoder::GetSpaced(Converter converter, return batch.total_read(); } + // Or the run stopped short because of invalid data (e.g. an out-of-bounds + // dictionary index), in which case we must not resume with the next run. + if (ARROW_PREDICT_FALSE(read.values_read < run_available)) { + return batch.total_read(); + } + // We finished the remaining run ARROW_DCHECK(run_remaining() == 0); } @@ -1301,6 +1315,7 @@ auto RleBitPackedDecoder::GetSpaced(Converter converter, typename RleBitPackedDecoderGetRunDecoder::type; RunDecoder decoder(run, value_bit_width_); + const auto run_available = decoder.remaining(); const auto read = internal::RunGetSpaced(&converter, out, batch.total_remaining(), batch.null_remaining(), value_bit_width_, @@ -1310,8 +1325,11 @@ auto RleBitPackedDecoder::GetSpaced(Converter converter, batch.AccrueReadValues(read.values_read); out += read.values_read + read.null_read; - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(read.values_read == 0 || batch.values_remaining() == 0)) { + // Stop reading and store remaining decoder if the batch values are all read, + // or if the run was not entirely consumed although some batch values remain, + // which means invalid data. + if (ARROW_PREDICT_FALSE(batch.values_remaining() == 0 || + read.values_read < run_available)) { decoder_ = std::move(decoder); return ControlFlow::Break; } diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index c795355f3ca4..6d9d8cf2e55b 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -898,6 +898,31 @@ TEST(RleBitPacked, BitWidthZeroLiteral) { EXPECT_FALSE(decoder.Get(&val)); } +TEST(RleBitPacked, GetBatchWithDictOutOfBoundsInPartialRun) { + const std::vector dict = {10, 11, 12}; + // Bit width 8: an RLE run of 8 out-of-bounds indices, then a bit-packed run of + // 8 valid indices. + const std::vector buffer = {8 << 1, 200, (1 << 1) | 1, 0, 1, 2, 0, 1, 2, 0, 1}; + for (bool spaced : {false, true}) { + ARROW_SCOPED_TRACE("spaced = ", spaced); + RleBitPackedDecoder decoder(buffer.data(), static_cast(buffer.size()), + 8); + // Leave the out-of-bounds run partially read + std::vector indices(4); + ASSERT_EQ(decoder.GetBatch(indices.data(), 4), 4); + + // Decoding must stop at the out-of-bounds run instead of resuming with the + // next run + std::vector out(8); + const std::vector valid_bits(1, 0xff); + const auto read = + spaced ? decoder.GetBatchWithDictSpaced(dict.data(), 3, out.data(), 8, + /*null_count=*/0, valid_bits.data(), 0) + : decoder.GetBatchWithDict(dict.data(), 3, out.data(), 8); + ASSERT_EQ(read, 0); + } +} + // Test that writes out a repeated group and then a literal // group but flush before finishing. TEST(BitRle, Flush) { diff --git a/cpp/src/parquet/decoder.cc b/cpp/src/parquet/decoder.cc index 8da344379c63..7410acfbcfd1 100644 --- a/cpp/src/parquet/decoder.cc +++ b/cpp/src/parquet/decoder.cc @@ -973,7 +973,11 @@ class DictDecoderImpl : public TypedDecoderImpl, public DictDecoder size_t i = 0; VisitNullBitmapInline( valid_bits, valid_bits_offset, num_values, null_count, - [&]() { valid_bytes[i++] = 1; }, [&]() { ++i; }); + [&]() { + PARQUET_THROW_NOT_OK(IndexInBounds(indices_buffer[i])); + valid_bytes[i++] = 1; + }, + [&]() { ++i; }); auto binary_builder = checked_cast<::arrow::BinaryDictionary32Builder*>(builder); PARQUET_THROW_NOT_OK( @@ -995,6 +999,7 @@ class DictDecoderImpl : public TypedDecoderImpl, public DictDecoder if (num_values != idx_decoder_.GetBatch(indices_buffer, num_values)) { ParquetException::EofException(); } + CheckIndicesInBounds(indices_buffer, num_values); auto binary_builder = checked_cast<::arrow::BinaryDictionary32Builder*>(builder); PARQUET_THROW_NOT_OK(binary_builder->AppendIndices(indices_buffer, num_values)); this->num_values_ -= num_values; @@ -1005,6 +1010,7 @@ class DictDecoderImpl : public TypedDecoderImpl, public DictDecoder if (num_values != idx_decoder_.GetBatch(indices, num_values)) { ParquetException::EofException(); } + CheckIndicesInBounds(indices, num_values); this->num_values_ -= num_values; return num_values; } @@ -1022,6 +1028,12 @@ class DictDecoderImpl : public TypedDecoderImpl, public DictDecoder return Status::Invalid("Index not in dictionary bounds"); } + void CheckIndicesInBounds(const int32_t* indices, int num_values) const { + for (int i = 0; i < num_values; ++i) { + PARQUET_THROW_NOT_OK(IndexInBounds(indices[i])); + } + } + inline void DecodeDict(TypedDecoder* dictionary) { dictionary_length_ = static_cast(dictionary->values_left()); PARQUET_THROW_NOT_OK( diff --git a/cpp/src/parquet/encoding_test.cc b/cpp/src/parquet/encoding_test.cc index 672e55a900dc..b7ed9cd528a2 100644 --- a/cpp/src/parquet/encoding_test.cc +++ b/cpp/src/parquet/encoding_test.cc @@ -1412,6 +1412,114 @@ TEST_F(DictEncoding, CheckDecodeIndicesNoNulls) { CheckDict(actual_num_values, *builder); } +// Out-of-bounds dictionary indices must be rejected by all decoding methods. +class DictDecodingOutOfBounds : public ::testing::Test { + public: + void SetUp() override { + descr_ = std::make_unique(schema::ByteArray("name"), 0, 0); + for (int i = 0; i < kDictLength; ++i) { + const std::string value = std::to_string(i); + const auto length = static_cast(value.size()); + dict_data_.append(reinterpret_cast(&length), sizeof(length)); + dict_data_ += value; + } + } + + // RLE_DICTIONARY data with a bit width of 8: a bit-packed run of + // `num_groups * 8` indices with an out-of-bounds index at `bad_position`, + // followed by an RLE run of 8 valid indices. + void MakeIndices(uint32_t num_groups, uint32_t bad_position) { + indices_ = {/*bit_width=*/8}; + // ULEB128-encoded bit-packed run header + for (uint32_t header = (num_groups << 1) | 1; header != 0; header >>= 7) { + indices_.push_back( + static_cast((header & 0x7f) | (header > 0x7f ? 0x80 : 0))); + } + for (uint32_t i = 0; i < num_groups * 8; ++i) { + indices_.push_back(i == bad_position ? 200 : static_cast(i % kDictLength)); + } + indices_.push_back(/*RLE run of 8 values=*/8 << 1); + indices_.push_back(/*index=*/3); + num_values_ = static_cast(num_groups * 8 + 8); + } + + std::unique_ptr> MakeDecoder() { + plain_decoder_ = MakeTypedDecoder(Encoding::PLAIN, descr_.get()); + plain_decoder_->SetData(kDictLength, + reinterpret_cast(dict_data_.data()), + static_cast(dict_data_.size())); + auto decoder = MakeDictDecoder(descr_.get()); + decoder->SetDict(plain_decoder_.get()); + decoder->SetData(num_values_, indices_.data(), static_cast(indices_.size())); + return decoder; + } + + protected: + static constexpr int kDictLength = 10; + std::unique_ptr descr_; + std::string dict_data_; + std::unique_ptr> plain_decoder_; + std::vector indices_; + int num_values_ = 0; +}; + +TEST_F(DictDecodingOutOfBounds, Decode) { + // The out-of-bounds index is in a bit-packed run left partially read by a previous + // call; decoding must not skip over it and resume with the next run. + MakeIndices(/*num_groups=*/2, /*bad_position=*/10); + std::vector out(8); + { + auto decoder = MakeDecoder(); + ASSERT_EQ(8, decoder->Decode(out.data(), 8)); + EXPECT_THROW(decoder->Decode(out.data(), 8), ParquetException); + } + { + const std::vector valid_bits(1, 0xff); + auto decoder = MakeDecoder(); + ASSERT_EQ(8, decoder->DecodeSpaced(out.data(), 8, 0, valid_bits.data(), 0)); + EXPECT_THROW(decoder->DecodeSpaced(out.data(), 8, 0, valid_bits.data(), 0), + ParquetException); + } + // The out-of-bounds index is in a long bit-packed run, past the values that + // were decoded successfully. + MakeIndices(/*num_groups=*/129, /*bad_position=*/1028); + out.resize(1032); + { + auto decoder = MakeDecoder(); + EXPECT_THROW(decoder->Decode(out.data(), 1032), ParquetException); + } + { + const std::vector valid_bits(1032 / 8, 0xff); + auto decoder = MakeDecoder(); + EXPECT_THROW(decoder->DecodeSpaced(out.data(), 1032, 0, valid_bits.data(), 0), + ParquetException); + } +} + +TEST_F(DictDecodingOutOfBounds, DecodeIndices) { + MakeIndices(/*num_groups=*/2, /*bad_position=*/10); + { + std::vector out(16); + auto decoder = MakeDecoder(); + EXPECT_THROW(decoder->DecodeIndices(16, out.data()), ParquetException); + } + { + ::arrow::BinaryDictionary32Builder builder(default_memory_pool()); + auto decoder = MakeDecoder(); + decoder->InsertDictionary(&builder); + EXPECT_THROW(decoder->DecodeIndices(16, &builder), ParquetException); + } + { + // 16 non-null slots, including the out-of-bounds index, then a null slot + const std::vector valid_bits = {0xff, 0xff, 0x00}; + ::arrow::BinaryDictionary32Builder builder(default_memory_pool()); + auto decoder = MakeDecoder(); + decoder->InsertDictionary(&builder); + EXPECT_THROW(decoder->DecodeIndicesSpaced(17, 1, valid_bits.data(), 0, &builder), + ParquetException); + } +} + // ---------------------------------------------------------------------- // BYTE_STREAM_SPLIT encode/decode tests.