From 6707eacfaffc012bdb2ac95eae37b4baceef11f0 Mon Sep 17 00:00:00 2001 From: Justin Gibbs Date: Sun, 19 Apr 2026 08:54:55 -0700 Subject: [PATCH] Add is_trained check to IndexIVF search and range_search to prevent querying untrained indexes (#5114) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5114 Add FAISS_THROW_IF_NOT(is_trained) to IndexIVF::search(), IndexIVF::search_preassigned(), IndexIVF::range_search(), and IndexIVF::range_search_preassigned(), mirroring the existing check in IndexScalarQuantizer::search(). This prevents querying untrained IVF indexes deserialized from corrupt data where the ScalarQuantizer trained vector is empty. The existing deserialization validation in read_ScalarQuantizer correctly allows untrained indexes (is_trained=false with empty trained) to be deserialized, since these are legitimately produced by index_factory before training. However, IndexIVF search methods lacked the is_trained guard that IndexScalarQuantizer::search() has, allowing a deserialized untrained IndexIVFScalarQuantizer to be queried, which causes null-deref in QuantizerTemplate when it indexes into the empty trained vector. Reviewed By: mnorris11 Differential Revision: D101243973 fbshipit-source-id: eca68dc82e5cca37d4c461b735c5d59a66349248 --- faiss/IndexIVF.cpp | 4 + tests/test_read_index_deserialize.cpp | 103 ++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index fabac7753..ad401a66b 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -312,6 +312,7 @@ void IndexIVF::search( const SearchParameters* params_in) const { FAISS_THROW_IF_NOT(k > 0); FAISS_THROW_IF_NOT_MSG(quantizer, "IVF quantizer must not be null"); + FAISS_THROW_IF_NOT_MSG(is_trained, "IVF index is not trained"); FAISS_THROW_IF_NOT_MSG(invlists, "IVF index has no inverted lists"); const IVFSearchParameters* params = nullptr; if (params_in) { @@ -409,6 +410,7 @@ void IndexIVF::search_preassigned( const IVFSearchParameters* params, IndexIVFStats* ivf_stats) const { FAISS_THROW_IF_NOT(k > 0); + FAISS_THROW_IF_NOT_MSG(is_trained, "IVF index is not trained"); FAISS_THROW_IF_NOT_MSG(invlists, "IVF index has no inverted lists"); idx_t cur_nprobe = params ? params->nprobe : this->nprobe; @@ -738,6 +740,7 @@ void IndexIVF::range_search( RangeSearchResult* result, const SearchParameters* params_in) const { FAISS_THROW_IF_NOT_MSG(quantizer, "IVF quantizer must not be null"); + FAISS_THROW_IF_NOT_MSG(is_trained, "IVF index is not trained"); const IVFSearchParameters* params = nullptr; const SearchParameters* quantizer_params = nullptr; if (params_in) { @@ -782,6 +785,7 @@ void IndexIVF::range_search_preassigned( bool store_pairs, const IVFSearchParameters* params, IndexIVFStats* stats) const { + FAISS_THROW_IF_NOT_MSG(is_trained, "IVF index is not trained"); idx_t cur_nprobe = params ? params->nprobe : this->nprobe; cur_nprobe = std::min((idx_t)nlist, cur_nprobe); FAISS_THROW_IF_NOT(cur_nprobe > 0); diff --git a/tests/test_read_index_deserialize.cpp b/tests/test_read_index_deserialize.cpp index ea7a65838..8c3feb171 100644 --- a/tests/test_read_index_deserialize.cpp +++ b/tests/test_read_index_deserialize.cpp @@ -1680,6 +1680,109 @@ TEST(ReadIndexDeserialize, IVFQuantizerUntrained) { EXPECT_NO_THROW(read_index_up(&reader)); } +// ----------------------------------------------------------------------- +// Test: IndexIVFScalarQuantizer with empty trained vector and +// is_trained=false deserializes successfully (legitimate untrained index), +// but searching it throws because IndexIVF::search checks is_trained. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IVFScalarQuantizerUntrainedSearchRejected) { + std::vector buf; + push_fourcc(buf, "IwSq"); + // IVF header: index_header + nlist + nprobe + quantizer + direct_map + push_index_header(buf, /*d=*/4, /*ntotal=*/0, /*is_trained=*/false); + push_val(buf, 1); // nlist + push_val(buf, 1); // nprobe + push_minimal_flat(buf, /*d=*/4); + push_empty_direct_map(buf); + // ScalarQuantizer fields: + push_val(buf, 0); // qtype = QT_8bit + push_val(buf, 0); // rangestat + push_val(buf, 0.0f); // rangestat_arg + push_val(buf, 4); // d + push_val(buf, 4); // code_size + push_vector(buf, {}); // trained (empty — untrained) + // IwSq additional fields: + push_val(buf, 4); // code_size + push_val(buf, false); // by_residual + push_null_invlists(buf); + + // Deserialization should succeed — untrained indexes are legitimate + VectorIOReader reader; + reader.data = buf; + auto idx = read_index_up(&reader); + ASSERT_NE(idx, nullptr); + EXPECT_FALSE(idx->is_trained); + + // search should throw — is_trained check in IndexIVF::search + std::vector xq(4, 0.0f); + std::vector distances(1); + std::vector labels(1); + EXPECT_THROW( + idx->search(1, xq.data(), 1, distances.data(), labels.data()), + FaissException); + + // range_search should throw — is_trained check in IndexIVF::range_search + RangeSearchResult rsr(1); + EXPECT_THROW(idx->range_search(1, xq.data(), 1.0f, &rsr), FaissException); + + // search_preassigned should throw directly + auto* ivf = dynamic_cast(idx.get()); + ASSERT_NE(ivf, nullptr); + idx_t key = 0; + float coarse_dis = 0.0f; + EXPECT_THROW( + ivf->search_preassigned( + 1, + xq.data(), + 1, + &key, + &coarse_dis, + distances.data(), + labels.data(), + false, + nullptr, + nullptr), + FaissException); + + // range_search_preassigned should throw directly + RangeSearchResult rsr2(1); + EXPECT_THROW( + ivf->range_search_preassigned( + 1, + xq.data(), + 1.0f, + &key, + &coarse_dis, + &rsr2, + false, + nullptr, + nullptr), + FaissException); +} + +// ----------------------------------------------------------------------- +// Test: IndexIVFScalarQuantizer with is_trained=true but empty trained +// is rejected at deserialization time — corrupt data. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IVFScalarQuantizerTrainedEmptyTrained) { + std::vector buf; + push_fourcc(buf, "IwSq"); + push_index_header(buf, /*d=*/4, /*ntotal=*/0, /*is_trained=*/true); + push_val(buf, 1); // nlist + push_val(buf, 1); // nprobe + push_minimal_flat(buf, /*d=*/4); + push_empty_direct_map(buf); + // ScalarQuantizer fields: + push_val(buf, 0); // qtype = QT_8bit + push_val(buf, 0); // rangestat + push_val(buf, 0.0f); // rangestat_arg + push_val(buf, 4); // d + push_val(buf, 4); // code_size + push_vector(buf, {}); // trained (empty — but is_trained=true!) + + expect_read_throws_with(buf, "ScalarQuantizer trained size"); +} + // ----------------------------------------------------------------------- // Test: initialize_IVFPQ_precomputed_table rejects a null quantizer. // Protects against null-deref from corrupt serialized data where the