GH-48334: [C++][Parquet] Support reading encrypted bloom filters

cpp/src/parquet/CMakeLists.txt

@@ -416,6 +416,7 @@ add_parquet_test(arrow-metadata-test SOURCES arrow/arrow_metadata_test.cc
 if(PARQUET_REQUIRE_ENCRYPTION)
   add_parquet_test(encryption-test
                    SOURCES
+                   encryption/bloom_filter_encryption_test.cc
                    encryption/encryption_internal_test.cc
                    encryption/write_configurations_test.cc
                    encryption/read_configurations_test.cc

cpp/src/parquet/arrow/fuzz_internal.cc

@@ -155,6 +155,23 @@ Status FuzzReadPageIndex(RowGroupPageIndexReader* reader, const SchemaDescriptor
   return st;
 }
 
+Status FuzzReadBloomFilter(RowGroupBloomFilterReader* reader, int column,
+                           std::uniform_int_distribution<uint64_t>& hash_dist,
+                           std::default_random_engine& rng) {
+  Status st;
+  BEGIN_PARQUET_CATCH_EXCEPTIONS
+  std::unique_ptr<BloomFilter> bloom;
+  bloom = reader->GetColumnBloomFilter(column);
+  // If the column has a bloom filter, find a bunch of random hashes
+  if (bloom != nullptr) {
+    for (int k = 0; k < 100; ++k) {
+      bloom->FindHash(hash_dist(rng));
+    }
+  }
+  END_PARQUET_CATCH_EXCEPTIONS
+  return st;
+}
+
 ReaderProperties MakeFuzzReaderProperties(MemoryPool* pool) {
   FileDecryptionProperties::Builder builder;
   builder.key_retriever(MakeKeyRetriever());
@@ -207,31 +224,12 @@ Status FuzzReader(const uint8_t* data, int64_t size) {
     }
     {
       // Read and decode bloom filters
-      try {
-        auto& bloom_reader = pq_file_reader->GetBloomFilterReader();
-        std::uniform_int_distribution<uint64_t> hash_dist;
-        for (int i = 0; i < num_row_groups; ++i) {
-          auto bloom_rg = bloom_reader.RowGroup(i);
-          for (int j = 0; j < num_columns; ++j) {
-            std::unique_ptr<BloomFilter> bloom;
-            bloom = bloom_rg->GetColumnBloomFilter(j);
-            // If the column has a bloom filter, find a bunch of random hashes
-            if (bloom != nullptr) {
-              for (int k = 0; k < 100; ++k) {
-                bloom->FindHash(hash_dist(rng));
-              }
-            }
-          }
-        }
-      } catch (const ParquetException& exc) {
-        // XXX we just want to ignore encrypted bloom filters and validate the
-        // rest of the file; there is no better way of doing this until GH-46597
-        // is done.
-        // (also see GH-48334 for reading encrypted bloom filters)
-        if (std::string_view(exc.what())
-                .find("BloomFilter decryption is not yet supported") ==
-            std::string_view::npos) {
-          throw;
+      auto& bloom_reader = pq_file_reader->GetBloomFilterReader();
+      std::uniform_int_distribution<uint64_t> hash_dist;
+      for (int i = 0; i < num_row_groups; ++i) {
+        auto bloom_rg = bloom_reader.RowGroup(i);
+        for (int j = 0; j < num_columns; ++j) {
+          st &= FuzzReadBloomFilter(bloom_rg.get(), j, hash_dist, rng);
         }
       }
     }

cpp/src/parquet/bloom_filter.cc

@@ -17,20 +17,42 @@
 
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <memory>
 
+#include "arrow/io/memory.h"
 #include "arrow/result.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/macros.h"
 
 #include "generated/parquet_types.h"
 
 #include "parquet/bloom_filter.h"
+#include "parquet/encryption/encryption_internal.h"
+#include "parquet/encryption/internal_file_decryptor.h"
 #include "parquet/exception.h"
 #include "parquet/thrift_internal.h"
 #include "parquet/xxhasher.h"
 
 namespace parquet {
+namespace {
+
+constexpr int32_t kCiphertextLengthSize = 4;
+
+/// Parse the 4-byte little-endian length prefix and return the total ciphertext size,
+/// including the 4-byte length field itself.
+int64_t ParseCiphertextTotalLength(const uint8_t* data, int64_t length) {
+  if (length < kCiphertextLengthSize) {
+    throw ParquetException("Ciphertext length buffer is too small");
+  }
+  uint32_t buffer_size =
+      (static_cast<uint32_t>(data[3]) << 24) | (static_cast<uint32_t>(data[2]) << 16) |
+      (static_cast<uint32_t>(data[1]) << 8) | (static_cast<uint32_t>(data[0]));
+  return static_cast<int64_t>(buffer_size) + kCiphertextLengthSize;
+}
+
+}  // namespace
+
 constexpr uint32_t BlockSplitBloomFilter::SALT[kBitsSetPerBlock];
 
 BlockSplitBloomFilter::BlockSplitBloomFilter(::arrow::MemoryPool* pool)
@@ -75,10 +97,11 @@ void BlockSplitBloomFilter::Init(const uint8_t* bitset, uint32_t num_bytes) {
   this->hasher_ = std::make_unique<XxHasher>();
 }
 
-static constexpr uint32_t kBloomFilterHeaderSizeGuess = 256;
+namespace {
 
-static ::arrow::Status ValidateBloomFilterHeader(
-    const format::BloomFilterHeader& header) {
+constexpr uint32_t kBloomFilterHeaderSizeGuess = 256;
+
+::arrow::Status ValidateBloomFilterHeader(const format::BloomFilterHeader& header) {
   if (!header.algorithm.__isset.BLOCK) {
     return ::arrow::Status::Invalid(
         "Unsupported Bloom filter algorithm: ", header.algorithm, ".");
@@ -104,6 +127,139 @@ static ::arrow::Status ValidateBloomFilterHeader(
   return ::arrow::Status::OK();
 }
 
+BlockSplitBloomFilter DeserializeEncryptedFromStream(
+    const ReaderProperties& properties, ArrowInputStream* input,
+    std::optional<int64_t> bloom_filter_length, Decryptor* decryptor,
+    int16_t row_group_ordinal, int16_t column_ordinal) {
+  ThriftDeserializer deserializer(properties);
+  format::BloomFilterHeader header;
+
+  // Read the length-prefixed ciphertext for the header.
+  PARQUET_ASSIGN_OR_THROW(auto length_buf, input->Read(kCiphertextLengthSize));
+  if (ARROW_PREDICT_FALSE(length_buf->size() < kCiphertextLengthSize)) {
+    std::stringstream ss;
+    ss << "Bloom filter header read failed: expected " << kCiphertextLengthSize
+       << " bytes, got " << length_buf->size();
+    throw ParquetException(ss.str());
+  }
+
+  const int64_t header_cipher_total_len =
+      ParseCiphertextTotalLength(length_buf->data(), length_buf->size());
+  if (ARROW_PREDICT_FALSE(header_cipher_total_len >
+                          std::numeric_limits<int32_t>::max())) {
+    throw ParquetException("Bloom filter header ciphertext length overflows int32");
+  }
+  if (bloom_filter_length && header_cipher_total_len > *bloom_filter_length) {
+    throw ParquetException(
+        "Bloom filter length less than encrypted bloom filter header length");
+  }
+
+  // Read the full header ciphertext and decrypt the Thrift header.
+  auto header_cipher_buf =
+      AllocateBuffer(properties.memory_pool(), header_cipher_total_len);
+  std::memcpy(header_cipher_buf->mutable_data(), length_buf->data(),
+              kCiphertextLengthSize);
+  const int64_t header_cipher_remaining = header_cipher_total_len - kCiphertextLengthSize;
+  PARQUET_ASSIGN_OR_THROW(auto read_size, input->Read(header_cipher_remaining,
+                                                      header_cipher_buf->mutable_data() +
+                                                          kCiphertextLengthSize));
+  if (ARROW_PREDICT_FALSE(read_size < header_cipher_remaining)) {
+    std::stringstream ss;
+    ss << "Bloom filter header read failed: expected " << header_cipher_remaining
+       << " bytes, got " << read_size;
+    throw ParquetException(ss.str());
+  }
+
+  // Bloom filter header and bitset are separate encrypted modules with different AADs.
+  UpdateDecryptor(decryptor, row_group_ordinal, column_ordinal,
+                  encryption::kBloomFilterHeader);
+  uint32_t header_cipher_len = static_cast<uint32_t>(header_cipher_total_len);
+  try {
+    deserializer.DeserializeMessage(header_cipher_buf->data(), &header_cipher_len,
+                                    &header, decryptor);
+  } catch (std::exception& e) {
+    std::stringstream ss;
+    ss << "Deserializing bloom filter header failed.\n" << e.what();
+    throw ParquetException(ss.str());
+  }
+  if (ARROW_PREDICT_FALSE(header_cipher_len != header_cipher_total_len)) {
+    std::stringstream ss;
+    ss << "Encrypted bloom filter header length mismatch: expected "
+       << header_cipher_total_len << " bytes, got " << header_cipher_len;
+    throw ParquetException(ss.str());
+  }
+  PARQUET_THROW_NOT_OK(ValidateBloomFilterHeader(header));
+
+  const int32_t bloom_filter_size = header.numBytes;
+  UpdateDecryptor(decryptor, row_group_ordinal, column_ordinal,
+                  encryption::kBloomFilterBitset);
+  const int32_t bitset_cipher_len = decryptor->CiphertextLength(bloom_filter_size);
+  const int64_t total_cipher_len =
+      header_cipher_total_len + static_cast<int64_t>(bitset_cipher_len);
+  if (bloom_filter_length && *bloom_filter_length != total_cipher_len) {
+    std::stringstream ss;
+    ss << "Bloom filter length (" << bloom_filter_length.value()
+       << ") does not match the actual bloom filter (size: " << total_cipher_len << ").";
+    throw ParquetException(ss.str());
+  }
+
+  // Read and decrypt the bitset bytes.
+  PARQUET_ASSIGN_OR_THROW(auto bitset_cipher_buf, input->Read(bitset_cipher_len));
+  if (ARROW_PREDICT_FALSE(bitset_cipher_buf->size() < bitset_cipher_len)) {
+    std::stringstream ss;
+    ss << "Bloom filter bitset read failed: expected " << bitset_cipher_len
+       << " bytes, got " << bitset_cipher_buf->size();
+    throw ParquetException(ss.str());
+  }
+
+  const int32_t bitset_plain_len =
+      decryptor->PlaintextLength(static_cast<int32_t>(bitset_cipher_len));
+  if (ARROW_PREDICT_FALSE(bitset_plain_len != bloom_filter_size)) {
+    throw ParquetException("Bloom filter bitset size does not match header");
+  }
+
+  auto bitset_plain_buf = AllocateBuffer(properties.memory_pool(), bitset_plain_len);
+  int32_t decrypted_len =
+      decryptor->Decrypt(bitset_cipher_buf->span_as<const uint8_t>(),
+                         bitset_plain_buf->mutable_span_as<uint8_t>());
+  if (ARROW_PREDICT_FALSE(decrypted_len != bitset_plain_len)) {
+    throw ParquetException("Bloom filter bitset decryption failed");
+  }
+
+  // Initialize the bloom filter from the decrypted bitset.
+  BlockSplitBloomFilter bloom_filter(properties.memory_pool());
+  bloom_filter.Init(bitset_plain_buf->data(), bloom_filter_size);
+  return bloom_filter;
+}
+
+}  // namespace
+
+BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
+    const ReaderProperties& properties, ArrowInputStream* input,
+    std::optional<int64_t> bloom_filter_length, Decryptor* decryptor,
+    int16_t row_group_ordinal, int16_t column_ordinal) {
+  if (decryptor == nullptr) {
+    throw ParquetException("Bloom filter decryptor must be provided");
+  }
+
+  // Read the full Bloom filter payload up front when the total length is known.
+  if (bloom_filter_length.has_value()) {
+    PARQUET_ASSIGN_OR_THROW(auto bloom_filter_buf, input->Read(*bloom_filter_length));
+    if (ARROW_PREDICT_FALSE(bloom_filter_buf->size() < *bloom_filter_length)) {
+      std::stringstream ss;
+      ss << "Bloom filter read failed: expected " << *bloom_filter_length
+         << " bytes, got " << bloom_filter_buf->size();
+      throw ParquetException(ss.str());
+    }
+    ::arrow::io::BufferReader reader(bloom_filter_buf);
+    return DeserializeEncryptedFromStream(properties, &reader, bloom_filter_length,
+                                          decryptor, row_group_ordinal, column_ordinal);
+  }
+
+  return DeserializeEncryptedFromStream(properties, input, bloom_filter_length, decryptor,
+                                        row_group_ordinal, column_ordinal);
+}
+
 BlockSplitBloomFilter BlockSplitBloomFilter::Deserialize(
     const ReaderProperties& properties, ArrowInputStream* input,
     std::optional<int64_t> bloom_filter_length) {
@@ -126,8 +282,7 @@ BlockSplitBloomFilter BlockSplitBloomFilter::Deserialize(
   // This gets used, then set by DeserializeThriftMsg
   uint32_t header_size = static_cast<uint32_t>(header_buf->size());
   try {
-    deserializer.DeserializeMessage(reinterpret_cast<const uint8_t*>(header_buf->data()),
-                                    &header_size, &header);
+    deserializer.DeserializeMessage(header_buf->data(), &header_size, &header);
     DCHECK_LE(header_size, header_buf->size());
   } catch (std::exception& e) {
     std::stringstream ss;

cpp/src/parquet/bloom_filter.h

@@ -23,6 +23,7 @@
 
 #include "arrow/util/bit_util.h"
 #include "arrow/util/logging.h"
+#include "parquet/encryption/type_fwd.h"
 #include "parquet/hasher.h"
 #include "parquet/platform.h"
 #include "parquet/types.h"
@@ -328,6 +329,24 @@ class PARQUET_EXPORT BlockSplitBloomFilter : public BloomFilter {
       const ReaderProperties& properties, ArrowInputStream* input_stream,
       std::optional<int64_t> bloom_filter_length = std::nullopt);
 
+  /// Deserialize an encrypted Bloom filter from an input stream.
+  ///
+  /// The same metadata decryptor is used for both the serialized header and bitset,
+  /// while switching module AADs between the two encrypted modules.
+  ///
+  /// @param properties The parquet reader properties.
+  /// @param input_stream The input stream from which to construct the bloom filter.
+  /// @param bloom_filter_length The length of the serialized bloom filter including
+  /// header.
+  /// @param decryptor Decryptor for encrypted Bloom filter modules.
+  /// @param row_group_ordinal Ordinal of the row group containing this Bloom filter.
+  /// @param column_ordinal Ordinal of the column containing this Bloom filter.
+  /// @return The BlockSplitBloomFilter.
+  static BlockSplitBloomFilter DeserializeEncrypted(
+      const ReaderProperties& properties, ArrowInputStream* input_stream,
+      std::optional<int64_t> bloom_filter_length, Decryptor* decryptor,
+      int16_t row_group_ordinal, int16_t column_ordinal);
+
  private:
   inline void InsertHashImpl(uint64_t hash);