Check the buffer size before reading compact theta sketch preamble fields

theta/include/compact_theta_sketch_parser_impl.hpp

@@ -80,8 +80,9 @@ auto compact_theta_sketch_parser<dummy>::parse(const void* ptr, size_t size, uin
         check_memory_size(ptr, size, 16, dump_on_error);
         return {false, true, seed_hash, 1, theta, reinterpret_cast<const uint64_t*>(ptr) + COMPACT_SKETCH_SINGLE_ENTRY_U64, 64};
       }
-      const uint32_t num_entries = reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32];
       const size_t entries_start_u64 = has_theta ? COMPACT_SKETCH_ENTRIES_ESTIMATION_U64 : COMPACT_SKETCH_ENTRIES_EXACT_U64;
+      check_memory_size(ptr, size, entries_start_u64 * sizeof(uint64_t), dump_on_error);
+      const uint32_t num_entries = reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32];
       const uint64_t* entries = reinterpret_cast<const uint64_t*>(ptr) + entries_start_u64;
       const size_t expected_size_bytes = (entries_start_u64 + num_entries) * sizeof(uint64_t);
       check_memory_size(ptr, size, expected_size_bytes, dump_on_error);
@@ -90,6 +91,7 @@ auto compact_theta_sketch_parser<dummy>::parse(const void* ptr, size_t size, uin
   }
   case 1:  {
       uint16_t seed_hash = compute_seed_hash(seed);
+      check_memory_size(ptr, size, COMPACT_SKETCH_ENTRIES_ESTIMATION_U64 * sizeof(uint64_t), dump_on_error);
       const uint32_t num_entries = reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32];
       uint64_t theta = reinterpret_cast<const uint64_t*>(ptr)[COMPACT_SKETCH_THETA_U64];
       bool is_empty = (num_entries == 0) && (theta == theta_constants::MAX_THETA);
@@ -106,6 +108,7 @@ auto compact_theta_sketch_parser<dummy>::parse(const void* ptr, size_t size, uin
       if (preamble_size == 1) {
           return {true, true, seed_hash, 0, theta_constants::MAX_THETA, nullptr, 64};
       } else if (preamble_size == 2) {
+          check_memory_size(ptr, size, COMPACT_SKETCH_ENTRIES_EXACT_U64 * sizeof(uint64_t), dump_on_error);
           const uint32_t num_entries = reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32];
           if (num_entries == 0) {
               return {true, true, seed_hash, 0, theta_constants::MAX_THETA, nullptr, 64};
@@ -116,6 +119,7 @@ auto compact_theta_sketch_parser<dummy>::parse(const void* ptr, size_t size, uin
               return {false, true, seed_hash, num_entries, theta_constants::MAX_THETA, entries, 64};
           }
       } else if (preamble_size == 3) {
+          check_memory_size(ptr, size, COMPACT_SKETCH_ENTRIES_ESTIMATION_U64 * sizeof(uint64_t), dump_on_error);
           const uint32_t num_entries = reinterpret_cast<const uint32_t*>(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32];
           uint64_t theta = reinterpret_cast<const uint64_t*>(ptr)[COMPACT_SKETCH_THETA_U64];
           bool is_empty = (num_entries == 0) && (theta == theta_constants::MAX_THETA);

theta/test/theta_sketch_test.cpp

@@ -570,6 +570,50 @@ TEST_CASE("theta sketch: deserialize estimation mode buffer overrun", "[theta_sk
   REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(bytes.data(), bytes.size() - 1), std::out_of_range);
 }
 
+// The header fields of serial versions 1 to 3 that follow the first 8 bytes must not be read
+// before the size of the buffer is checked. Each buffer below holds exactly the preamble bytes,
+// so that reading beyond them would be an out-of-bounds read, and is followed by the bytes
+// that would make the sketch look empty if they were read.
+TEST_CASE("theta sketch: deserialize truncated preamble", "[theta_sketch]") {
+  const uint16_t seed_hash = compute_seed_hash(DEFAULT_SEED);
+  const uint8_t seed_hash_lo = seed_hash & 0xff;
+  const uint8_t seed_hash_hi = seed_hash >> 8;
+  const uint8_t max_theta[8] = {0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f};
+
+  // serial version 1: three preamble longs, num_entries = 0 and theta = max would be empty
+  std::vector<uint8_t> v1 = {3, 1, 3, 0, 0, 0x1a, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  v1.insert(v1.end(), max_theta, max_theta + 8);
+  for (size_t size: {8, 16}) {
+    const std::vector<uint8_t> bytes(v1.begin(), v1.begin() + size);
+    REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(bytes.data(), bytes.size()), std::out_of_range);
+    REQUIRE_THROWS_AS(wrapped_compact_theta_sketch::wrap(bytes.data(), bytes.size()), std::out_of_range);
+  }
+
+  // serial version 2: two preamble longs, num_entries = 0 would be empty
+  std::vector<uint8_t> v2_exact = {2, 2, 3, 0, 0, 0x1a, seed_hash_lo, seed_hash_hi, 0, 0, 0, 0, 0, 0, 0, 0};
+  {
+    const std::vector<uint8_t> bytes(v2_exact.begin(), v2_exact.begin() + 8);
+    REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(bytes.data(), bytes.size()), std::out_of_range);
+    REQUIRE_THROWS_AS(wrapped_compact_theta_sketch::wrap(bytes.data(), bytes.size()), std::out_of_range);
+  }
+  REQUIRE(compact_theta_sketch::deserialize(v2_exact.data(), v2_exact.size()).is_empty());
+
+  // serial version 2: three preamble longs, num_entries = 0 and theta = max would be empty
+  std::vector<uint8_t> v2_estimation = {3, 2, 3, 0, 0, 0x1a, seed_hash_lo, seed_hash_hi, 0, 0, 0, 0, 0, 0, 0, 0};
+  v2_estimation.insert(v2_estimation.end(), max_theta, max_theta + 8);
+  for (size_t size: {8, 16}) {
+    const std::vector<uint8_t> bytes(v2_estimation.begin(), v2_estimation.begin() + size);
+    REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(bytes.data(), bytes.size()), std::out_of_range);
+    REQUIRE_THROWS_AS(wrapped_compact_theta_sketch::wrap(bytes.data(), bytes.size()), std::out_of_range);
+  }
+  REQUIRE(compact_theta_sketch::deserialize(v2_estimation.data(), v2_estimation.size()).is_empty());
+
+  // serial version 3: two preamble longs, not empty
+  const std::vector<uint8_t> v3 = {2, 3, 3, 0, 0, 0x1a, seed_hash_lo, seed_hash_hi};
+  REQUIRE_THROWS_AS(compact_theta_sketch::deserialize(v3.data(), v3.size()), std::out_of_range);
+  REQUIRE_THROWS_AS(wrapped_compact_theta_sketch::wrap(v3.data(), v3.size()), std::out_of_range);
+}
+
 TEST_CASE("theta sketch: conversion constructor and wrapped compact", "[theta_sketch]") {
   update_theta_sketch update_sketch = update_theta_sketch::builder().build();
   const int n = 8192;