Check KLL, REQ and Quantiles images when deserializing

NOTICE

@@ -1,9 +1,9 @@
 Apache DataSketches C++
-Copyright 2025-2026 The Apache Software Foundation
+Copyright 2019-2026 The Apache Software Foundation
 
 Copyright 2015-2018 Yahoo Inc.
 Copyright 2019-2020 Verizon Media
-Copyright 2021- Yahoo Inc.
+Copyright 2021-2025 Yahoo Inc.
 
 This product includes software developed at
 The Apache Software Foundation (http://www.apache.org/).

kll/include/kll_sketch.hpp

@@ -591,10 +591,13 @@ class kll_sketch {
     uint32_t safe_level_size(uint8_t level) const;
     uint32_t get_num_retained_above_level_zero() const;
 
+    static void check_k(uint16_t k);
     static void check_m(uint8_t m);
     static void check_preamble_ints(uint8_t preamble_ints, uint8_t flags_byte);
     static void check_serial_version(uint8_t serial_version);
     static void check_family_id(uint8_t family_id);
+    static void check_num_levels(uint8_t num_levels);
+    static void check_levels(const vector_u32& levels);
 
     void check_sorting() const;
 

kll/include/kll_sketch_impl.hpp

@@ -450,6 +450,7 @@ kll_sketch<T, C, A> kll_sketch<T, C, A>::deserialize(std::istream& is, const Ser
   const auto m = read<uint8_t>(is);
   read<uint8_t>(is); // skip unused byte
 
+  check_k(k);
   check_m(m);
   check_preamble_ints(preamble_ints, flags_byte);
   check_serial_version(serial_version);
@@ -472,6 +473,7 @@ kll_sketch<T, C, A> kll_sketch<T, C, A>::deserialize(std::istream& is, const Ser
     min_k = read<uint16_t>(is);
     num_levels = read<uint8_t>(is);
     read<uint8_t>(is); // skip unused byte
+    check_num_levels(num_levels);
   }
   vector_u32 levels(num_levels + 1, 0, allocator);
   const uint32_t capacity(kll_helper::compute_total_capacity(k, m, num_levels));
@@ -480,8 +482,10 @@ kll_sketch<T, C, A> kll_sketch<T, C, A>::deserialize(std::istream& is, const Ser
   } else {
     // the last integer in levels_ is not serialized because it can be derived
     read(is, levels.data(), sizeof(levels[0]) * num_levels);
+    if (!is.good()) throw std::runtime_error("error reading from std::istream");
   }
   levels[num_levels] = capacity;
+  check_levels(levels);
   optional<T> min_item;
   optional<T> max_item;
   if (!is_single_item) {
@@ -537,6 +541,7 @@ kll_sketch<T, C, A> kll_sketch<T, C, A>::deserialize(const void* bytes, size_t s
   ptr += copy_from_mem(ptr, m);
   ptr += sizeof(uint8_t); // skip unused byte
 
+  check_k(k);
   check_m(m);
   check_preamble_ints(preamble_ints, flags_byte);
   check_serial_version(serial_version);
@@ -560,16 +565,19 @@ kll_sketch<T, C, A> kll_sketch<T, C, A>::deserialize(const void* bytes, size_t s
     ptr += copy_from_mem(ptr, min_k);
     ptr += copy_from_mem(ptr, num_levels);
     ptr += sizeof(uint8_t); // skip unused byte
+    check_num_levels(num_levels);
   }
   vector_u32 levels(num_levels + 1, 0, allocator);
   const uint32_t capacity(kll_helper::compute_total_capacity(k, m, num_levels));
   if (is_single_item) {
     levels[0] = capacity - 1;
   } else {
     // the last integer in levels_ is not serialized because it can be derived
+    ensure_minimum_memory(size, (ptr - static_cast<const char*>(bytes)) + sizeof(levels[0]) * num_levels);
     ptr += copy_from_mem(ptr, levels.data(), sizeof(levels[0]) * num_levels);
   }
   levels[num_levels] = capacity;
+  check_levels(levels);
   optional<T> min_item;
   optional<T> max_item;
   if (!is_single_item) {
@@ -867,6 +875,14 @@ uint32_t kll_sketch<T, C, A>::get_num_retained_above_level_zero() const {
   return levels_[num_levels_] - levels_[1];
 }
 
+template<typename T, typename C, typename A>
+void kll_sketch<T, C, A>::check_k(uint16_t k) {
+  if (k < kll_constants::MIN_K) {
+    throw std::invalid_argument("Possible corruption: K must be >= " + std::to_string(kll_constants::MIN_K)
+        + ": " + std::to_string(k));
+  }
+}
+
 template<typename T, typename C, typename A>
 void kll_sketch<T, C, A>::check_m(uint8_t m) {
   if (m != kll_constants::DEFAULT_M) {
@@ -909,6 +925,25 @@ void kll_sketch<T, C, A>::check_family_id(uint8_t family_id) {
   }
 }
 
+template<typename T, typename C, typename A>
+void kll_sketch<T, C, A>::check_num_levels(uint8_t num_levels) {
+  if (num_levels == 0) {
+    throw std::invalid_argument("Possible corruption: number of levels must be > 0");
+  }
+}
+
+// The last entry is the capacity. Each level must start at or after the previous one,
+// so every level lies within the items buffer.
+template<typename T, typename C, typename A>
+void kll_sketch<T, C, A>::check_levels(const vector_u32& levels) {
+  for (size_t i = 1; i < levels.size(); ++i) {
+    if (levels[i - 1] > levels[i]) {
+      throw std::invalid_argument("Possible corruption: levels[" + std::to_string(i - 1) + "] = "
+          + std::to_string(levels[i - 1]) + " > levels[" + std::to_string(i) + "] = " + std::to_string(levels[i]));
+    }
+  }
+}
+
 template <typename T, typename C, typename A>
 string<A> kll_sketch<T, C, A>::to_string(bool print_levels, bool print_items) const {
   // Using a temporary stream for implementation here does not comply with AllocatorAwareContainer requirements.

kll/test/kll_sketch_test.cpp

@@ -835,4 +835,74 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
   REQUIRE(test_allocator_total_bytes == 0);
 }
 
+
+TEST_CASE("kll sketch: deserialize rejects inconsistent levels", "[kll_sketch]") {
+  using vector_bytes = kll_sketch<int>::vector_bytes;
+  const size_t k_offset = 4;
+  const size_t num_levels_offset = 18;
+  const size_t levels_offset = 20;
+  auto put_u32 = [](vector_bytes& bytes, size_t offset, uint32_t value) {
+    std::memcpy(bytes.data() + offset, &value, sizeof(value));
+  };
+  auto get_u32 = [](const vector_bytes& bytes, size_t offset) {
+    uint32_t value;
+    std::memcpy(&value, bytes.data() + offset, sizeof(value));
+    return value;
+  };
+  auto check_rejected = [](const vector_bytes& bytes) {
+    REQUIRE_THROWS_AS(kll_sketch<int>::deserialize(bytes.data(), bytes.size()), std::invalid_argument);
+    std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
+    s.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
+    REQUIRE_THROWS_AS(kll_sketch<int>::deserialize(s), std::invalid_argument);
+  };
+
+  kll_sketch<int> one_level(8);
+  for (int i = 0; i < 3; ++i) one_level.update(i);
+  kll_sketch<int> three_levels(8);
+  for (int i = 0; i < 40; ++i) three_levels.update(i);
+  REQUIRE(three_levels.serialize()[num_levels_offset] == 3);
+
+  SECTION("first level beyond capacity") {
+    auto bytes = one_level.serialize();
+    put_u32(bytes, levels_offset, 1009);
+    check_rejected(bytes);
+  }
+
+  SECTION("level beyond capacity") {
+    auto bytes = three_levels.serialize();
+    put_u32(bytes, levels_offset + sizeof(uint32_t), 100000);
+    check_rejected(bytes);
+  }
+
+  SECTION("decreasing levels") {
+    auto bytes = three_levels.serialize();
+    REQUIRE(get_u32(bytes, levels_offset) > 0);
+    put_u32(bytes, levels_offset + sizeof(uint32_t), 0);
+    check_rejected(bytes);
+  }
+
+  SECTION("zero levels") {
+    auto bytes = three_levels.serialize();
+    bytes[num_levels_offset] = 0;
+    check_rejected(bytes);
+  }
+
+  SECTION("k below minimum") {
+    auto bytes = three_levels.serialize();
+    const uint16_t k = 4;
+    std::memcpy(bytes.data() + k_offset, &k, sizeof(k));
+    check_rejected(bytes);
+  }
+
+  SECTION("levels beyond the input") {
+    auto bytes = three_levels.serialize();
+    bytes[num_levels_offset] = 40;
+    bytes.resize(levels_offset);
+    REQUIRE_THROWS_AS(kll_sketch<int>::deserialize(bytes.data(), bytes.size()), std::out_of_range);
+    std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
+    s.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
+    REQUIRE_THROWS_AS(kll_sketch<int>::deserialize(s), std::runtime_error);
+  }
+}
+
 } /* namespace datasketches */

quantiles/include/quantiles_sketch_impl.hpp

@@ -531,6 +531,7 @@ auto quantiles_sketch<T, C, A>::deserialize(const void* bytes, size_t size, cons
 
   if (serial_version == 1) {
     uint64_t unused_long;
+    ensure_minimum_memory(end_ptr - ptr, sizeof(unused_long));
     ptr += copy_from_mem(ptr, unused_long); // no longer used
   }
 

quantiles/test/quantiles_sketch_test.cpp

@@ -956,4 +956,17 @@ TEST_CASE("quantiles sketch", "[quantiles_sketch]") {
   }
 }
 
+
+TEST_CASE("quantiles sketch: deserialize serial version 1 truncated after min and max", "[quantiles_sketch]") {
+  quantiles_sketch<int> sketch(16);
+  for (int i = 0; i < 1000; ++i) sketch.update(i);
+  auto bytes = sketch.serialize();
+  // serial version 1 has an unused 8-byte field after min and max
+  bytes[0] = 5; // preamble longs
+  bytes[1] = 1; // serial version
+  bytes[3] &= ~(1 << 3); // not compact
+  bytes.resize(8 + sizeof(uint64_t) + 2 * sizeof(int)); // preamble, n, min and max
+  REQUIRE_THROWS_AS(quantiles_sketch<int>::deserialize(bytes.data(), bytes.size()), std::out_of_range);
+}
+
 } /* namespace datasketches */

req/include/req_compactor.hpp

@@ -87,11 +87,11 @@ class req_compactor {
 
   template<typename S>
   static req_compactor deserialize(std::istream& is, const S& serde, const Comparator& comparator,
-      const Allocator& allocator, bool sorted, bool hra);
+      const Allocator& allocator, bool sorted, bool hra, uint16_t k, uint8_t level);
 
   template<typename S>
   static std::pair<req_compactor, size_t> deserialize(const void* bytes, size_t size, const S& serde,
-      const Comparator& comparator, const Allocator& allocator, bool sorted, bool hra);
+      const Comparator& comparator, const Allocator& allocator, bool sorted, bool hra, uint16_t k, uint8_t level);
 
   template<typename S>
   static req_compactor deserialize(std::istream& is, const S& serde, const Comparator& comparator,
@@ -122,6 +122,7 @@ class req_compactor {
   void ensure_space(uint32_t num);
 
   static uint32_t nearest_even(float value);
+  static void check_header(uint16_t k, uint8_t level, uint8_t lg_weight, float section_size_raw, uint8_t num_sections);
 
   template<typename InIter, typename OutIter>
   static void promote_evens_or_odds(InIter from, InIter to, bool flag, OutIter dst);

req/include/req_compactor_impl.hpp

@@ -326,6 +326,30 @@ uint32_t req_compactor<T, C, A>::nearest_even(float value) {
   return static_cast<uint32_t>(round(value / 2)) << 1;
 }
 
+template<typename T, typename C, typename A>
+void req_compactor<T, C, A>::check_header(uint16_t k, uint8_t level, uint8_t lg_weight, float section_size_raw,
+    uint8_t num_sections) {
+  // the sketch creates the compactor at level h with lg_weight h
+  if (lg_weight != level) {
+    throw std::invalid_argument("Possible corruption: compactor " + std::to_string(level) + " has lg_weight "
+        + std::to_string(lg_weight));
+  }
+  // the section size starts at k and shrinks to no less than MIN_K
+  // (the first check is false for NaN and keeps nearest_even() in range)
+  const uint32_t section_size = (section_size_raw > 0 && section_size_raw <= k + 1) ? nearest_even(section_size_raw) : 0;
+  if (section_size < req_constants::MIN_K || section_size > k) {
+    throw std::invalid_argument("Possible corruption: section size " + std::to_string(section_size_raw)
+        + " is out of range for k " + std::to_string(k));
+  }
+  // the number of sections starts at 3 and doubles, and must keep 1 << (num_sections - 1) defined
+  // the limit uses the raw section size: when that has shrunk to 3, the even section size is 4 and can exceed it
+  const uint32_t rounded_section_size = static_cast<uint32_t>(round(section_size_raw));
+  if (num_sections < 1 || num_sections > 64 || rounded_section_size * num_sections > static_cast<uint32_t>(k) * k) {
+    throw std::invalid_argument("Possible corruption: number of sections " + std::to_string(num_sections)
+        + " is out of range for section size " + std::to_string(section_size_raw) + " and k " + std::to_string(k));
+  }
+}
+
 template<typename T, typename C, typename A>
 template<typename InIter, typename OutIter>
 void req_compactor<T, C, A>::promote_evens_or_odds(InIter from, InIter to, bool odds, OutIter dst) {
@@ -394,13 +418,15 @@ size_t req_compactor<T, C, A>::serialize(void* dst, size_t capacity, const S& se
 template<typename T, typename C, typename A>
 template<typename S>
 req_compactor<T, C, A> req_compactor<T, C, A>::deserialize(std::istream& is, const S& serde,
-    const C& comparator, const A& allocator, bool sorted, bool hra) {
+    const C& comparator, const A& allocator, bool sorted, bool hra, uint16_t k, uint8_t level) {
   auto state = read<decltype(state_)>(is);
   auto section_size_raw = read<decltype(section_size_raw_)>(is);
   auto lg_weight = read<decltype(lg_weight_)>(is);
   auto num_sections = read<decltype(num_sections_)>(is);
   read<uint16_t>(is); // padding
   auto num_items = read<uint32_t>(is);
+  if (!is.good()) throw std::runtime_error("error reading from std::istream");
+  check_header(k, level, lg_weight, section_size_raw, num_sections);
   auto items = deserialize_items(is, serde, allocator, num_items);
   return req_compactor(hra, lg_weight, sorted, section_size_raw, num_sections, state, std::move(items), num_items,
       comparator, allocator);
@@ -431,8 +457,9 @@ auto req_compactor<T, C, A>::deserialize_items(std::istream& is, const S& serde,
 template<typename T, typename C, typename A>
 template<typename S>
 std::pair<req_compactor<T, C, A>, size_t> req_compactor<T, C, A>::deserialize(const void* bytes, size_t size,
-    const S& serde, const C& comparator, const A& allocator, bool sorted, bool hra) {
-  ensure_minimum_memory(size, 8);
+    const S& serde, const C& comparator, const A& allocator, bool sorted, bool hra, uint16_t k, uint8_t level) {
+  ensure_minimum_memory(size, sizeof(state_) + sizeof(section_size_raw_) + sizeof(lg_weight_) + sizeof(num_sections_)
+      + sizeof(uint16_t) + sizeof(num_items_));
   const char* ptr = static_cast<const char*>(bytes);
   const char* end_ptr = static_cast<const char*>(bytes) + size;
 
@@ -447,6 +474,7 @@ std::pair<req_compactor<T, C, A>, size_t> req_compactor<T, C, A>::deserialize(co
   ptr += 2; // padding
   uint32_t num_items;
   ptr += copy_from_mem(ptr, num_items);
+  check_header(k, level, lg_weight, section_size_raw, num_sections);
   auto pair = deserialize_items(ptr, end_ptr - ptr, serde, allocator, num_items);
   ptr += pair.second;
   return std::pair<req_compactor, size_t>(

req/include/req_sketch.hpp

@@ -465,7 +465,9 @@ class req_sketch {
       optional<T>&& min_item, optional<T>&& max_item,
       std::vector<Compactor, AllocCompactor>&& compactors, const Comparator& comparator);
 
+  static void check_k(uint16_t k);
   static void check_preamble_ints(uint8_t preamble_ints, uint8_t num_levels);
+  static void check_num_levels(uint8_t num_levels, bool raw_items);
   static void check_serial_version(uint8_t serial_version);
   static void check_family_id(uint8_t family_id);