Validate IndexIDMap id_map and inner-index ntotal consistency to prevent search-time null deref (#5203)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5203

IndexIDMap::search_ex translates each label returned by the inner index through this->id_map.  The outer IDMap header and the inner index header carry independent ntotal fields on disk, and the existing read_index check only enforces id_map.size() == idxmap->ntotal — the inner index's ntotal is not cross-validated.  When the two disagree, the translation loop can index past the end of id_map.  With id_map empty its data() pointer can be null, so the offending read dereferences a null pointer and aborts under ASAN; without ASAN the access is a heap-buffer overflow.  The same gap exists for the binary counterpart at the IBMp/IBM2 read site.

Two layers of protection:

1. read_index now validates idxmap->index->ntotal == idxmap->ntotal in addition to the existing id_map.size() == ntotal check, for both IxMp/IxM2 and IBMp/IBM2.  The new check rejects corrupt or maliciously constructed serialized payloads that pass each per-index validator individually but disagree at the IDMap layer.

2. IndexIDMap::search_ex and IndexIDMap::range_search bound-check each translated label against id_map.size() and remap out-of-range labels to the -1 "no result" sentinel rather than indexing past the end of id_map.  This is needed because the same invariant can be transiently violated by ordinary in-memory code paths outside of deserialization — for example mid-construction (vectors added directly to the inner index before the IDMap layer has caught up), or after a recoverable allocation failure in add_with_ids_ex, add_sa_codes, or merge_from leaves the IDMap behind the inner index.  The added comparison sits adjacent to the existing label sign check on the search hot path and adds a single branch.

Reviewed By: mnorris11

Differential Revision: D104468890

fbshipit-source-id: d7323adb62b9c5c18173fd43f029a8adbe4ba986

faiss/IndexIDMap.cpp

@@ -196,9 +196,10 @@ void IndexIDMapTemplate<IndexT>::search_ex(
     }
     index->search_ex(n, x, numeric_type, k, distances, labels, params);
     idx_t* li = labels;
+    const idx_t id_map_size = static_cast<idx_t>(id_map.size());
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; i++) {
-        li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
+        li[i] = (li[i] < 0 || li[i] >= id_map_size) ? idx_t(-1) : id_map[li[i]];
     }
 }
 
@@ -237,10 +238,12 @@ void IndexIDMapTemplate<IndexT>::range_search(
         index->range_search(n, x, radius, result);
     }
 
+    const idx_t id_map_size = static_cast<idx_t>(id_map.size());
 #pragma omp parallel for
     for (idx_t i = 0; i < static_cast<idx_t>(result->lims[result->nq]); i++) {
-        result->labels[i] = result->labels[i] < 0 ? result->labels[i]
-                                                  : id_map[result->labels[i]];
+        const idx_t label = result->labels[i];
+        result->labels[i] =
+                (label < 0 || label >= id_map_size) ? idx_t(-1) : id_map[label];
     }
 }
 

faiss/impl/index_read.cpp

@@ -2127,6 +2127,12 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
                 ") does not match ntotal (%" PRId64 ")",
                 int64_t(idxmap->id_map.size()),
                 idxmap->ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                idxmap->index->ntotal == idxmap->ntotal,
+                "IndexIDMap inner index ntotal (%" PRId64
+                ") does not match IndexIDMap ntotal (%" PRId64 ")",
+                idxmap->index->ntotal,
+                idxmap->ntotal);
         if (is_map2) {
             static_cast<IndexIDMap2*>(idxmap.get())->construct_rev_map();
         }
@@ -3022,7 +3028,18 @@ std::unique_ptr<IndexBinary> read_index_binary_up(IOReader* f, int io_flags) {
         idxmap->index = read_index_binary(f, io_flags);
         idxmap->own_fields = true;
         READVECTOR(idxmap->id_map);
-        FAISS_THROW_IF_NOT(idxmap->id_map.size() == idxmap->ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                idxmap->id_map.size() == idxmap->ntotal,
+                "IndexBinaryIDMap id_map size (%" PRId64
+                ") does not match ntotal (%" PRId64 ")",
+                int64_t(idxmap->id_map.size()),
+                idxmap->ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                idxmap->index->ntotal == idxmap->ntotal,
+                "IndexBinaryIDMap inner index ntotal (%" PRId64
+                ") does not match IndexBinaryIDMap ntotal (%" PRId64 ")",
+                idxmap->index->ntotal,
+                idxmap->ntotal);
         if (is_map2) {
             static_cast<IndexBinaryIDMap2*>(idxmap.get())->construct_rev_map();
         }

tests/test_read_index_deserialize.cpp

@@ -22,6 +22,7 @@
 #include <faiss/IndexBinaryIVF.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
+#include <faiss/IndexIDMap.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
 #include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
@@ -4153,3 +4154,127 @@ TEST(ReadIndexDeserialize, SVSVamanaFourccRejected) {
 }
 
 #endif // FAISS_ENABLE_SVS
+
+// -----------------------------------------------------------------------
+// IndexIDMap deserialization invariant: id_map.size(), idxmap.ntotal,
+// and the inner index's ntotal must all agree. The outer IDMap header
+// and the inner index header carry independent ntotal fields on disk,
+// so deserialization must reject payloads where they disagree.
+// -----------------------------------------------------------------------
+
+namespace {
+
+/// Build a serialized IxMp/IxM2 payload. The inner IxF2 (IndexFlat) carries
+/// codes sized to match its own ntotal, so it passes the inner deserializer's
+/// codes-vs-ntotal check and the IDMap-vs-inner-ntotal check is exercised
+/// in isolation.
+std::vector<uint8_t> make_idmap_payload(
+        const char fourcc[4],
+        int64_t outer_ntotal,
+        int64_t inner_ntotal,
+        size_t id_map_size) {
+    constexpr int d = 4;
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, fourcc);
+    push_index_header(buf, d, outer_ntotal);
+    push_fourcc(buf, "IxF2");
+    push_index_header(buf, d, inner_ntotal);
+    // IndexFlat codes are written via WRITEXBVECTOR: a size_t prefix of
+    // (byte_count / 4) followed by byte_count bytes. For inner_ntotal
+    // float vectors of dimension d, byte_count = inner_ntotal * d * 4.
+    const size_t byte_count =
+            static_cast<size_t>(inner_ntotal) * d * sizeof(float);
+    push_val<size_t>(buf, byte_count / 4);
+    buf.insert(buf.end(), byte_count, uint8_t{0});
+    std::vector<int64_t> id_map(id_map_size, 0);
+    push_vector<int64_t>(buf, id_map);
+    return buf;
+}
+
+std::vector<uint8_t> make_binary_idmap_payload(
+        const char fourcc[4],
+        int64_t outer_ntotal,
+        int64_t inner_ntotal,
+        size_t id_map_size) {
+    constexpr int d = 8;
+    constexpr int code_size = d / 8;
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, fourcc);
+    push_binary_index_header(buf, d, outer_ntotal);
+    push_fourcc(buf, "IBxF");
+    push_binary_index_header(buf, d, inner_ntotal);
+    std::vector<uint8_t> codes(inner_ntotal * code_size, 0);
+    push_vector<uint8_t>(buf, codes);
+    std::vector<int64_t> id_map(id_map_size, 0);
+    push_vector<int64_t>(buf, id_map);
+    return buf;
+}
+
+} // namespace
+
+TEST(ReadIndexDeserialize, IndexIDMapInnerNtotalMismatchRejected) {
+    // Outer IDMap ntotal=0 agrees with id_map.size()=0, but the inner
+    // IndexFlat reports ntotal=3.
+    auto buf = make_idmap_payload(
+            "IxMp", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0);
+    expect_read_throws_with(buf, "inner index ntotal");
+}
+
+TEST(ReadIndexDeserialize, IndexIDMap2InnerNtotalMismatchRejected) {
+    auto buf = make_idmap_payload(
+            "IxM2", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0);
+    expect_read_throws_with(buf, "inner index ntotal");
+}
+
+TEST(ReadIndexDeserialize, IndexBinaryIDMapInnerNtotalMismatchRejected) {
+    auto buf = make_binary_idmap_payload(
+            "IBMp", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0);
+    expect_binary_read_throws_with(buf, "inner index ntotal");
+}
+
+TEST(ReadIndexDeserialize, IndexBinaryIDMap2InnerNtotalMismatchRejected) {
+    auto buf = make_binary_idmap_payload(
+            "IBM2", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0);
+    expect_binary_read_throws_with(buf, "inner index ntotal");
+}
+
+// A well-formed IxMp payload with matching ntotals deserializes cleanly.
+TEST(ReadIndexDeserialize, IndexIDMapMatchingNtotalsAccepted) {
+    auto buf = make_idmap_payload(
+            "IxMp", /*outer_ntotal=*/0, /*inner_ntotal=*/0, /*id_map_size=*/0);
+    VectorIOReader reader;
+    reader.data = buf;
+    EXPECT_NO_THROW({
+        auto idx = read_index_up(&reader);
+        EXPECT_NE(idx, nullptr);
+    });
+}
+
+// -----------------------------------------------------------------------
+// IndexIDMap::search_ex maps each inner-index label through id_map. When
+// inner.ntotal exceeds id_map.size() (e.g. mid-construction, or after a
+// recoverable failure in add_with_ids_ex / add_sa_codes / merge_from
+// leaves the IDMap layer behind the inner index), inner labels can land
+// outside id_map. Such labels must be remapped to the -1 "no result"
+// sentinel.
+// -----------------------------------------------------------------------
+TEST(IndexIDMapSafety, SearchRemapsOutOfRangeLabelsToSentinel) {
+    constexpr int d = 4;
+    constexpr int nb = 3;
+    std::vector<float> xb(nb * d, 0.0f);
+    for (int i = 0; i < nb; i++) {
+        xb[i * d] = static_cast<float>(i);
+    }
+
+    IndexFlat inner(d, METRIC_L2);
+    IndexIDMap idmap(&inner);
+    // inner.ntotal == nb, idmap.id_map is empty.
+    inner.add(nb, xb.data());
+
+    std::vector<float> xq(d, 0.0f);
+    std::vector<float> distances(1);
+    std::vector<faiss::idx_t> labels(1, 0);
+    EXPECT_NO_THROW(
+            idmap.search(1, xq.data(), 1, distances.data(), labels.data()));
+    EXPECT_EQ(labels[0], -1);
+}