From a7ab8270f160122c00856ead973274b9a2f25538 Mon Sep 17 00:00:00 2001 From: Mikhail Chichvarin Date: Sun, 24 May 2026 13:45:08 +0100 Subject: [PATCH] Fix checked arithmetic for allocation sizes Add checked size arithmetic helpers and wire them into allocation, reservation, serialization, and dense metadata paths that previously relied on unchecked size_t math. Co-authored-by: Mikhail Chichvarin <6496186+desertfury@users.noreply.github.com> Co-authored-by: Mikhail Chichvarin Co-authored-by: Ash Vardanian <1983160+ashvardanian@users.noreply.github.com> --- cpp/test.cpp | 43 ++++++ include/usearch/index.hpp | 233 ++++++++++++++++++++++++------ include/usearch/index_dense.hpp | 65 ++++++--- include/usearch/index_plugins.hpp | 104 +++++++++---- 4 files changed, 349 insertions(+), 96 deletions(-) diff --git a/cpp/test.cpp b/cpp/test.cpp index fba9ed78..3ea2f86e 100644 --- a/cpp/test.cpp +++ b/cpp/test.cpp @@ -17,6 +17,7 @@ #include // `std::signal`, `SIGSEGV`, ... #include // `std::fprintf` #include // `std::_Exit` +#include // `std::numeric_limits` #include // `std::shuffle` #include // `std::default_random_engine` @@ -216,6 +217,47 @@ void test_uint40() { expect_eq(u40_default, uint40_t(0u)); } +void test_checked_size_arithmetic() { + + std::printf("Testing checked size arithmetic\n"); + std::size_t max = (std::numeric_limits::max)(); + + checked_size_result_t cast = checked_size_from_u64(42); + expect(cast); + expect_eq(cast.value, static_cast(42)); + if (sizeof(std::size_t) < sizeof(std::uint64_t)) + expect(!checked_size_from_u64((std::numeric_limits::max)())); + + checked_size_result_t sum = checked_add(max - 1, 1); + expect(sum); + expect_eq(sum.value, max); + expect(!checked_add(max, 1)); + + checked_size_result_t product = checked_mul(max / 2, 2); + expect(product); + expect_eq(product.value, max - 1); + expect(!checked_mul(max / 2 + 1, 2)); + + checked_size_result_t fused = checked_mul_add(10, 20, 30); + expect(fused); + expect_eq(fused.value, static_cast(230)); + expect(!checked_mul_add(max / 2 + 1, 2, 0)); + expect(!checked_mul_add(max / 2, 2, 2)); + + checked_size_result_t rounded = checked_round_up(17, 8); + expect(rounded); + expect_eq(rounded.value, static_cast(24)); + expect(!checked_round_up(max, 8)); + + checked_size_result_t power = checked_ceil2(17); + expect(power); + expect_eq(power.value, static_cast(32)); + checked_size_result_t zero_power = checked_ceil2(0); + expect(zero_power); + expect_eq(zero_power.value, static_cast(0)); + expect(!checked_ceil2(max)); +} + /** * @brief Tests the functionality of the custom float16_t type ensuring consistent. */ @@ -1327,6 +1369,7 @@ int main(int, char**) { // Non-default floating-point types may result in many compilation & rounding issues. test_uint40(); + test_checked_size_arithmetic(); test_cosine(10, 10); test_cosine(10, 10); test_cosine(10, 10); diff --git a/include/usearch/index.hpp b/include/usearch/index.hpp index 05ec5e48..de12eab9 100644 --- a/include/usearch/index.hpp +++ b/include/usearch/index.hpp @@ -176,6 +176,40 @@ namespace usearch { using byte_t = char; +struct checked_size_result_t { + std::size_t value; + bool overflow; + constexpr checked_size_result_t(std::size_t value = 0, bool overflow = false) noexcept + : value(value), overflow(overflow) {} + constexpr explicit operator bool() const noexcept { return !overflow; } +}; + +constexpr checked_size_result_t checked_size_overflow() noexcept { return {0, true}; } + +constexpr checked_size_result_t checked_size_from_u64(std::uint64_t value) noexcept { + return value > static_cast((std::numeric_limits::max)()) + ? checked_size_overflow() + : checked_size_result_t{static_cast(value), false}; +} + +constexpr checked_size_result_t checked_add(std::size_t a, std::size_t b) noexcept { + return (std::numeric_limits::max)() - a < b ? checked_size_overflow() + : checked_size_result_t{a + b, false}; +} + +constexpr checked_size_result_t checked_mul(std::size_t a, std::size_t b) noexcept { + return a && b > (std::numeric_limits::max)() / a ? checked_size_overflow() + : checked_size_result_t{a * b, false}; +} + +constexpr checked_size_result_t checked_mul_add_(checked_size_result_t product, std::size_t c) noexcept { + return product ? checked_add(product.value, c) : product; +} + +constexpr checked_size_result_t checked_mul_add(std::size_t a, std::size_t b, std::size_t c) noexcept { + return checked_mul_add_(checked_mul(a, b), c); +} + template std::size_t divide_round_up(std::size_t num) noexcept { return (num + multiple_ak - 1) / multiple_ak; } @@ -184,6 +218,21 @@ inline std::size_t divide_round_up(std::size_t num, std::size_t denominator) noe return (num + denominator - 1) / denominator; } +constexpr checked_size_result_t checked_divide_round_up(std::size_t num, std::size_t denominator) noexcept { + return !denominator ? checked_size_overflow() + : (std::numeric_limits::max)() - num < denominator - 1 + ? checked_size_overflow() + : checked_size_result_t{(num + denominator - 1) / denominator, false}; +} + +constexpr checked_size_result_t checked_round_up_(checked_size_result_t quotient, std::size_t multiple) noexcept { + return quotient ? checked_mul(quotient.value, multiple) : quotient; +} + +constexpr checked_size_result_t checked_round_up(std::size_t num, std::size_t multiple) noexcept { + return checked_round_up_(checked_divide_round_up(num, multiple), multiple); +} + inline std::size_t ceil2(std::size_t v) noexcept { v--; v |= v >> 1; @@ -198,6 +247,14 @@ inline std::size_t ceil2(std::size_t v) noexcept { return v; } +inline checked_size_result_t checked_ceil2(std::size_t v) noexcept { + if (!v) + return checked_size_result_t{0, false}; + if (v > (std::size_t{1} << ((sizeof(std::size_t) * CHAR_BIT) - 1))) + return checked_size_overflow(); + return checked_size_result_t{ceil2(v), false}; +} + /// @brief Simply dereferencing misaligned pointers can be dangerous. template void misaligned_store(void* ptr, at v) noexcept { static_assert(!std::is_reference::value, "Can't store a reference"); @@ -525,9 +582,12 @@ template > class bitset_gt { count_ = 0; } - bitset_gt(std::size_t capacity) noexcept - : slots_((compressed_slot_t*)allocator_t{}.allocate(bits_slots(capacity) * sizeof(compressed_slot_t))), - count_(slots_ ? bits_slots(capacity) : 0u) { + bitset_gt(std::size_t capacity) noexcept { + checked_size_result_t slots_count = checked_divide_round_up(capacity, bits_per_slot()); + checked_size_result_t bytes = + slots_count ? checked_mul(slots_count.value, sizeof(compressed_slot_t)) : slots_count; + slots_ = bytes ? (compressed_slot_t*)allocator_t{}.allocate(bytes.value) : nullptr; + count_ = slots_ ? slots_count.value : 0u; clear(); } @@ -654,10 +714,19 @@ class striped_locks_gt { } striped_locks_gt(std::size_t threads, std::size_t connectivity) noexcept { - std::size_t desired = threads * connectivity * 4; - if (desired < 256) - desired = 256; - count_ = ceil2(desired); + checked_size_result_t desired = checked_mul(threads, connectivity); + desired = desired ? checked_mul(desired.value, std::size_t{4}) : desired; + if (!desired) { + shift_ = 64; + return; + } + + checked_size_result_t count = checked_ceil2((std::max)(desired.value, 256)); + if (!count) { + shift_ = 64; + return; + } + count_ = count.value; shift_ = 64; for (std::size_t n = count_; n > 1; n >>= 1) shift_--; @@ -665,7 +734,13 @@ class striped_locks_gt { // `cache_line_ak`-aligned address inside the allocation, regardless of // what the underlying allocator returns. constexpr std::size_t alignment_k = alignof(padded_lock_t); - raw_bytes_ = count_ * sizeof(padded_lock_t) + alignment_k; + checked_size_result_t raw_bytes = checked_mul_add(count_, sizeof(padded_lock_t), alignment_k); + if (!raw_bytes) { + count_ = 0; + shift_ = 64; + return; + } + raw_bytes_ = raw_bytes.value; raw_ = allocator_t{}.allocate(raw_bytes_); if (!raw_) { raw_bytes_ = 0; @@ -850,10 +925,14 @@ class max_heap_gt { if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - if (new_capacity == 0) + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) + return false; + checked_size_result_t doubled_capacity = checked_mul(capacity_, std::size_t{2}); + if (!doubled_capacity) return false; - new_capacity = (std::max)(new_capacity, (std::max)(capacity_ * 2u, 16u)); + new_capacity = + (std::max)(rounded_capacity.value, (std::max)(doubled_capacity.value, 16u)); auto allocator = allocator_t{}; auto new_elements = allocator.allocate(new_capacity); if (!new_elements) @@ -1018,10 +1097,14 @@ class sorted_buffer_gt { if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - if (new_capacity == 0) + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) + return false; + checked_size_result_t doubled_capacity = checked_mul(capacity_, std::size_t{2}); + if (!doubled_capacity) return false; - new_capacity = (std::max)(new_capacity, (std::max)(capacity_ * 2u, 16u)); + new_capacity = + (std::max)(rounded_capacity.value, (std::max)(doubled_capacity.value, 16u)); auto allocator = allocator_t{}; auto new_elements = allocator.allocate(new_capacity); if (!new_elements) @@ -1246,9 +1329,11 @@ class growing_hash_set_gt { count_ = 0; } - growing_hash_set_gt(std::size_t capacity) noexcept - : slots_((element_t*)allocator_t{}.allocate(ceil2(capacity) * sizeof(element_t))), - capacity_(slots_ ? ceil2(capacity) : 0u), count_(0u) { + growing_hash_set_gt(std::size_t capacity) noexcept : count_(0u) { + checked_size_result_t slots_count = checked_ceil2(capacity); + checked_size_result_t bytes = slots_count ? checked_mul(slots_count.value, sizeof(element_t)) : slots_count; + slots_ = bytes ? (element_t*)allocator_t{}.allocate(bytes.value) : nullptr; + capacity_ = slots_ ? slots_count.value : 0u; clear(); } @@ -1306,12 +1391,21 @@ class growing_hash_set_gt { * @return `true` if enough capacity is available, `false` if memory allocation failed. */ bool reserve(std::size_t new_capacity) noexcept { - new_capacity = (new_capacity * 5u) / 3u; + checked_size_result_t scaled_capacity = checked_mul(new_capacity, std::size_t{5}); + if (!scaled_capacity) + return false; + new_capacity = scaled_capacity.value / 3u; if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - element_t* new_slots = (element_t*)allocator_t{}.allocate(new_capacity * sizeof(element_t)); + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) + return false; + new_capacity = rounded_capacity.value; + checked_size_result_t new_bytes = checked_mul(new_capacity, sizeof(element_t)); + if (!new_bytes) + return false; + element_t* new_slots = (element_t*)allocator_t{}.allocate(new_bytes.value); if (!new_slots) return false; @@ -1414,7 +1508,10 @@ class ring_gt { return false; // prevent data loss if (n <= capacity()) return true; - n = (std::max)(ceil2(n), 64u); + checked_size_result_t rounded_capacity = checked_ceil2(n); + if (!rounded_capacity) + return false; + n = (std::max)(rounded_capacity.value, 64u); element_t* elements = allocator_.allocate(n); if (!elements) return false; @@ -1503,12 +1600,22 @@ struct index_config_t { inline error_t validate() noexcept { if (connectivity == 0) connectivity = default_connectivity(); - if (connectivity_base == 0) - connectivity_base = connectivity * 2; + if (connectivity_base == 0) { + checked_size_result_t default_base = checked_mul(connectivity, std::size_t{2}); + if (!default_base) + return "Connectivity is too large"; + connectivity_base = default_base.value; + } if (connectivity < 2) return "Connectivity must be at least 2, otherwise the index degenerates into ropes"; if (connectivity_base < connectivity) return "Base layer should be at least as connected as the rest of the graph"; + checked_size_result_t neighbors_bytes = + checked_mul_add(connectivity, sizeof(std::uint64_t), sizeof(std::uint32_t)); + checked_size_result_t neighbors_base_bytes = + checked_mul_add(connectivity_base, sizeof(std::uint64_t), sizeof(std::uint32_t)); + if (!neighbors_bytes || !neighbors_base_bytes) + return "Connectivity is too large"; return {}; } @@ -3555,26 +3662,31 @@ class index_gt { // Progress status std::size_t processed = 0; - std::size_t const total = 2 * header.size; + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large to serialize"); + checked_size_result_t total = checked_mul(std::size_t{2}, header_size.value); + if (!total) + return result.failed("Index is too large to serialize"); // Export the number of levels per node // That is both enough to estimate the overall memory consumption, // and to be able to estimate the offsets of every entry in the file. - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { node_t node = node_at_(i); level_t level = node.level(); if (!output(&level, sizeof(level))) return result.failed("Failed to serialize into stream"); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return result.failed("Terminated by user"); } // After that dump the nodes themselves - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { span_bytes_t node_bytes = node_bytes_(node_at_(i)); if (!output(node_bytes.data(), node_bytes.size())) return result.failed("Failed to serialize into stream"); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return result.failed("Terminated by user"); } @@ -3606,10 +3718,16 @@ class index_gt { // Allocate some dynamic memory to read all the levels using levels_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc; - buffer_gt levels(header.size); + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large"); + buffer_gt levels(header_size.value); if (!levels) return result.failed("Out of memory"); - if (!input(levels, header.size * sizeof(level_t))) + checked_size_result_t levels_bytes = checked_mul(header_size.value, sizeof(level_t)); + if (!levels_bytes) + return result.failed("Index is too large"); + if (!input(levels, levels_bytes.value)) return result.failed("Failed to pull nodes levels from the stream"); // Submit metadata @@ -3621,26 +3739,26 @@ class index_gt { pre_ = precompute_(config_); index_limits_t limits; - limits.members = header.size; + limits.members = header_size.value; limits.threads_add = (std::max)(1, old_limits.threads_add); limits.threads_search = (std::max)(1, old_limits.threads_search); if (!reserve(limits)) { reset(); return result.failed("Out of memory"); } - nodes_count_ = header.size; + nodes_count_ = header_size.value; max_level_ = static_cast(header.max_level); entry_slot_ = static_cast(header.entry_slot); // Load the nodes - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { span_bytes_t node_bytes = node_malloc_(levels[i]); if (!input(node_bytes.data(), node_bytes.size())) { reset(); return result.failed("Failed to pull nodes from the stream"); } nodes_[i] = node_t{node_bytes.data()}; - if (!progress(i + 1, header.size)) + if (!progress(i + 1, header_size.value)) return result.failed("Terminated by user"); } return {}; @@ -3787,11 +3905,14 @@ class index_gt { reset(); return result; } + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large"); // Precompute offsets of every node, but before that we need to update the configs // This could have been done with `std::exclusive_scan`, but it's only available from C++17. using offsets_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc; - buffer_gt offsets(header.size); + buffer_gt offsets(header_size.value); if (!offsets) return result.failed("Out of memory"); @@ -3803,33 +3924,47 @@ class index_gt { pre_ = precompute_(config_); misaligned_ptr_gt levels{(byte_t*)file.data() + offset + sizeof(header)}; - offsets[0u] = offset + sizeof(header) + sizeof(level_t) * header.size; - for (std::size_t i = 1; i < header.size; ++i) - offsets[i] = offsets[i - 1] + node_bytes_(levels[i - 1]); + checked_size_result_t levels_bytes = checked_mul(sizeof(level_t), header_size.value); + checked_size_result_t offset_after_header = checked_add(offset, sizeof(header)); + checked_size_result_t first_offset = levels_bytes && offset_after_header + ? checked_add(offset_after_header.value, levels_bytes.value) + : checked_size_overflow(); + if (!first_offset) + return result.failed("Index is too large"); + offsets[0u] = first_offset.value; + for (std::size_t i = 1; i < header_size.value; ++i) { + checked_size_result_t next_offset = checked_add(offsets[i - 1], node_bytes_(levels[i - 1])); + if (!next_offset) + return result.failed("Index is too large"); + offsets[i] = next_offset.value; + } - std::size_t total_bytes = offsets[header.size - 1] + node_bytes_(levels[header.size - 1]); - if (file.size() < total_bytes) { + checked_size_result_t total_bytes = + checked_add(offsets[header_size.value - 1], node_bytes_(levels[header_size.value - 1])); + if (!total_bytes) + return result.failed("Index is too large"); + if (file.size() < total_bytes.value) { reset(); return result.failed("File is corrupted and can't fit all the nodes"); } // Submit metadata and reserve memory index_limits_t limits; - limits.members = header.size; + limits.members = header_size.value; limits.threads_add = (std::max)(1, old_limits.threads_add); limits.threads_search = (std::max)(1, old_limits.threads_search); if (!reserve(limits)) { reset(); return result.failed("Out of memory"); } - nodes_count_ = header.size; + nodes_count_ = header_size.value; max_level_ = static_cast(header.max_level); entry_slot_ = static_cast(header.entry_slot); // Rapidly address all the nodes - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { nodes_[i] = node_t{(byte_t*)file.data() + offsets[i]}; - if (!progress(i + 1, header.size)) + if (!progress(i + 1, header_size.value)) return result.failed("Terminated by user"); } viewed_file_ = std::move(file); @@ -3881,7 +4016,9 @@ class index_gt { // Progress status std::atomic do_tasks{true}; std::atomic processed{0}; - std::size_t const total = 3 * slots_and_levels.size(); + checked_size_result_t total = checked_mul(std::size_t{3}, slots_and_levels.size()); + if (!total) + return; // For every bottom level node, determine its parent cluster executor.dynamic(slots_and_levels.size(), [&](std::size_t thread_idx, std::size_t old_slot_as_uint) { @@ -3894,7 +4031,7 @@ class index_gt { slots_and_levels[old_slot] = {old_slot, cluster, node_at_(old_slot).level()}; ++processed; if (thread_idx == 0) - do_tasks = progress(processed.load(), total); + do_tasks = progress(processed.load(), total.value); return do_tasks.load(); }); if (!do_tasks.load()) @@ -3928,7 +4065,7 @@ class index_gt { neighbor = static_cast(old_slot_to_new[compressed_slot_t(neighbor)]); reordered_nodes[new_slot] = new_node; - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return; } @@ -3937,7 +4074,7 @@ class index_gt { slot_transition(node_at_(old_slot).ckey(), // static_cast(old_slot), // static_cast(new_slot)); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return; } diff --git a/include/usearch/index_dense.hpp b/include/usearch/index_dense.hpp index f59d59a1..050094e4 100644 --- a/include/usearch/index_dense.hpp +++ b/include/usearch/index_dense.hpp @@ -276,15 +276,22 @@ inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* std::uint32_t dimensions_u32[2]{0}; std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32)); - std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32); + checked_size_result_t offset_if_u32 = + checked_mul_add(std::size_t(dimensions_u32[0]), std::size_t(dimensions_u32[1]), sizeof(dimensions_u32)); std::uint64_t dimensions_u64[2]{0}; std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64)); - std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64); + checked_size_result_t rows_if_u64 = checked_size_from_u64(dimensions_u64[0]); + checked_size_result_t columns_if_u64 = checked_size_from_u64(dimensions_u64[1]); + checked_size_result_t offset_if_u64 = + rows_if_u64 && columns_if_u64 ? checked_mul_add(rows_if_u64.value, columns_if_u64.value, sizeof(dimensions_u64)) + : checked_size_overflow(); // Check if it starts with 32-bit - if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) { - if (std::fseek(file.get(), static_cast(offset_if_u32), SEEK_SET) != 0) + checked_size_result_t head_offset_if_u32 = + offset_if_u32 ? checked_add(offset_if_u32.value, sizeof(index_dense_head_buffer_t)) : offset_if_u32; + if (head_offset_if_u32 && head_offset_if_u32.value < file_size) { + if (std::fseek(file.get(), static_cast(offset_if_u32.value), SEEK_SET) != 0) return result.failed(std::strerror(errno)); read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get()); if (!read) @@ -299,8 +306,10 @@ inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* } // Check if it starts with 64-bit - if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) { - if (std::fseek(file.get(), static_cast(offset_if_u64), SEEK_SET) != 0) + checked_size_result_t head_offset_if_u64 = + offset_if_u64 ? checked_add(offset_if_u64.value, sizeof(index_dense_head_buffer_t)) : offset_if_u64; + if (head_offset_if_u64 && head_offset_if_u64.value < file_size) { + if (std::fseek(file.get(), static_cast(offset_if_u64.value), SEEK_SET) != 0) return result.failed(std::strerror(errno)); read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get()); if (!read) @@ -341,15 +350,22 @@ inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_map // Check if it starts with 32-bit std::uint32_t dimensions_u32[2]{0}; std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32)); - std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32); + checked_size_result_t offset_if_u32 = + checked_mul_add(std::size_t(dimensions_u32[0]), std::size_t(dimensions_u32[1]), sizeof(dimensions_u32)); std::uint64_t dimensions_u64[2]{0}; std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64)); - std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64); + checked_size_result_t rows_if_u64 = checked_size_from_u64(dimensions_u64[0]); + checked_size_result_t columns_if_u64 = checked_size_from_u64(dimensions_u64[1]); + checked_size_result_t offset_if_u64 = + rows_if_u64 && columns_if_u64 ? checked_mul_add(rows_if_u64.value, columns_if_u64.value, sizeof(dimensions_u64)) + : checked_size_overflow(); // Check if it starts with 32-bit - if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) { - std::memcpy(&result.head_buffer, file_data + offset_if_u32, sizeof(index_dense_head_buffer_t)); + checked_size_result_t head_offset_if_u32 = + offset_if_u32 ? checked_add(offset_if_u32.value, sizeof(index_dense_head_buffer_t)) : offset_if_u32; + if (head_offset_if_u32 && head_offset_if_u32.value < file_size) { + std::memcpy(&result.head_buffer, file_data + offset_if_u32.value, sizeof(index_dense_head_buffer_t)); result.config.exclude_vectors = false; result.config.use_64_bit_dimensions = false; if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0) @@ -357,8 +373,10 @@ inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_map } // Check if it starts with 64-bit - if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) { - std::memcpy(&result.head_buffer, file_data + offset_if_u64, sizeof(index_dense_head_buffer_t)); + checked_size_result_t head_offset_if_u64 = + offset_if_u64 ? checked_add(offset_if_u64.value, sizeof(index_dense_head_buffer_t)) : offset_if_u64; + if (head_offset_if_u64 && head_offset_if_u64.value < file_size) { + std::memcpy(&result.head_buffer, file_data + offset_if_u64.value, sizeof(index_dense_head_buffer_t)); result.config.exclude_vectors = false; result.config.use_64_bit_dimensions = true; if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0) @@ -710,9 +728,11 @@ class index_dense_gt { * left untouched in that case so the index stays consistent. */ bool try_change_metric(metric_t metric) noexcept { - std::size_t needed_bytes = limits().threads() * metric.bytes_per_vector(); - if (needed_bytes > cast_buffer_.size()) { - cast_buffer_t new_buffer(needed_bytes); + checked_size_result_t needed_bytes = checked_mul(limits().threads(), metric.bytes_per_vector()); + if (!needed_bytes) + return false; + if (needed_bytes.value > cast_buffer_.size()) { + cast_buffer_t new_buffer(needed_bytes.value); if (!new_buffer) return false; cast_buffer_ = std::move(new_buffer); @@ -1055,7 +1075,10 @@ class index_dense_gt { available_threads_.push(i); // Allocate a buffer for the casted vectors. - cast_buffer_t cast_buffer(limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return false; + cast_buffer_t cast_buffer(cast_buffer_bytes.value); if (!cast_buffer) return false; cast_buffer_ = std::move(cast_buffer); @@ -1269,7 +1292,10 @@ class index_dense_gt { config_.multi = head.multi; metric_ = metric_t::builtin(head.dimensions, head.kind_metric, head.kind_scalar); - cast_buffer_ = cast_buffer_t(new_limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(new_limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return result.failed("Failed to allocate memory for the casts"); + cast_buffer_ = cast_buffer_t(cast_buffer_bytes.value); if (!cast_buffer_) return result.failed("Failed to allocate memory for the casts"); casts_ = casts_punned_t::make(head.kind_scalar); @@ -1384,7 +1410,10 @@ class index_dense_gt { config_.multi = head.multi; metric_ = metric_t::builtin(head.dimensions, head.kind_metric, head.kind_scalar); // available_threads_.size() will be updated to old_limits.threads() later in this - cast_buffer_ = cast_buffer_t(new_limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(new_limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return result.failed("Failed to allocate memory for the casts"); + cast_buffer_ = cast_buffer_t(cast_buffer_bytes.value); if (!cast_buffer_) return result.failed("Failed to allocate memory for the casts"); casts_ = casts_punned_t::make(head.kind_scalar); diff --git a/include/usearch/index_plugins.hpp b/include/usearch/index_plugins.hpp index 38fa2670..6877e9c0 100644 --- a/include/usearch/index_plugins.hpp +++ b/include/usearch/index_plugins.hpp @@ -1572,21 +1572,24 @@ class aligned_allocator_gt { constexpr std::size_t alignment() const { return alignment_ak; } pointer allocate(size_type length) const { - std::size_t length_bytes = alignment_ak * divide_round_up(length * sizeof(value_type)); - // Avoid overflow - if (length > length_bytes) + checked_size_result_t bytes = checked_mul(length, sizeof(value_type)); + if (!bytes) return nullptr; + checked_size_result_t length_bytes = checked_round_up(bytes.value, alignment_ak); + if (!length_bytes) + return nullptr; + std::size_t alignment = alignment_ak; #if defined(USEARCH_DEFINED_WINDOWS) - return (pointer)_aligned_malloc(length_bytes, alignment); + return (pointer)_aligned_malloc(length_bytes.value, alignment); #elif defined(USEARCH_DEFINED_APPLE) || defined(USEARCH_DEFINED_ANDROID) // Apple Clang keeps complaining that `aligned_alloc` is only available // with macOS 10.15 and newer or Android API >= 28, so let's use `posix_memalign` there. void* result = nullptr; - int status = posix_memalign(&result, alignment, length_bytes); + int status = posix_memalign(&result, alignment, length_bytes.value); return status == 0 ? (pointer)result : nullptr; #else - return (pointer)aligned_alloc(alignment, length_bytes); + return (pointer)aligned_alloc(alignment, length_bytes.value); #endif } @@ -1615,7 +1618,10 @@ class page_allocator_t { * @return A pointer to the allocated memory block, or `nullptr` if allocation fails. */ byte_t* allocate(std::size_t count_bytes) const noexcept { - count_bytes = divide_round_up(count_bytes, page_size()) * page_size(); + checked_size_result_t rounded_bytes = checked_round_up(count_bytes, page_size()); + if (!rounded_bytes) + return nullptr; + count_bytes = rounded_bytes.value; #if defined(USEARCH_DEFINED_WINDOWS) return (byte_t*)(::VirtualAlloc(NULL, count_bytes, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)); #else @@ -1628,7 +1634,10 @@ class page_allocator_t { #if defined(USEARCH_DEFINED_WINDOWS) ::VirtualFree(page_pointer, 0, MEM_RELEASE); #else - count_bytes = divide_round_up(count_bytes, page_size()) * page_size(); + checked_size_result_t rounded_bytes = checked_round_up(count_bytes, page_size()); + if (!rounded_bytes) + return; + count_bytes = rounded_bytes.value; munmap(page_pointer, count_bytes); #endif } @@ -1724,25 +1733,39 @@ template class memory_mapping_allocator_gt { * @return A pointer to the allocated memory block, or `nullptr` if allocation fails. */ inline byte_t* allocate(std::size_t count_bytes) noexcept { - std::size_t extended_bytes = divide_round_up(count_bytes) * alignment_ak; + checked_size_result_t extended_bytes = checked_round_up(count_bytes, alignment_ak); + if (!extended_bytes) + return nullptr; std::unique_lock lock(mutex_); - if (!last_arena_ || (last_usage_ + extended_bytes >= last_capacity_)) { - std::size_t new_cap = (std::max)(last_capacity_, ceil2(extended_bytes)) * capacity_multiplier(); - byte_t* new_arena = page_allocator_t{}.allocate(new_cap); + checked_size_result_t next_usage = checked_add(last_usage_, extended_bytes.value); + if (!next_usage) + return nullptr; + if (!last_arena_ || (next_usage.value >= last_capacity_)) { + checked_size_result_t rounded_bytes = checked_ceil2(extended_bytes.value); + if (!rounded_bytes) + return nullptr; + checked_size_result_t new_cap = + checked_mul((std::max)(last_capacity_, rounded_bytes.value), capacity_multiplier()); + if (!new_cap) + return nullptr; + checked_size_result_t new_total_allocated = checked_add(total_allocated_, new_cap.value); + if (!new_total_allocated) + return nullptr; + byte_t* new_arena = page_allocator_t{}.allocate(new_cap.value); if (!new_arena) return nullptr; std::memcpy(new_arena, &last_arena_, sizeof(byte_t*)); - std::memcpy(new_arena + sizeof(byte_t*), &new_cap, sizeof(std::size_t)); + std::memcpy(new_arena + sizeof(byte_t*), &new_cap.value, sizeof(std::size_t)); wasted_space_ += total_reserved(); last_arena_ = new_arena; - last_capacity_ = new_cap; + last_capacity_ = new_cap.value; last_usage_ = head_size(); - total_allocated_ += new_cap; + total_allocated_ = new_total_allocated.value; } - wasted_space_ += extended_bytes - count_bytes; - return last_arena_ + exchange(last_usage_, last_usage_ + extended_bytes); + wasted_space_ += extended_bytes.value - count_bytes; + return last_arena_ + exchange(last_usage_, last_usage_ + extended_bytes.value); } /** @@ -3774,7 +3797,10 @@ class flat_hash_multi_set_gt { } // Allocate new memory - data_ = (char*)allocator_t{}.allocate(other.buckets_ * bytes_per_bucket()); + checked_size_result_t bytes = checked_mul(other.buckets_, bytes_per_bucket()); + if (!bytes) + usearch_raise_runtime_error("failed memory allocation"); + data_ = (char*)allocator_t{}.allocate(bytes.value); if (!data_) usearch_raise_runtime_error("failed memory allocation"); @@ -3814,7 +3840,10 @@ class flat_hash_multi_set_gt { allocator_t{}.deallocate(data_, buckets_ * bytes_per_bucket()); // Allocate new memory - data_ = (char*)allocator_t{}.allocate(other.buckets_ * bytes_per_bucket()); + checked_size_result_t bytes = checked_mul(other.buckets_, bytes_per_bucket()); + if (!bytes) + usearch_raise_runtime_error("failed memory allocation"); + data_ = (char*)allocator_t{}.allocate(bytes.value); if (!data_) usearch_raise_runtime_error("failed memory allocation"); @@ -3863,22 +3892,37 @@ class flat_hash_multi_set_gt { } bool try_reserve(std::size_t capacity) noexcept { - if (capacity * 3u <= capacity_slots_ * 2u) + if (capacity <= (capacity_slots_ / 3u) * 2u) return true; // Calculate new sizes - std::size_t new_slots = ceil2((capacity * 3ul) / 2ul); - std::size_t new_buckets = divide_round_up(new_slots); - new_slots = new_buckets * slots_per_bucket(); // This must be a power of two! - std::size_t new_bytes = new_buckets * bytes_per_bucket(); + checked_size_result_t scaled_capacity = checked_mul(capacity, std::size_t{3}); + if (!scaled_capacity) + return false; + checked_size_result_t slots_needed = checked_divide_round_up(scaled_capacity.value, std::size_t{2}); + if (!slots_needed) + return false; + checked_size_result_t new_slots_checked = checked_ceil2(slots_needed.value); + if (!new_slots_checked) + return false; + checked_size_result_t new_buckets_checked = + checked_divide_round_up(new_slots_checked.value, slots_per_bucket()); + if (!new_buckets_checked) + return false; + checked_size_result_t new_slots = checked_mul(new_buckets_checked.value, slots_per_bucket()); + if (!new_slots) + return false; + checked_size_result_t new_bytes = checked_mul(new_buckets_checked.value, bytes_per_bucket()); + if (!new_bytes) + return false; // Allocate new memory - char* new_data = (char*)allocator_t{}.allocate(new_bytes); + char* new_data = (char*)allocator_t{}.allocate(new_bytes.value); if (!new_data) return false; // Initialize new buckets to empty - std::memset(new_data, 0, new_bytes); + std::memset(new_data, 0, new_bytes.value); // Rehash and copy existing elements to new_data hash_t hasher; @@ -3889,7 +3933,7 @@ class flat_hash_multi_set_gt { // Rehash std::size_t hash_value = hasher(old_slot.element); - std::size_t new_slot_index = hash_value & (new_slots - 1); + std::size_t new_slot_index = hash_value & (new_slots.value - 1); // Linear probing to find an empty slot in new_data while (true) { @@ -3899,7 +3943,7 @@ class flat_hash_multi_set_gt { new_slot.header.populated |= new_slot.mask; break; } - new_slot_index = (new_slot_index + 1) & (new_slots - 1); + new_slot_index = (new_slot_index + 1) & (new_slots.value - 1); } } @@ -3907,8 +3951,8 @@ class flat_hash_multi_set_gt { if (data_) allocator_t{}.deallocate(data_, buckets_ * bytes_per_bucket()); data_ = new_data; - buckets_ = new_buckets; - capacity_slots_ = new_slots; + buckets_ = new_buckets_checked.value; + capacity_slots_ = new_slots.value; return true; }