Theta: validate the preamble of empty compact sketches as Java does

theta/include/compact_theta_sketch_parser.hpp

@@ -41,6 +41,31 @@ class compact_theta_sketch_parser {
   static void check_v4_entry_bits(uint8_t entry_bits);
   static void check_v4_num_entries_bytes(uint8_t num_entries_bytes);
 
+  /**
+   * Checks the first 8 bytes of a serial version 3 compact sketch that has the empty flag set.
+   * Accepts the current 8-byte form (preamble longs 1), and the legacy 24-byte form
+   * (preamble longs 3) written by Java before 1.0.0 for sketches with p < 1.
+   * The legacy form must also carry the expected seed hash.
+   * The caller must check that a legacy image is 24 bytes long and has no entries.
+   * @param pre0 the first 8 bytes of the image, little-endian
+   * @param expected_seed_hash the seed hash computed from the expected seed
+   * @return the number of preamble longs: 1 or 3
+   * @throw std::invalid_argument if the image is not a valid empty compact sketch
+   */
+  static uint8_t check_empty_v3(uint64_t pre0, uint16_t expected_seed_hash);
+
+  // The MASK selects which bits of the first 8 bytes of an empty image are examined.
+  // The TEST gives the required values of the examined bits; it must lie within the MASK.
+  // Examined: preamble longs, serial version 3, sketch type 3, bytes 3 and 4 = 0 (not used by compact);
+  //  flags read-only, empty, compact set; single-item and reserved bits 0, 6, 7 clear.
+  // The ordered flag is ignored: C++ before 3.3.0 and Java before 1.0.0 wrote empty images without it.
+  // The reserved flag bits must be zero for serial version 3. Any future use of them requires a new one.
+  // The seed hash is ignored in the 8-byte form: it may be 0 or the seed hash, depending on the writer.
+  // These must stay the same as in Java EmptyCompactSketch.
+  static const uint64_t EMPTY_SKETCH_MASK = 0x0000EFFFFFFFFFFFULL;
+  static const uint64_t EMPTY_SKETCH_TEST = 0x00000E0000030301ULL;
+  static const uint64_t EMPTY_SKETCH_TEST_LEGACY = 0x00000E0000030303ULL;
+
 private:
   // offsets are in sizeof(type)
   static const size_t COMPACT_SKETCH_PRE_LONGS_BYTE = 0;

theta/include/compact_theta_sketch_parser_impl.hpp

@@ -70,6 +70,12 @@ auto compact_theta_sketch_parser<dummy>::parse(const void* ptr, size_t size, uin
       uint64_t theta = theta_constants::MAX_THETA;
       const uint16_t seed_hash = reinterpret_cast<const uint16_t*>(ptr)[COMPACT_SKETCH_SEED_HASH_U16];
       if (reinterpret_cast<const uint8_t*>(ptr)[COMPACT_SKETCH_FLAGS_BYTE] & (1 << COMPACT_SKETCH_IS_EMPTY_FLAG)) {
+        if (check_empty_v3(reinterpret_cast<const uint64_t*>(ptr)[0], compute_seed_hash(seed)) == 3) {
+          check_memory_size(ptr, size, COMPACT_SKETCH_ENTRIES_ESTIMATION_U64 * sizeof(uint64_t), dump_on_error);
+          if (reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32] != 0) {
+            throw std::invalid_argument("legacy empty compact sketch has entries");
+          }
+        }
         return {true, true, seed_hash, 0, theta, nullptr, 64};
       }
       checker<true>::check_seed_hash(seed_hash, compute_seed_hash(seed));
@@ -146,6 +152,26 @@ void compact_theta_sketch_parser<dummy>::check_memory_size(const void* ptr, size
       + (dump_on_error ? (", sketch dump: " + hex_dump(reinterpret_cast<const uint8_t*>(ptr), actual_bytes)) : ""));
 }
 
+template<bool dummy>
+const uint64_t compact_theta_sketch_parser<dummy>::EMPTY_SKETCH_MASK;
+template<bool dummy>
+const uint64_t compact_theta_sketch_parser<dummy>::EMPTY_SKETCH_TEST;
+template<bool dummy>
+const uint64_t compact_theta_sketch_parser<dummy>::EMPTY_SKETCH_TEST_LEGACY;
+
+template<bool dummy>
+uint8_t compact_theta_sketch_parser<dummy>::check_empty_v3(uint64_t pre0, uint16_t expected_seed_hash) {
+  const uint64_t masked = pre0 & EMPTY_SKETCH_MASK;
+  if (masked == EMPTY_SKETCH_TEST) return 1;
+  if (masked == EMPTY_SKETCH_TEST_LEGACY) {
+    checker<true>::check_seed_hash(static_cast<uint16_t>(pre0 >> 48), expected_seed_hash);
+    return 3;
+  }
+  std::stringstream s;
+  s << "empty compact sketch preamble does not match: 0x" << std::hex << masked;
+  throw std::invalid_argument(s.str());
+}
+
 template<bool dummy>
 void compact_theta_sketch_parser<dummy>::check_v4_entry_bits(uint8_t entry_bits) {
   // deltas between ordered hashes below 2^63 need 1 to 63 bits

theta/include/theta_sketch_impl.hpp

@@ -668,11 +668,26 @@ template<typename A>
 compact_theta_sketch_alloc<A> compact_theta_sketch_alloc<A>::deserialize_v3(
     uint8_t preamble_longs, std::istream& is, uint64_t seed, const A& allocator)
 {
-  read<uint16_t>(is); // unused
+  const auto unused16 = read<uint16_t>(is);
   const auto flags_byte = read<uint8_t>(is);
   const auto seed_hash = read<uint16_t>(is);
   const bool is_empty = flags_byte & (1 << flags::IS_EMPTY);
-  if (!is_empty) checker<true>::check_seed_hash(seed_hash, compute_seed_hash(seed));
+  if (is_empty) {
+    const uint64_t pre0 = static_cast<uint64_t>(preamble_longs)
+        | (static_cast<uint64_t>(UNCOMPRESSED_SERIAL_VERSION) << 8)
+        | (static_cast<uint64_t>(SKETCH_TYPE) << 16)
+        | (static_cast<uint64_t>(unused16) << 24)
+        | (static_cast<uint64_t>(flags_byte) << 40)
+        | (static_cast<uint64_t>(seed_hash) << 48);
+    if (compact_theta_sketch_parser<true>::check_empty_v3(pre0, compute_seed_hash(seed)) == 3) {
+      const auto num_entries = read<uint32_t>(is);
+      read<uint32_t>(is); // unused
+      read<uint64_t>(is); // theta is ignored, empty implies theta = 1.0
+      if (num_entries != 0) throw std::invalid_argument("legacy empty compact sketch has entries");
+    }
+  } else {
+    checker<true>::check_seed_hash(seed_hash, compute_seed_hash(seed));
+  }
   uint64_t theta = theta_constants::MAX_THETA;
   uint32_t num_entries = 0;
   if (!is_empty) {

theta/test/theta_sketch_test.cpp

@@ -940,4 +940,66 @@ TEST_CASE("theta sketch: deserialize v4 entry bits overflow", "[theta_sketch]")
   REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(bytes.data(), bytes.size()), std::out_of_range);
 }
 
+// Empty images: same rules as Java EmptyCompactSketch
+
+static void check_empty_accepted(const std::vector<uint8_t>& bytes) {
+  std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
+  s.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
+  REQUIRE(compact_theta_sketch::deserialize(s).is_empty());
+  REQUIRE(s.tellg() == static_cast<std::streampos>(bytes.size())); // whole image consumed
+  REQUIRE(compact_theta_sketch::deserialize(bytes.data(), bytes.size()).is_empty());
+  REQUIRE(wrapped_compact_theta_sketch::wrap(bytes.data(), bytes.size()).is_empty());
+}
+
+static void check_empty_rejected(const std::vector<uint8_t>& bytes) {
+  std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
+  s.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
+  REQUIRE_THROWS(compact_theta_sketch::deserialize(s));
+  REQUIRE_THROWS(compact_theta_sketch::deserialize(bytes.data(), bytes.size()));
+  REQUIRE_THROWS(wrapped_compact_theta_sketch::wrap(bytes.data(), bytes.size()));
+}
+
+// sketches-core 0.9.0, p = 0.5: full 24-byte form, count, p = 1.0f, theta = 0.5
+static std::vector<uint8_t> legacy_empty_p05(uint8_t flags, uint8_t count) {
+  return {3, 3, 3, 0, 0, flags, 0xCC, 0x93,
+          count, 0, 0, 0, 0, 0, 0x80, 0x3F,
+          0, 0, 0, 0, 0, 0, 0, 0x40};
+}
+
+TEST_CASE("theta sketch: empty mask and tests", "[theta_sketch]") {
+  using parser = compact_theta_sketch_parser<true>;
+  REQUIRE((parser::EMPTY_SKETCH_TEST & ~parser::EMPTY_SKETCH_MASK) == 0);
+  REQUIRE((parser::EMPTY_SKETCH_TEST_LEGACY & ~parser::EMPTY_SKETCH_MASK) == 0);
+  // our own empty image must pass
+  auto bytes = update_theta_sketch::builder().build().compact().serialize();
+  uint64_t pre0;
+  std::memcpy(&pre0, bytes.data(), sizeof(pre0));
+  REQUIRE((pre0 & parser::EMPTY_SKETCH_MASK) == parser::EMPTY_SKETCH_TEST);
+}
+
+TEST_CASE("theta sketch: deserialize historical empty images", "[theta_sketch]") {
+  check_empty_accepted({1, 3, 3, 0, 0, 0x1E, 0, 0});       // Java and current C++
+  check_empty_accepted({1, 3, 3, 0, 0, 0x1E, 0xCC, 0x93}); // C++ 1.0.0 to 5.2.0, ordered
+  check_empty_accepted({1, 3, 3, 0, 0, 0x0E, 0xCC, 0x93}); // C++ before 3.3.0, unordered
+  check_empty_accepted(legacy_empty_p05(0x1E, 0));         // Java before 1.0.0, ordered
+  check_empty_accepted(legacy_empty_p05(0x0E, 0));         // Java before 1.0.0, unordered
+}
+
+TEST_CASE("theta sketch: deserialize mangled empty images", "[theta_sketch]") {
+  check_empty_rejected({1, 3, 3, 1, 0, 0x1E, 0, 0}); // byte 3 not zero
+  check_empty_rejected({1, 3, 3, 0, 1, 0x1E, 0, 0}); // byte 4 not zero
+  check_empty_rejected({1, 3, 3, 0, 0, 0x1F, 0, 0}); // reserved bit 0 set
+  check_empty_rejected({1, 3, 3, 0, 0, 0x5E, 0, 0}); // reserved bit 6 set
+  check_empty_rejected({1, 3, 3, 0, 0, 0x9E, 0, 0}); // reserved bit 7 set
+  check_empty_rejected({1, 3, 3, 0, 0, 0x3E, 0, 0}); // single-item set
+  check_empty_rejected({1, 3, 3, 0, 0, 0x1C, 0, 0}); // read-only not set
+  check_empty_rejected({1, 3, 3, 0, 0, 0x16, 0, 0}); // compact not set
+  check_empty_rejected({2, 3, 3, 0, 0, 0x1E, 0, 0}); // preamble longs 2
+  check_empty_rejected(legacy_empty_p05(0x1E, 1));   // preamble longs 3, count not 0
+  check_empty_rejected({3, 3, 3, 0, 0, 0x1E, 0, 0}); // preamble longs 3, too short
+  auto wrong_seed = legacy_empty_p05(0x1E, 0);
+  wrong_seed[6] = 0; wrong_seed[7] = 0;
+  check_empty_rejected(wrong_seed);                  // preamble longs 3, seed hash mismatch
+}
+
 } /* namespace datasketches */