refactor: strict validation from bytes bloom filter

filters/include/bloom_filter_impl.hpp

@@ -37,6 +37,19 @@
 
 namespace datasketches {
 
+static uint64_t validate_bloom_filter_num_bits_set(uint64_t num_bits_set,
+                                                   uint8_t* bit_array,
+                                                   uint64_t length_bytes,
+                                                   uint64_t dirty_bits_value)
+{
+  const uint64_t counted_bits_set = bit_array_ops::count_num_bits_set(bit_array, length_bytes);
+  if (num_bits_set != dirty_bits_value && num_bits_set != counted_bits_set) {
+    throw std::invalid_argument("Possible corruption: invalid number of bits set. Expected "
+      + std::to_string(counted_bits_set) + ", found " + std::to_string(num_bits_set));
+  }
+  return counted_bits_set;
+}
+
 template<typename A>
 bloom_filter_alloc<A>::bloom_filter_alloc(uint64_t num_bits, uint16_t num_hashes, uint64_t seed, const A& allocator) :
   allocator_(allocator),
@@ -143,9 +156,6 @@ bloom_filter_alloc<A>::bloom_filter_alloc(uint64_t seed,
 {
   // private constructor
   // no consistency checks since we should have done those prior to calling this
-  if (is_read_only_ && memory_ != nullptr && num_bits_set == DIRTY_BITS_VALUE) {
-    num_bits_set_ = bit_array_ops::count_num_bits_set(bit_array_, capacity_bits_ >> 3);
-  }
 }
 
 template<typename A>
@@ -298,8 +308,7 @@ bloom_filter_alloc<A> bloom_filter_alloc<A>::deserialize(std::istream& is, const
     return bloom_filter_alloc<A>(num_longs << 6, num_hashes, seed, allocator);
   }
 
-  const uint64_t num_bits_set = read<uint64_t>(is);
-  const bool is_dirty = (num_bits_set == DIRTY_BITS_VALUE);
+  const uint64_t raw_num_bits_set = read<uint64_t>(is);
 
   // allocate memory
   const uint64_t num_bytes = num_longs << 3;
@@ -310,8 +319,17 @@ bloom_filter_alloc<A> bloom_filter_alloc<A>::deserialize(std::istream& is, const
   }
   read(is, bit_array, num_bytes);
 
+  uint64_t num_bits_set;
+  try {
+    num_bits_set = validate_bloom_filter_num_bits_set(
+      raw_num_bits_set, bit_array, num_bytes, DIRTY_BITS_VALUE);
+  } catch (...) {
+    alloc.deallocate(bit_array, num_bytes);
+    throw;
+  }
+
   // pass to constructor
-  return bloom_filter_alloc<A>(seed, num_hashes, is_dirty, true, false, num_longs << 6, num_bits_set, bit_array, nullptr, allocator);
+  return bloom_filter_alloc<A>(seed, num_hashes, false, true, false, num_longs << 6, num_bits_set, bit_array, nullptr, allocator);
 }
 
 template<typename A>
@@ -374,9 +392,13 @@ bloom_filter_alloc<A> bloom_filter_alloc<A>::internal_deserialize_or_wrap(void*
     return bloom_filter_alloc<A>(num_longs << 6, num_hashes, seed, allocator);
   }
 
-  uint64_t num_bits_set;
-  ptr += copy_from_mem(ptr, num_bits_set);
-  const bool is_dirty = (num_bits_set == DIRTY_BITS_VALUE);
+  uint64_t raw_num_bits_set;
+  ptr += copy_from_mem(ptr, raw_num_bits_set);
+
+  const uint64_t num_bytes = num_longs << 3;
+  ensure_minimum_memory(end_ptr - ptr, num_bytes);
+  const uint64_t num_bits_set = validate_bloom_filter_num_bits_set(
+    raw_num_bits_set, const_cast<uint8_t*>(ptr), num_bytes, DIRTY_BITS_VALUE);
 
   uint8_t* bit_array;
   uint8_t* memory;
@@ -386,8 +408,6 @@ bloom_filter_alloc<A> bloom_filter_alloc<A>::internal_deserialize_or_wrap(void*
   } else {
     // allocate memory
     memory = nullptr;
-    const uint64_t num_bytes = num_longs << 3;
-    ensure_minimum_memory(end_ptr - ptr, num_bytes);
     AllocUint8 alloc(allocator);
     bit_array = alloc.allocate(num_bytes);
     if (bit_array == nullptr) {
@@ -397,7 +417,7 @@ bloom_filter_alloc<A> bloom_filter_alloc<A>::internal_deserialize_or_wrap(void*
   }
 
   // pass to constructor -- !wrap == is_owned_
-  return bloom_filter_alloc<A>(seed, num_hashes, is_dirty, !wrap, read_only, num_longs << 6, num_bits_set, bit_array, memory, allocator);
+  return bloom_filter_alloc<A>(seed, num_hashes, false, !wrap, read_only, num_longs << 6, num_bits_set, bit_array, memory, allocator);
 }
 
 template<typename A>

filters/test/bloom_filter_test.cpp

@@ -19,6 +19,9 @@
 
 #include <catch2/catch.hpp>
 
+#include <cstring>
+#include <limits>
+
 #include "bloom_filter.hpp"
 
 #ifdef TEST_BINARY_INPUT_PATH
@@ -403,4 +406,28 @@ TEST_CASE("bloom_filter: non-empty serialization", "[bloom_filter]") {
   REQUIRE(bf_wrap.query(-1.0));
 }
 
+TEST_CASE("bloom_filter: inconsistent num bits set is rejected", "[bloom_filter]") {
+  const size_t num_bits_set_offset = 24;
+
+  auto bf = bloom_filter::builder::create_by_accuracy(100, 0.01);
+  bf.update("apple");
+  bf.update("banana");
+  const uint64_t actual_bits_set = bf.get_bits_used();
+  REQUIRE(actual_bits_set > 1);
+
+  const uint64_t invalid_counts[] = {0, actual_bits_set - 1, actual_bits_set + 1};
+  for (const uint64_t serialized_count: invalid_counts) {
+    auto bytes = bf.serialize();
+    std::memcpy(bytes.data() + num_bits_set_offset, &serialized_count, sizeof(serialized_count));
+
+    REQUIRE_THROWS_AS(bloom_filter::deserialize(bytes.data(), bytes.size()), std::invalid_argument);
+    REQUIRE_THROWS_AS(bloom_filter::wrap(bytes.data(), bytes.size()), std::invalid_argument);
+    REQUIRE_THROWS_AS(bloom_filter::writable_wrap(bytes.data(), bytes.size()), std::invalid_argument);
+
+    const std::string serialized(reinterpret_cast<const char*>(bytes.data()), bytes.size());
+    std::istringstream is(serialized, std::ios::in | std::ios::binary);
+    REQUIRE_THROWS_AS(bloom_filter::deserialize(is), std::invalid_argument);
+  }
+}
+
 } // namespace datasketches