Fix DeltaBitPackDecoder fuzzer issue.

apache · Feb 8, 2022 · 5f5cb6a · 5f5cb6a
1 parent 11caf00
commit 5f5cb6a
Showing 1 changed file with 18 additions and 11 deletions.
diff --git a/cpp/src/parquet/encoding.cc b/cpp/src/parquet/encoding.cc
@@ -2109,9 +2109,9 @@ class DeltaBitPackDecoder : public DecoderImpl, virtual public TypedDecoder<DTyp
       ParquetException::NYI("Delta bit pack DecodeArrow with null slots");
     }
     std::vector<T> values(num_values);
-    GetInternal(values.data(), num_values);
-    PARQUET_THROW_NOT_OK(out->AppendValues(values));
-    return num_values;
+    int decoded_count = GetInternal(values.data(), num_values);
+    PARQUET_THROW_NOT_OK(out->AppendValues(values.data(), decoded_count));
+    return decoded_count;
   }
 
   int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits,
@@ -2121,12 +2121,12 @@ class DeltaBitPackDecoder : public DecoderImpl, virtual public TypedDecoder<DTyp
       ParquetException::NYI("Delta bit pack DecodeArrow with null slots");
     }
     std::vector<T> values(num_values);
-    GetInternal(values.data(), num_values);
-    PARQUET_THROW_NOT_OK(out->Reserve(num_values));
-    for (T value : values) {
-      PARQUET_THROW_NOT_OK(out->Append(value));
+    int decoded_count = GetInternal(values.data(), num_values);
+    PARQUET_THROW_NOT_OK(out->Reserve(decoded_count));
+    for (int i = 0; i < decoded_count; ++i) {
+      PARQUET_THROW_NOT_OK(out->Append(values[i]));
     }
-    return num_values;
+    return decoded_count;
   }
 
  private:
@@ -2181,12 +2181,14 @@ class DeltaBitPackDecoder : public DecoderImpl, virtual public TypedDecoder<DTyp
   }
 
   int GetInternal(T* buffer, int max_values) {
-    max_values = std::min(max_values, this->num_values_);
+    int total_value_count = static_cast<int>(std::min(
+        total_value_count_, static_cast<uint32_t>(std::numeric_limits<int>::max())));
+    max_values = std::min(max_values, total_value_count);
+    DCHECK_GE(max_values, 0);
     if (max_values == 0) {
       return 0;
     }
 
-    DCHECK_LE(static_cast<uint32_t>(max_values), total_value_count_);
     int i = 0;
     while (i < max_values) {
       if (ARROW_PREDICT_FALSE(values_current_mini_block_ == 0)) {
@@ -2424,7 +2426,12 @@ class DeltaByteArrayDecoder : public DecoderImpl,
       return max_values;
     }
 
-    suffix_decoder_.Decode(buffer, max_values);
+    int suffix_read = suffix_decoder_.Decode(buffer, max_values);
+    if (ARROW_PREDICT_FALSE(suffix_read != max_values)) {
+      ParquetException::EofException("Read " + std::to_string(suffix_read) +
+                                     "Expecting " + std::to_string(max_values) +
+                                     " from suffix decoder");
+    }
 
     int64_t data_size = 0;
     const int32_t* prefix_len_ptr =