diff --git a/theta/include/compact_theta_sketch_parser_impl.hpp b/theta/include/compact_theta_sketch_parser_impl.hpp index d728c0a6..a801d1ea 100644 --- a/theta/include/compact_theta_sketch_parser_impl.hpp +++ b/theta/include/compact_theta_sketch_parser_impl.hpp @@ -80,8 +80,9 @@ auto compact_theta_sketch_parser::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(ptr) + COMPACT_SKETCH_SINGLE_ENTRY_U64, 64}; } - const uint32_t num_entries = reinterpret_cast(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(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32]; const uint64_t* entries = reinterpret_cast(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::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(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32]; uint64_t theta = reinterpret_cast(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::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(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::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(ptr)[COMPACT_SKETCH_NUM_ENTRIES_U32]; uint64_t theta = reinterpret_cast(ptr)[COMPACT_SKETCH_THETA_U64]; bool is_empty = (num_entries == 0) && (theta == theta_constants::MAX_THETA); diff --git a/theta/test/theta_sketch_test.cpp b/theta/test/theta_sketch_test.cpp index 96ff7166..adc0713f 100644 --- a/theta/test/theta_sketch_test.cpp +++ b/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 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 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 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 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 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 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 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;