Fix heap-buffer-overflow in HNSW visited bitmap

vector/src/function/query_hnsw_index.cpp

@@ -369,7 +369,11 @@ static std::unique_ptr<TableFuncSharedState> initQueryHNSWSharedState(
     auto nodeTable = storage::StorageManager::Get(*context)
                          ->getTable(bindData->nodeTableEntry->getTableID())
                          ->ptrCast<storage::NodeTable>();
-    auto numNodes = nodeTable->getStats(transaction::Transaction::Get(*context)).getTableCard();
+    // NB: TableStats::getTableCard() is an estimated (possibly stale) cardinality and must not
+    // be used to size the visited bitmap. Use the exact row count instead. VisitedState is
+    // additionally bounds-checked and auto-growing, so a stale size cannot overflow.
+    auto numNodes =
+        nodeTable->getNumTotalRows(transaction::Transaction::Get(*context));
     return std::make_unique<QueryHNSWIndexSharedState>(nodeTable, numNodes);
 }
 

vector/src/include/index/hnsw_index.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cstring>
 #include <memory>
 #include <queue>
 
@@ -55,9 +56,42 @@ struct VisitedState {
 
     // NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const.
     void reset() { memset(visited.get(), 0, size); }
+    // Grow the bitmap so that `offset` is addressable. Preserves existing marks.
+    void ensureCapacity(common::offset_t offset) {
+        if (offset < size) {
+            return;
+        }
+        // Grow geometrically to amortize repeated growth when the table grows
+        // between sizing and search (e.g. stale cardinality estimates).
+        common::offset_t newSize = size == 0 ? offset + 1 : size;
+        while (newSize <= offset) {
+            newSize *= 2;
+        }
+        auto newVisited = std::make_unique<uint8_t[]>(newSize);
+        memset(newVisited.get(), 0, newSize);
+        if (size > 0) {
+            memcpy(newVisited.get(), visited.get(), size);
+        }
+        visited = std::move(newVisited);
+        size = newSize;
+    }
+    void resize(common::offset_t newSize) {
+        if (newSize <= size) {
+            return;
+        }
+        ensureCapacity(newSize - 1);
+    }
     // NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const.
-    void add(common::offset_t offset) { visited[offset] = 1; }
-    bool contains(common::offset_t offset) const { return visited[offset]; }
+    void add(common::offset_t offset) {
+        ensureCapacity(offset);
+        visited[offset] = 1;
+    }
+    bool contains(common::offset_t offset) const {
+        if (offset >= size) {
+            return false;
+        }
+        return visited[offset];
+    }
 };
 
 struct HNSWStorageInfo final : storage::IndexStorageInfo {

vector/src/index/hnsw_index.cpp

@@ -958,6 +958,10 @@ common::offset_t OnDiskHNSWIndex::searchNNInUpperLayer(const EmbeddingHandle& qu
 
 void OnDiskHNSWIndex::initLayerSearchState(Transaction* transaction, HNSWSearchState& searchState,
     bool isUpperLayer) const {
+    // The visited bitmap is sized at query init time and may be smaller than the table at
+    // search time (e.g. concurrent inserts). Grow it proactively; VisitedState::add/contains
+    // are additionally bounds-safe so a stale size can never cause an out-of-bounds access.
+    searchState.visited.resize(nodeTable.getNumTotalRows(transaction));
     searchState.visited.reset();
     const auto& hnswStorageInfo = storageInfo->cast<HNSWStorageInfo>();
     const auto& hnswGraph = isUpperLayer ? searchState.upperGraph : searchState.lowerGraph;