From 408786bec9d016f2b9893373a2ad40b82868eeb4 Mon Sep 17 00:00:00 2001 From: Justin Gibbs Date: Fri, 3 Apr 2026 09:58:37 -0700 Subject: [PATCH] Validate AdditiveQuantizer dimension consistency during deserialization (#5023) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5023 Add deserialization-time validation to catch AdditiveQuantizer dimension mismatches that would cause crashes during search operations. 1. **AQ dimension vs index dimension cross-check**: Added `validate_aq_dimension_match()` helper that verifies `aq.d == idx.d` for all AdditiveQuantizer-based index types. Without this check, a corrupted index where AQ.d != Index.d causes `compute_LUT()` to read codebooks out of bounds (query vectors are sized for Index.d while codebooks are laid out for AQ.d). 2. **Coverage across all AQ index types**: - IndexResidualQuantizer (IxRq/IxRQ) - IndexLocalSearchQuantizer (IxLS) - IndexProductResidualQuantizer (IxPR) - IndexProductLocalSearchQuantizer (IxPL) - ResidualCoarseQuantizer (ImRQ) - IndexAdditiveQuantizerFastScan (IRfs/ILfs/IPRf/IPLf) - IndexIVFAdditiveQuantizerFastScan (IVLf/IVRf/NPLf/NPRf) - IndexIVFAdditiveQuantizer (IwLS/IwRQ/IwPL/IwPR) Reviewed By: mnorris11 Differential Revision: D99137777 fbshipit-source-id: bb4201f32b1268d66907f9a3f34eff869c80dc9c --- faiss/impl/index_read.cpp | 57 +++++++++++++++------------ tests/test_read_index_deserialize.cpp | 52 ++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 25 deletions(-) diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index d2682aec4..85d16e832 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -697,6 +697,22 @@ static void validate_codebooks_size( aq.total_codebook_size); } +// Validate that the AdditiveQuantizer dimension matches the index header +// dimension. compute_LUT() treats codebooks as a (d, total_codebook_size) +// matrix and query vectors are sized for idx_d, so a mismatch leads to +// out-of-bounds reads. +static void validate_aq_dimension_match( + const AdditiveQuantizer& aq, + int idx_d, + const char* index_type) { + FAISS_THROW_IF_NOT_FMT( + aq.d == static_cast(idx_d), + "%s: AdditiveQuantizer d=%zd does not match index d=%d", + index_type, + aq.d, + idx_d); +} + static void read_ResidualQuantizer( ResidualQuantizer& rq, IOReader* f, @@ -1359,11 +1375,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { } else { read_ResidualQuantizer(idxr->rq, f, io_flags); } - FAISS_THROW_IF_NOT_FMT( - idxr->rq.d == static_cast(idxr->d), - "IndexResidualQuantizer d mismatch: rq.d=%zd vs idx.d=%d", - idxr->rq.d, - idxr->d); + validate_aq_dimension_match( + idxr->rq, idxr->d, "IndexResidualQuantizer"); READ1(idxr->code_size); read_vector(idxr->codes, f); FAISS_THROW_IF_NOT( @@ -1373,11 +1386,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { auto idxr = std::make_unique(); read_index_header(*idxr, f); read_LocalSearchQuantizer(idxr->lsq, f); - FAISS_THROW_IF_NOT_FMT( - idxr->lsq.d == static_cast(idxr->d), - "IndexLocalSearchQuantizer d mismatch: lsq.d=%zd vs idx.d=%d", - idxr->lsq.d, - idxr->d); + validate_aq_dimension_match( + idxr->lsq, idxr->d, "IndexLocalSearchQuantizer"); READ1(idxr->code_size); read_vector(idxr->codes, f); FAISS_THROW_IF_NOT( @@ -1387,11 +1397,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { auto idxpr = std::make_unique(); read_index_header(*idxpr, f); read_ProductResidualQuantizer(idxpr->prq, f, io_flags); - FAISS_THROW_IF_NOT_FMT( - idxpr->prq.d == static_cast(idxpr->d), - "IndexProductResidualQuantizer d mismatch: prq.d=%zd vs idx.d=%d", - idxpr->prq.d, - idxpr->d); + validate_aq_dimension_match( + idxpr->prq, idxpr->d, "IndexProductResidualQuantizer"); READ1(idxpr->code_size); read_vector(idxpr->codes, f); FAISS_THROW_IF_NOT( @@ -1401,11 +1408,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { auto idxpl = std::make_unique(); read_index_header(*idxpl, f); read_ProductLocalSearchQuantizer(idxpl->plsq, f); - FAISS_THROW_IF_NOT_FMT( - idxpl->plsq.d == static_cast(idxpl->d), - "IndexProductLocalSearchQuantizer d mismatch: plsq.d=%zd vs idx.d=%d", - idxpl->plsq.d, - idxpl->d); + validate_aq_dimension_match( + idxpl->plsq, idxpl->d, "IndexProductLocalSearchQuantizer"); READ1(idxpl->code_size); read_vector(idxpl->codes, f); FAISS_THROW_IF_NOT( @@ -1415,11 +1419,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { auto idxr = std::make_unique(); read_index_header(*idxr, f); read_ResidualQuantizer(idxr->rq, f, io_flags); - FAISS_THROW_IF_NOT_FMT( - idxr->rq.d == static_cast(idxr->d), - "ResidualCoarseQuantizer d mismatch: rq.d=%zd vs idx.d=%d", - idxr->rq.d, - idxr->d); + validate_aq_dimension_match( + idxr->rq, idxr->d, "ResidualCoarseQuantizer"); READ1(idxr->beam_factor); if (io_flags & IO_FLAG_SKIP_PRECOMPUTE_TABLE) { // then we force the beam factor to -1 @@ -1490,6 +1491,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductResidualQuantizer( *(ProductResidualQuantizer*)idxaqfs->aq, f, io_flags); } + validate_aq_dimension_match( + *idxaqfs->aq, idxaqfs->d, "IndexAdditiveQuantizerFastScan"); READ1(idxaqfs->implem); READ1(idxaqfs->bbs); @@ -1563,6 +1566,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductResidualQuantizer( *(ProductResidualQuantizer*)ivaqfs->aq, f, io_flags); } + validate_aq_dimension_match( + *ivaqfs->aq, ivaqfs->d, "IndexIVFAdditiveQuantizerFastScan"); READ1(ivaqfs->by_residual); READ1(ivaqfs->implem); @@ -1720,6 +1725,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductResidualQuantizer( *(ProductResidualQuantizer*)iva->aq, f, io_flags); } + validate_aq_dimension_match( + *iva->aq, iva->d, "IndexIVFAdditiveQuantizer"); READ1(iva->by_residual); READ1(iva->use_precomputed_table); read_InvertedLists(*iva, f, io_flags); diff --git a/tests/test_read_index_deserialize.cpp b/tests/test_read_index_deserialize.cpp index 64677d14a..c961fa0db 100644 --- a/tests/test_read_index_deserialize.cpp +++ b/tests/test_read_index_deserialize.cpp @@ -2868,3 +2868,55 @@ TEST(ReadIndexDeserialize, SQUntrainedEmptyTrainedAccepted) { auto idx = read_index_up(&reader); EXPECT_FALSE(idx->is_trained); } + +// Test: IndexResidualQuantizer rejects AQ dimension != index dimension. +TEST(ReadIndexDeserialize, IndexResidualQuantizerAQDimensionMismatch) { + std::vector buf; + push_fourcc(buf, "IxRq"); + // Index header says d=4, but AQ will have d=8 → mismatch. + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4}); + push_val(buf, 1); // code_size + push_vector(buf, {}); // codes + + expect_read_throws_with(buf, "does not match index d"); +} + +// Test: IndexLocalSearchQuantizer rejects AQ dimension != index dimension. +TEST(ReadIndexDeserialize, IndexLocalSearchQuantizerAQDimensionMismatch) { + std::vector buf; + push_fourcc(buf, "IxLS"); + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + push_local_search_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4}); + push_val(buf, 1); // code_size + push_vector(buf, {}); // codes + + expect_read_throws_with(buf, "does not match index d"); +} + +// Test: ResidualCoarseQuantizer rejects AQ dimension != index dimension. +TEST(ReadIndexDeserialize, ResidualCoarseQuantizerAQDimensionMismatch) { + std::vector buf; + push_fourcc(buf, "ImRQ"); + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4}); + push_val(buf, -1.0f); // beam_factor + + expect_read_throws_with(buf, "does not match index d"); +} + +// Test: IndexResidualQuantizerFastScan rejects AQ dimension != index dimension. +TEST(ReadIndexDeserialize, IndexRQFastScanAQDimensionMismatch) { + std::vector buf; + push_fourcc(buf, "IRfs"); + // Index header says d=4, but AQ will have d=8 → mismatch. + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4}); + // FastScan fields won't be reached due to early validation. + // But include a few just in case: + push_val(buf, 0); // implem + push_val(buf, 32); // bbs + push_val(buf, 0); // qbs + + expect_read_throws_with(buf, "does not match index d"); +}