Skip to content

Commit 5522e86

Browse files
asg017claude
andcommitted
Validate validity/rowids blob sizes in rescore KNN path
The rescore KNN loop read validity and rowids blobs from the chunks iterator without checking their sizes matched chunk_size expectations. A truncated or corrupt blob could cause OOB reads in bitmap_copy or rowid array access. The flat KNN path already had these checks. Adds corruption tests: truncated rowids blob and truncated validity blob both produce errors instead of crashes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent f2c9fb8 commit 5522e86

2 files changed

Lines changed: 36 additions & 5 deletions

File tree

sqlite-vec-rescore.c

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,10 +426,18 @@ static int rescore_knn(vec0_vtab *p, vec0_cursor *pCur,
426426
unsigned char *chunkValidity =
427427
(unsigned char *)sqlite3_column_blob(stmtChunks, 1);
428428
i64 *chunkRowids = (i64 *)sqlite3_column_blob(stmtChunks, 2);
429+
int validityBytes = sqlite3_column_bytes(stmtChunks, 1);
430+
int rowidsBytes = sqlite3_column_bytes(stmtChunks, 2);
429431
if (!chunkValidity || !chunkRowids) {
430432
rc = SQLITE_ERROR;
431433
goto cleanup;
432434
}
435+
// Validate blob sizes match chunk_size expectations
436+
if (validityBytes < (p->chunk_size + 7) / 8 ||
437+
rowidsBytes < p->chunk_size * (int)sizeof(i64)) {
438+
rc = SQLITE_ERROR;
439+
goto cleanup;
440+
}
433441

434442
memset(chunk_distances, 0, p->chunk_size * sizeof(f32));
435443
memset(chunk_topk_idxs, 0, k_oversample * sizeof(i32));

tests/test-rescore.py

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -587,14 +587,37 @@ def test_corrupt_zeroblob_validity(db):
587587
# Corrupt: replace rowids with a truncated blob (wrong size)
588588
db.execute("UPDATE t_chunks SET rowids = x'00'")
589589

590-
# Should not crash — may return wrong results or error
591-
try:
592-
rows = db.execute(
590+
# Should error, not crash — blob size validation catches the mismatch
591+
with pytest.raises(sqlite3.OperationalError):
592+
db.execute(
593593
"SELECT rowid FROM t WHERE embedding MATCH ? ORDER BY distance LIMIT 1",
594594
[float_vec([1, 0, 0, 0, 0, 0, 0, 0])],
595595
).fetchall()
596-
except sqlite3.OperationalError:
597-
pass # Error is acceptable — crash is not
596+
597+
598+
def test_corrupt_truncated_validity_blob(db):
599+
"""KNN should error when rescore chunk validity blob is truncated."""
600+
db.execute(
601+
"CREATE VIRTUAL TABLE t USING vec0("
602+
" embedding float[128] indexed by rescore(quantizer=bit)"
603+
")"
604+
)
605+
for i in range(5):
606+
import random
607+
random.seed(i)
608+
db.execute(
609+
"INSERT INTO t(rowid, embedding) VALUES (?, ?)",
610+
[i + 1, float_vec([random.gauss(0, 1) for _ in range(128)])],
611+
)
612+
613+
# Corrupt: truncate validity blob to 1 byte (should be chunk_size/8 = 128 bytes)
614+
db.execute("UPDATE t_chunks SET validity = x'FF'")
615+
616+
with pytest.raises(sqlite3.OperationalError):
617+
db.execute(
618+
"SELECT rowid FROM t WHERE embedding MATCH ? ORDER BY distance LIMIT 1",
619+
[float_vec([1.0] * 128)],
620+
).fetchall()
598621

599622

600623
def test_rescore_text_pk_insert_knn_delete(db):

0 commit comments

Comments
 (0)