GH-55388: [C++][Parquet] Reject out-of-bounds dictionary indices in all RLE_DICTIONARY decode paths

cpp/src/arrow/util/rle_encoding_internal.h

@@ -928,12 +928,16 @@ auto RleBitPackedDecoder<T>::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<T>::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<T>::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<T>::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<T>::GetSpaced(Converter converter,
         typename RleBitPackedDecoderGetRunDecoder<value_type, decltype(run)>::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<T>::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;
     }

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<int32_t> 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<uint8_t> 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<int32_t> decoder(buffer.data(), static_cast<int>(buffer.size()),
+                                         8);
+    // Leave the out-of-bounds run partially read
+    std::vector<int32_t> 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<int32_t> out(8);
+    const std::vector<uint8_t> 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) {

cpp/src/parquet/decoder.cc

@@ -973,7 +973,11 @@ class DictDecoderImpl : public TypedDecoderImpl<Type>, public DictDecoder<Type>
     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<Type>, public DictDecoder<Type>
     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<Type>, public DictDecoder<Type>
     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<Type>, public DictDecoder<Type>
     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<Type>* dictionary) {
     dictionary_length_ = static_cast<int32_t>(dictionary->values_left());
     PARQUET_THROW_NOT_OK(

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<ColumnDescriptor>(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<uint32_t>(value.size());
+      dict_data_.append(reinterpret_cast<const char*>(&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<uint8_t>((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<uint8_t>(i % kDictLength));
+    }
+    indices_.push_back(/*RLE run of 8 values=*/8 << 1);
+    indices_.push_back(/*index=*/3);
+    num_values_ = static_cast<int>(num_groups * 8 + 8);
+  }
+
+  std::unique_ptr<DictDecoder<ByteArrayType>> MakeDecoder() {
+    plain_decoder_ = MakeTypedDecoder<ByteArrayType>(Encoding::PLAIN, descr_.get());
+    plain_decoder_->SetData(kDictLength,
+                            reinterpret_cast<const uint8_t*>(dict_data_.data()),
+                            static_cast<int>(dict_data_.size()));
+    auto decoder = MakeDictDecoder<ByteArrayType>(descr_.get());
+    decoder->SetDict(plain_decoder_.get());
+    decoder->SetData(num_values_, indices_.data(), static_cast<int>(indices_.size()));
+    return decoder;
+  }
+
+ protected:
+  static constexpr int kDictLength = 10;
+  std::unique_ptr<ColumnDescriptor> descr_;
+  std::string dict_data_;
+  std::unique_ptr<TypedDecoder<ByteArrayType>> plain_decoder_;
+  std::vector<uint8_t> 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<ByteArray> out(8);
+  {
+    auto decoder = MakeDecoder();
+    ASSERT_EQ(8, decoder->Decode(out.data(), 8));
+    EXPECT_THROW(decoder->Decode(out.data(), 8), ParquetException);
+  }
+  {
+    const std::vector<uint8_t> 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<uint8_t> 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<int32_t> 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<uint8_t> 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.