Add is_trained check to IndexIVF search and range_search to prevent querying untrained indexes (#5114)

Summary:
Pull Request resolved: #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

Author: scsiguy

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);

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<uint8_t> 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<size_t>(buf, 1); // nlist
+    push_val<size_t>(buf, 1); // nprobe
+    push_minimal_flat(buf, /*d=*/4);
+    push_empty_direct_map(buf);
+    // ScalarQuantizer fields:
+    push_val<int>(buf, 0);       // qtype = QT_8bit
+    push_val<int>(buf, 0);       // rangestat
+    push_val<float>(buf, 0.0f);  // rangestat_arg
+    push_val<size_t>(buf, 4);    // d
+    push_val<size_t>(buf, 4);    // code_size
+    push_vector<float>(buf, {}); // trained (empty — untrained)
+    // IwSq additional fields:
+    push_val<size_t>(buf, 4);   // code_size
+    push_val<bool>(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<float> xq(4, 0.0f);
+    std::vector<float> distances(1);
+    std::vector<idx_t> 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<IndexIVF*>(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<uint8_t> buf;
+    push_fourcc(buf, "IwSq");
+    push_index_header(buf, /*d=*/4, /*ntotal=*/0, /*is_trained=*/true);
+    push_val<size_t>(buf, 1); // nlist
+    push_val<size_t>(buf, 1); // nprobe
+    push_minimal_flat(buf, /*d=*/4);
+    push_empty_direct_map(buf);
+    // ScalarQuantizer fields:
+    push_val<int>(buf, 0);       // qtype = QT_8bit
+    push_val<int>(buf, 0);       // rangestat
+    push_val<float>(buf, 0.0f);  // rangestat_arg
+    push_val<size_t>(buf, 4);    // d
+    push_val<size_t>(buf, 4);    // code_size
+    push_vector<float>(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