From efde1ab6bf1bcfe3f61e55be39b973bff3933190 Mon Sep 17 00:00:00 2001
From: Arun Sharma <arun@sharma-home.net>
Date: Wed, 9 Sep 2026 08:05:20 -0700
Subject: [PATCH] Fix heap-buffer-overflow in HNSW visited bitmap

The visited bitmap was sized from TableStats::getTableCard(), an
estimated (possibly stale) cardinality, and indexed without bounds
checks by graph neighbor offsets in oneHopSearch and related paths.
When the estimate was smaller than the actual max node offset this
caused an ASan heap-buffer-overflow READ in
VisitedState::contains (LadybugDB/ladybug#940).

- Size the bitmap from the exact NodeTable::getNumTotalRows() instead
  of the estimated table cardinality.
- Make VisitedState bounds-safe and auto-growing: contains() returns
  false for out-of-range offsets, add() grows geometrically while
  preserving existing marks.
- Proactively resize the bitmap to the current row count in
  initLayerSearchState to cover growth between query init and search.
---
 vector/src/function/query_hnsw_index.cpp |  6 +++-
 vector/src/include/index/hnsw_index.h    | 38 ++++++++++++++++++++++--
 vector/src/index/hnsw_index.cpp          |  4 +++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/vector/src/function/query_hnsw_index.cpp b/vector/src/function/query_hnsw_index.cpp
index 48dfb84b..21525658 100644
--- a/vector/src/function/query_hnsw_index.cpp
+++ b/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);
 }
 
diff --git a/vector/src/include/index/hnsw_index.h b/vector/src/include/index/hnsw_index.h
index a3ac95bd..fdcc3705 100644
--- a/vector/src/include/index/hnsw_index.h
+++ b/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 {
diff --git a/vector/src/index/hnsw_index.cpp b/vector/src/index/hnsw_index.cpp
index 100a7ad1..8b1e07cb 100644
--- a/vector/src/index/hnsw_index.cpp
+++ b/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;