Support limits on index deserialization loops — useful for tests (#4902)

Summary:
Pull Request resolved: #4902

When enabled, throw a FAISS exception when index deserialization loops
driven by read data fields (nlist, nsplits, VT chain length, nhash, etc.)
exceed a configured limit.  This can be set by input validation tests to
prevent wasting time on pathological, OOM inducting, inputs.

Default: 0 (no limit).

Reviewed By: mnorris11

Differential Revision: D96016580

Author: scsiguy

faiss/impl/index_read.cpp

@@ -78,6 +78,32 @@
 
 namespace faiss {
 
+namespace {
+size_t deserialization_loop_limit_ = 0;
+} // namespace
+
+size_t get_deserialization_loop_limit() {
+    return deserialization_loop_limit_;
+}
+
+void set_deserialization_loop_limit(size_t value) {
+    deserialization_loop_limit_ = value;
+}
+
+#define FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(val, field_name) \
+    do {                                                        \
+        auto limit_ = get_deserialization_loop_limit();         \
+        if (limit_ > 0) {                                       \
+            FAISS_THROW_IF_NOT_FMT(                             \
+                    static_cast<size_t>(val) <= limit_,         \
+                    "%s=%zd exceeds deserialization_loop_limit" \
+                    " of %zd",                                  \
+                    field_name,                                 \
+                    static_cast<size_t>(val),                   \
+                    limit_);                                    \
+        }                                                       \
+    } while (0)
+
 /*************************************************************
  * Mmap-ing and viewing facilities
  **************************************************************/
@@ -371,6 +397,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
     } else if (h == fourcc("ilpn") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
         size_t nlist, code_size, n_levels;
         READ1(nlist);
+        FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nlist, "ilpn nlist");
         READ1(code_size);
         READ1(n_levels);
         auto ailp = std::make_unique<ArrayInvertedListsPanorama>(
@@ -400,6 +427,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
     } else if (h == fourcc("ilar") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
         auto ails = std::make_unique<ArrayInvertedLists>(0, 0);
         READ1(ails->nlist);
+        FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ails->nlist, "ilar nlist");
         READ1(ails->code_size);
         ails->ids.resize(ails->nlist);
         ails->codes.resize(ails->nlist);
@@ -430,6 +458,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
         int h2 = (io_flags & 0xffff0000) | (fourcc("il__") & 0x0000ffff);
         size_t nlist, code_size;
         READ1(nlist);
+        FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nlist, "ilar skip nlist");
         READ1(code_size);
         std::vector<size_t> sizes(nlist);
         read_ArrayInvertedLists_sizes(f, sizes);
@@ -557,6 +586,7 @@ static void read_ProductAdditiveQuantizer(
             paq.nsplits > 0,
             "invalid ProductAdditiveQuantizer nsplits %zd (must be > 0)",
             paq.nsplits);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(paq.nsplits, "nsplits");
 }
 
 static void read_ProductResidualQuantizer(
@@ -735,6 +765,8 @@ static void read_HNSW(HNSW& hnsw, IOReader* f) {
 static void read_NSG(NSG& nsg, IOReader* f) {
     READ1(nsg.ntotal);
     READ1(nsg.R);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nsg.ntotal, "nsg.ntotal");
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nsg.R, "nsg.R");
     FAISS_THROW_IF_NOT_FMT(nsg.R > 0, "invalid NSG R %d (must be > 0)", nsg.R);
     READ1(nsg.L);
     READ1(nsg.C);
@@ -858,6 +890,7 @@ void read_ivf_header(
         std::vector<std::vector<idx_t>>* ids) {
     read_index_header(*ivf, f);
     READ1(ivf->nlist);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ivf->nlist, "nlist");
     READ1(ivf->nprobe);
     ivf->quantizer = read_index(f);
     ivf->own_fields = true;
@@ -1368,6 +1401,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
                 nt >= 0,
                 "invalid VectorTransform chain length %d (must be >= 0)",
                 nt);
+        FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(
+                nt, "VectorTransform chain length");
         for (int i = 0; i < nt; i++) {
             ixpt->chain.push_back(read_VectorTransform(f));
         }
@@ -1898,6 +1933,7 @@ static void read_binary_ivf_header(
         std::vector<std::vector<idx_t>>* ids = nullptr) {
     read_index_binary_header(ivf, f);
     READ1(ivf.nlist);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ivf.nlist, "nlist");
     READ1(ivf.nprobe);
     ivf.quantizer = read_index_binary(f);
     ivf.own_fields = true;
@@ -1915,6 +1951,7 @@ static void read_binary_hash_invlists(
         IOReader* f) {
     size_t sz;
     READ1(sz);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(sz, "binary hash invlists sz");
     int il_nbit = 0;
     READ1(il_nbit);
     FAISS_THROW_IF_NOT_FMT(
@@ -1965,6 +2002,7 @@ static void read_binary_multi_hash_map(
     size_t sz;
     READ1(id_bits);
     READ1(sz);
+    FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(sz, "multi hash map sz");
     std::vector<uint8_t> buf;
     READVECTOR(buf);
     size_t nbit = add_no_overflow(
@@ -2063,6 +2101,7 @@ std::unique_ptr<IndexBinary> read_index_binary_up(IOReader* f, int io_flags) {
                 idxmh->nhash > 0,
                 "invalid IndexBinaryMultiHash nhash %d (must be > 0)",
                 idxmh->nhash);
+        FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(idxmh->nhash, "nhash");
         READ1(idxmh->nflip);
         idxmh->maps.resize(idxmh->nhash);
         for (int i = 0; i < idxmh->nhash; i++) {

faiss/index_io.h

@@ -113,6 +113,17 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
         IOReader* reader,
         int io_flags = 0);
 
+// Returns the current deserialization loop limit.
+// When nonzero, deserialization rejects loop-driving fields (nlist,
+// nsplits, VT chain length, nhash, etc.) that exceed this value.
+// Default: 0 (no limit).
+size_t get_deserialization_loop_limit();
+
+// Sets the deserialization loop limit.
+// NOT thread-safe: set before any concurrent deserialization calls
+// and do not modify while deserialization is in progress on other threads.
+void set_deserialization_loop_limit(size_t value);
+
 } // namespace faiss
 
 #endif