Protect against integer overflow with sanisizer.

In particular, the calculation of search_k needs some care because it is
computed from a floating-point multiplication. We don't throw an error
if the product overflows but instead we just cap it at the int's max.

Author: LTLA

CMakeLists.txt

@@ -23,10 +23,11 @@ if(KNNCOLLE_ANNOY_FETCH_EXTERN)
     add_subdirectory(extern)
 else()
     find_package(knncolle_knncolle 3.0.0 CONFIG REQUIRED)
+    find_package(ltla_sanisizer 0.1.0 CONFIG REQUIRED)
     find_package(Annoy CONFIG REQUIRED)
 endif()
 
-target_link_libraries(knncolle_annoy INTERFACE knncolle::knncolle Annoy::Annoy)
+target_link_libraries(knncolle_annoy INTERFACE knncolle::knncolle ltla::sanisizer Annoy::Annoy)
 
 # Tests
 if(CMAKE_PROJECT_NAME STREQUAL PROJECT_NAME)

cmake/Config.cmake.in

@@ -2,6 +2,7 @@
 
 include(CMakeFindDependencyMacro)
 find_dependency(knncolle_knncolle 3.0.0 CONFIG REQUIRED)
+find_dependency(ltla_sanisizer 0.1.0 CONFIG REQUIRED)
 find_dependency(Annoy CONFIG REQUIRED)
 
 include("${CMAKE_CURRENT_LIST_DIR}/knncolle_knncolle_annoyTargets.cmake")

extern/CMakeLists.txt

@@ -6,6 +6,13 @@ FetchContent_Declare(
   GIT_TAG master # ^3.0.0
 )
 
+FetchContent_Declare(
+  sanisizer 
+  GIT_REPOSITORY https://github.com/LTLA/sanisizer
+  GIT_TAG master # ^0.1.0
+)
+
+
 FetchContent_Declare(
   Annoy 
   GIT_REPOSITORY https://github.com/spotify/Annoy

include/knncolle_annoy/Annoy.hpp

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <cstring>
 #include <stdexcept>
+#include <limits>
 
 #include "knncolle/knncolle.hpp"
 #include "annoy/annoylib.h"
@@ -68,18 +69,28 @@ class AnnoySearcher final : public knncolle::Searcher<Index_, Data_, Distance_>
     static constexpr bool same_internal_distance = std::is_same<Distance_, AnnoyData_>::value;
     typename std::conditional<!same_internal_distance, std::vector<AnnoyData_>, bool>::type my_distances;
 
-    int get_search_k(int k) const {
+    typedef int SearchKType;
+
+    Index_ my_capped_k = 0;
+
+    SearchKType get_search_k(Index_ k) const {
         if (my_parent.my_search_mult < 1) {
             return -1; // instructs Annoy to use k * num_trees. 
-        } else {
+        } else if (k <= my_capped_k) {
             return my_parent.my_search_mult * static_cast<double>(k) + 0.5; // rounded.
+        } else {
+            return std::numeric_limits<SearchKType>::max(); // cap to avoid overflow.
         }
     }
 
 public:
     AnnoySearcher(const AnnoyPrebuilt<Index_, Data_, Distance_, AnnoyDistance_, AnnoyIndex_, AnnoyData_, AnnoyRng_, AnnoyThreadPolicy_>& parent) : my_parent(parent) {
         if constexpr(!same_internal_data) {
-            my_buffer.resize(my_parent.my_dim);
+            sanisizer::resize(my_buffer, my_parent.my_dim);
+        }
+
+        if (my_parent.my_search_mult >= 1) {
+            my_capped_k = static_cast<double>(std::numeric_limits<SearchKType>::max()) / my_parent.my_search_mult;
         }
     }
 
@@ -139,18 +150,26 @@ class AnnoySearcher final : public knncolle::Searcher<Index_, Data_, Distance_>
 
 public:
     void search(Index_ i, Index_ k, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
-        Index_ kp1 = k + 1; // +1, as it forgets to discard 'self'.
+        // +1, as it forgets to discard 'self'. This should not overflow as k < num_obs.
+        const auto kp1 = k + 1;
+
         auto ptrs = obtain_pointers(output_indices, output_distances, kp1);
         auto icopy_ptr = ptrs.first;
         auto dcopy_ptr = ptrs.second;
 
-        my_parent.my_index.get_nns_by_item(i, kp1, get_search_k(kp1), icopy_ptr, dcopy_ptr);
+        my_parent.my_index.get_nns_by_item(
+            i,
+            sanisizer::cast<std::size_t>(sanisizer::attest_gez(kp1)),
+            get_search_k(kp1),
+            icopy_ptr,
+            dcopy_ptr
+        );
 
         std::size_t at;
         {
             const auto& cur_i = *icopy_ptr;
             at = cur_i.size();
-            AnnoyIndex_ icopy = i;
+            const AnnoyIndex_ icopy = i; // cast to AnnoyIndex_ to avoid signed/unsigned comparisons.
             for (std::size_t x = 0, end = cur_i.size(); x < end; ++x) {
                 if (cur_i[x] == icopy) {
                     at = x;
@@ -182,7 +201,13 @@ class AnnoySearcher final : public knncolle::Searcher<Index_, Data_, Distance_>
         auto icopy_ptr = ptrs.first;
         auto dcopy_ptr = ptrs.second;
 
-        my_parent.my_index.get_nns_by_vector(query, k, get_search_k(k), icopy_ptr, dcopy_ptr);
+        my_parent.my_index.get_nns_by_vector(
+            query,
+            sanisizer::cast<std::size_t>(sanisizer::attest_gez(k)),
+            get_search_k(k),
+            icopy_ptr,
+            dcopy_ptr
+        );
 
         if (output_indices) {
             if constexpr(!same_internal_index) {
@@ -229,14 +254,17 @@ class AnnoyPrebuilt final : public knncolle::Prebuilt<Index_, Data_, Distance_>
         my_search_mult(options.search_mult),
         my_index(my_dim)
     {
+        // check that we can, in fact, represent the number of observations in the specified AnnoyIndex_ type.
+        sanisizer::cast<AnnoyIndex_>(sanisizer::attest_gez(my_obs));
+
         auto work = data.new_known_extractor();
         if constexpr(std::is_same<Data_, AnnoyData_>::value) {
             for (Index_ i = 0; i < my_obs; ++i) {
                 auto ptr = work->next();
                 my_index.add_item(i, ptr);
             }
         } else {
-            std::vector<AnnoyData_> incoming(my_dim);
+            auto incoming = sanisizer::create<std::vector<AnnoyData_> >(my_dim);
             for (Index_ i = 0; i < my_obs; ++i) {
                 auto ptr = work->next();
                 std::copy_n(ptr, my_dim, incoming.begin());
@@ -310,8 +338,11 @@ class AnnoyPrebuilt final : public knncolle::Prebuilt<Index_, Data_, Distance_>
 
         // Not bothering to save anything else; the RNG and thread policy
         // should not be relevant once the index is built.
+
+        // For reasons unknown to us, AnnoyIndex::save() is not const, so we have to do it manually.
+        // Hopefully this will be fixed in the future.
         const auto idxpath = prefix + "index";
-        knncolle::quick_save(idxpath, reinterpret_cast<char*>(my_index.get_nodes()), static_cast<std::size_t>(my_index.get_s()) * my_index.get_n_nodes());
+        knncolle::quick_save(idxpath, reinterpret_cast<char*>(my_index.get_nodes()), sanisizer::product<std::streamsize>(my_index.get_s(), my_index.get_n_nodes()));
     }
 
     AnnoyPrebuilt(const std::string prefix, std::size_t ndim) : my_dim(ndim), my_index(ndim) {

tests/src/Annoy.cpp

@@ -248,16 +248,40 @@ class AnnoyMiscTest : public TestCore, public ::testing::Test {
 TEST_F(AnnoyMiscTest, SearchMult) {
     knncolle_annoy::AnnoyBuilder<int, double, double, Annoy::Euclidean> ab;
     auto& an_opt = ab.get_options();
-    an_opt.search_mult = 20;
 
     knncolle::SimpleMatrix<int, double> mat(ndim, nobs, data.data());
-    auto ptr = ab.build_known_shared(mat); // test coverage for the known override.
-    auto sptr = ptr->initialize();
+    auto default_ptr = ab.build_known_shared(mat); // test coverage for the known override.
+
+    // Manually setting the search multiplier and checking we get the same result.
+    {
+        an_opt.search_mult = an_opt.num_trees;
+        auto manual_ptr = ab.build_known_shared(mat);
+        auto dptr = default_ptr->initialize();
+        auto mptr = manual_ptr->initialize();
+
+        std::vector<int> output_i, output_i2;
+        std::vector<double> output_d, output_d2;
+        for (int x = 0; x < nobs; ++x) {
+            dptr->search(x, 5, &output_i, &output_d);
+            mptr->search(x, 5, &output_i2, &output_d2);
+            EXPECT_EQ(output_i, output_i2);
+            EXPECT_EQ(output_d, output_d2);
+        }
+    }
 
-    std::vector<int> ires;
-    std::vector<double> dres;
-    sptr->search(0, 10, &ires, &dres);
-    sanity_checks(ires, dres, 10, 0);
+    // Using an obscenely large multiplier.
+    {
+        an_opt.search_mult = std::numeric_limits<double>::max();
+        auto super_ptr = ab.build_known_shared(mat);
+        auto sptr = super_ptr->initialize();
+
+        std::vector<int> output_i;
+        std::vector<double> output_d;
+        for (int x = 0; x < nobs; ++x) {
+            sptr->search(x, 5, &output_i, &output_d);
+            sanity_checks(output_i, output_d, 5, x);
+        }
+    }
 }
 
 TEST(Annoy, Duplicates) {