Skip to content

Commit b02677e

Browse files
MayCXCmceachen
authored andcommitted
fix(sqlite-vec.c): keep a static blob KNN inside its candidate bitmap
vec_static_blob_entriesFilter() rounds the blob's vector count up to a multiple of 8 (bsize) for its candidate bitmap, then clears the bits past the last vector with a loop that starts at bsize itself, one bit past the bitmap's end. When the vector count fills the bitmap's bytes exactly, as 64 vectors do, SQLite's allocator has no padding to absorb the write, and AddressSanitizer reports a heap-buffer-overflow in bitmap_set() from the filter. The loop now clears the bits from the vector count up to bsize. The memory test's static blob runs a KNN query over its 64 vectors; under `make test-asan` it reports the overflow without the fix and passes with it. test-loadable.py checks the same query's rows; against the AddressSanitizer build it also reports the overflow without the fix.
1 parent 01c827e commit b02677e

4 files changed

Lines changed: 37 additions & 1 deletion

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
- A KNN query that takes its query vector, `k`, or `mmr_lambda` from another table of a join now uses it. `select v.rowid from p join v on v.embedding match ? and v.k = p.k` failed with `A LIMIT or 'k = ?' constraint is required`; the same join with `v.k = 3 and v.mmr_lambda = p.lambda` returned no rows; and `select v.rowid from v join p on p.id = v.tag where v.embedding match p.vector and v.k = 2`, with `p.id` an integer primary key, could return no rows. These joins now run one KNN per row of the other table, so a filter on the `vec0` table whose value comes from that table applies before the k nearest rows are chosen. Ported in part from the unmerged [vlasky#14](https://github.com/vlasky/sqlite-vec/pull/14) ([asg017#331](https://github.com/asg017/sqlite-vec/pull/331)), which also applies filters whose value comes from a joined table before the k nearest rows are chosen; that part is not ported, and ARCHITECTURE.md's idxStr section says why.
1414
- Fixed a memory leak: closing a connection freed the `vec_static_blobs` registry but not the name of each blob registered in it.
15+
- Fixed a write one byte past the end of a heap allocation in KNN queries on a `vec_static_blob_entries` table. SQLite's allocator rounds the allocation up, which hides the write unless the blob's vector count rounds up to a multiple of 64, as 64 vectors do.
1516

1617
## [2.1.1] - 2026-10-09
1718

‎sqlite-vec.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12606,7 +12606,7 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
1260612606
}
1260712607

1260812608
bitmap_fill(candidates, bsize);
12609-
for (size_t i = bsize; i >= p->blob->nvectors; i--) {
12609+
for (size_t i = p->blob->nvectors; i < bsize; i++) {
1261012610
bitmap_set(candidates, i, 0);
1261112611
}
1261212612
i32 k_used = 0;

‎tests/memory-test.c‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -995,12 +995,15 @@ static int test_static_blobs(void) {
995995
int rc;
996996
int result = 1;
997997

998+
// 64 vectors fill a KNN's candidate bitmap to its last byte, so a write past
999+
// its end lands outside the allocation.
9981000
static float vectors[64 * 4];
9991001
for (int i = 0; i < 64; i++) {
10001002
for (int j = 0; j < 4; j++) {
10011003
vectors[i * 4 + j] = i / 64.0f;
10021004
}
10031005
}
1006+
float query[4] = {0.5f, 0.5f, 0.5f, 0.5f};
10041007

10051008
rc = sqlite3_open(":memory:", &db);
10061009
CHECK_OK(rc, "open database");
@@ -1036,6 +1039,25 @@ static int test_static_blobs(void) {
10361039
sqlite3_finalize(stmt);
10371040
stmt = NULL;
10381041

1042+
rc = sqlite3_prepare_v2(db,
1043+
"SELECT rowid FROM w WHERE vector MATCH ?1 AND k = 3 "
1044+
"ORDER BY distance",
1045+
-1, &stmt, NULL);
1046+
CHECK_OK(rc, "prepare static blob KNN");
1047+
sqlite3_bind_blob(stmt, 1, query, sizeof(query), SQLITE_STATIC);
1048+
sqlite3_int64 expected[3] = {32, 33, 31};
1049+
for (int i = 0; i < 3; i++) {
1050+
rc = sqlite3_step(stmt);
1051+
CHECK_ROW(rc, "static blob KNN row");
1052+
if (sqlite3_column_int64(stmt, 0) != expected[i]) {
1053+
fprintf(stderr, "FAILED: static blob KNN row %d is %lld, not %lld\n", i,
1054+
sqlite3_column_int64(stmt, 0), expected[i]);
1055+
goto cleanup;
1056+
}
1057+
}
1058+
rc = sqlite3_step(stmt);
1059+
CHECK_DONE(rc, "static blob KNN end");
1060+
10391061
printf(" PASS: static blobs\n");
10401062
result = 0;
10411063

‎tests/test-loadable.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,19 @@ def test_vec_static_blob_entries():
265265
]
266266

267267

268+
def test_vec_static_blob_entries_knn_with_whole_candidate_bytes():
269+
# 64 vectors fill the KNN's candidate bitmap to its last byte. The blob
270+
# reads the array's memory, so the array stays referenced while it is queried.
271+
db = connect(EXT_PATH, extra_entrypoint="sqlite3_vec_static_blobs_init")
272+
vectors = np.array([[i / 64] * 4 for i in range(64)], np.float32)
273+
register_numpy(db, "w", vectors)
274+
rows = db.execute(
275+
"select rowid from w where vector match ? and k = 3 order by distance",
276+
[np.array([0.5] * 4, dtype=np.float32)],
277+
).fetchall()
278+
assert [r[0] for r in rows] == [32, 33, 31]
279+
280+
268281
def test_limits():
269282
db = connect(EXT_PATH)
270283
with _raises(

0 commit comments

Comments
 (0)