Validate ProductQuantizer M*ksub during deserialization to prevent oversized allocations (#5187)

Summary:
Pull Request resolved: #5187

ProductQuantizer::set_derived_values only checks d % M == 0, which is satisfied by any M when d == 0. A corrupt or maliciously constructed serialized index can therefore carry an enormous M alongside an empty centroids vector and pass read_ProductQuantizer's existing centroids validation. Downstream allocations sized M * ksub — most notably the r_norms vector in initialize_IVFPQ_precomputed_table when an IVFPQ is precomputed at deserialization time — then exceed std::vector::max_size() and raise std::length_error. In code paths invoked from noexcept callers this aborts the process via std::terminate.

Two layers of protection:

1. read_ProductQuantizer now bounds M * ksub by the configurable deserialization vector byte limit, mirroring the existing d * ksub check that protects the centroids vector. This caps every PQ-derived M * ksub allocation (residual norms, search-time distance tables) at the same byte budget the caller has already chosen for vector deserialization, and rejects pathological M values up front with a FaissException rather than letting them propagate to allocation sites.

2. initialize_IVFPQ_precomputed_table now uses mul_no_overflow when computing table_size = pq.M * pq.ksub * nlist * sizeof(float). The existing precomputed_table_max_bytes guard depended on this product not wrapping size_t; with raw multiplication a sufficiently large M, ksub, or nlist silently wrapped to a small value, bypassed the guard, and proceeded to the std::vector::max_size() failure described above. Switching to mul_no_overflow makes the guard work correctly for any input that reaches it, and also defends the in-memory IVFPQ construction paths (IndexIVFPQ::train, IndexHNSW post-init, IndexIVFPQFastScan) without changing their behavior on well-formed inputs.

Reviewed By: mnorris11

Differential Revision: D104129118

fbshipit-source-id: c273aecd22572d734f535df98127dbb7a437ac4b

Author: scsiguy

faiss/IndexIVFPQ.cpp

@@ -404,6 +404,9 @@ void initialize_IVFPQ_precomputed_table(
         return;
     }
 
+    const size_t m_ksub =
+            mul_no_overflow(pq.M, pq.ksub, "IVFPQ precomputed_table");
+
     if (use_precomputed_table == 0) { // then choose the type of table
         if (!(quantizer->metric_type == METRIC_L2 && by_residual)) {
             if (verbose) {
@@ -418,7 +421,10 @@ void initialize_IVFPQ_precomputed_table(
         if (miq && pq.M % miq->pq.M == 0) {
             use_precomputed_table = 2;
         } else {
-            size_t table_size = pq.M * pq.ksub * nlist * sizeof(float);
+            size_t table_size = mul_no_overflow(
+                    mul_no_overflow(m_ksub, nlist, "IVFPQ precomputed_table"),
+                    sizeof(float),
+                    "IVFPQ precomputed_table");
             if (table_size > precomputed_table_max_bytes) {
                 if (verbose) {
                     printf("IndexIVFPQ::precompute_table: not precomputing table, "
@@ -438,7 +444,7 @@ void initialize_IVFPQ_precomputed_table(
     }
 
     // squared norms of the PQ centroids
-    std::vector<float> r_norms(pq.M * pq.ksub, NAN);
+    std::vector<float> r_norms(m_ksub, NAN);
     for (size_t m = 0; m < pq.M; m++) {
         for (size_t j = 0; j < pq.ksub; j++) {
             r_norms[m * pq.ksub + j] =
@@ -447,15 +453,16 @@ void initialize_IVFPQ_precomputed_table(
     }
 
     if (use_precomputed_table == 1) {
-        precomputed_table.resize(nlist * pq.M * pq.ksub);
+        precomputed_table.resize(
+                mul_no_overflow(nlist, m_ksub, "IVFPQ precomputed_table"));
         std::vector<float> centroid(d);
 
         for (size_t i = 0; i < nlist; i++) {
             quantizer->reconstruct(i, centroid.data());
 
-            float* tab = &precomputed_table[i * pq.M * pq.ksub];
+            float* tab = &precomputed_table[i * m_ksub];
             pq.compute_inner_prod_table(centroid.data(), tab);
-            fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd_dispatch(m_ksub, r_norms.data(), 2.0, tab, tab);
         }
     } else if (use_precomputed_table == 2) {
         const MultiIndexQuantizer* miq =
@@ -464,7 +471,8 @@ void initialize_IVFPQ_precomputed_table(
         const ProductQuantizer& cpq = miq->pq;
         FAISS_THROW_IF_NOT(pq.M % cpq.M == 0);
 
-        precomputed_table.resize(cpq.ksub * pq.M * pq.ksub);
+        precomputed_table.resize(
+                mul_no_overflow(cpq.ksub, m_ksub, "IVFPQ precomputed_table"));
 
         // reorder PQ centroid table
         std::vector<float> centroids(d * cpq.ksub, NAN);
@@ -481,8 +489,8 @@ void initialize_IVFPQ_precomputed_table(
                 cpq.ksub, centroids.data(), precomputed_table.data());
 
         for (size_t i = 0; i < cpq.ksub; i++) {
-            float* tab = &precomputed_table[i * pq.M * pq.ksub];
-            fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            float* tab = &precomputed_table[i * m_ksub];
+            fvec_madd_dispatch(m_ksub, r_norms.data(), 2.0, tab, tab);
         }
     }
 }

faiss/impl/index_read.cpp

@@ -712,6 +712,13 @@ void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
         FAISS_THROW_IF_NOT_MSG(
                 n < get_deserialization_vector_byte_limit() / sizeof(float),
                 "PQ centroids allocation would exceed deserialization byte limit");
+        // Per-subquantizer tables (e.g. IVFPQ residual norms, search-time
+        // distance tables) are sized M * ksub.
+        size_t m_ksub = mul_no_overflow(pq->M, ksub, "PQ M*ksub");
+        FAISS_THROW_IF_NOT_MSG(
+                m_ksub <
+                        get_deserialization_vector_byte_limit() / sizeof(float),
+                "PQ M*ksub allocation would exceed deserialization byte limit");
     }
     pq->set_derived_values();
     READVECTOR(pq->centroids);

tests/test_read_index_deserialize.cpp

@@ -3798,6 +3798,107 @@ TEST(ReadIndexDeserialize, IndexLatticeCodeSizeTooLarge) {
     expect_read_throws_with(buf, "code_size");
 }
 
+// -----------------------------------------------------------------------
+// Test: ProductQuantizer with M*ksub exceeding the deserialization byte
+// limit is rejected. ProductQuantizer::set_derived_values only validates
+// d % M == 0, which is satisfied by any M when d == 0, so a crafted PQ
+// header could carry an enormous M alongside an empty centroids vector.
+// Without the read-side bound, downstream allocations sized M*ksub (e.g.
+// IVFPQ residual norms in initialize_IVFPQ_precomputed_table) reach
+// std::vector::max_size() and abort the process via std::length_error.
+// -----------------------------------------------------------------------
+TEST(ReadIndexDeserialize, ProductQuantizerMKsubExceedsByteLimit) {
+    const size_t old_limit = get_deserialization_vector_byte_limit();
+    set_deserialization_vector_byte_limit(1 << 20); // 1 MB
+
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IxPq");
+    push_index_header(buf, /*d=*/0, /*ntotal=*/0);
+    // PQ: d=0 (so 0 % M == 0 holds), M=1<<20, nbits=8 -> ksub=256
+    // M * ksub = 2^28 floats = 1 GB, well above the 1 MB byte limit.
+    push_pq(buf,
+            /*d=*/0,
+            /*M=*/(size_t{1} << 20),
+            /*nbits=*/8,
+            /*centroids=*/{});
+
+    expect_read_throws_with(buf, "M*ksub");
+
+    set_deserialization_vector_byte_limit(old_limit);
+}
+
+// -----------------------------------------------------------------------
+// Test: initialize_IVFPQ_precomputed_table rejects parameter combinations
+// whose nlist * pq.M * pq.ksub * sizeof(float) computation overflows
+// size_t. Without overflow-checked multiplication, the wrapped value
+// silently bypasses the precomputed_table_max_bytes guard and the
+// subsequent r_norms allocation aborts the process.
+// -----------------------------------------------------------------------
+TEST(ReadIndexDeserialize, InitializeIVFPQPrecomputedTableOverflowRejected) {
+    // Construct a PQ with M*ksub already saturating most of size_t. The
+    // nlist factor (taken from the quantizer's ntotal) then forces overflow
+    // when multiplied in.
+    ProductQuantizer pq;
+    pq.d = 0;
+    pq.M = size_t{1} << 40;
+    pq.nbits = 24;
+    pq.ksub = size_t{1} << pq.nbits;
+    pq.dsub = 0;
+    pq.code_size = 0;
+    // Skip set_derived_values (which would resize centroids); we are
+    // exercising initialize_IVFPQ_precomputed_table's arithmetic guard, not
+    // PQ training.
+
+    IndexFlatL2 quantizer(/*d_in=*/0);
+    quantizer.ntotal = size_t{1} << 20; // nlist
+    AlignedTable<float> precomputed_table;
+    int use_precomputed_table = 0;
+
+    EXPECT_THROW(
+            initialize_IVFPQ_precomputed_table(
+                    use_precomputed_table,
+                    &quantizer,
+                    pq,
+                    precomputed_table,
+                    /*by_residual=*/true,
+                    /*verbose=*/false),
+            faiss::FaissException);
+}
+
+// -----------------------------------------------------------------------
+// Test: initialize_IVFPQ_precomputed_table rejects an overflowing M*ksub
+// even when the caller has pre-set use_precomputed_table to 1, which
+// bypasses the in-function size-class branch. Protects in-memory IVFPQ
+// callers (training, IndexHNSW, IndexIVFPQFastScan) and the
+// IndexIVFIndependentQuantizer deserialization path that reads
+// use_precomputed_table directly from the stream.
+// -----------------------------------------------------------------------
+TEST(ReadIndexDeserialize,
+     InitializeIVFPQPrecomputedTableUserSetFlagOverflowRejected) {
+    ProductQuantizer pq;
+    pq.d = 0;
+    pq.M = size_t{1} << 40;
+    pq.nbits = 24;
+    pq.ksub = size_t{1} << pq.nbits;
+    pq.dsub = 0;
+    pq.code_size = 0;
+
+    IndexFlatL2 quantizer(/*d_in=*/0);
+    quantizer.ntotal = 1; // nlist
+    AlignedTable<float> precomputed_table;
+    int use_precomputed_table = 1; // caller pre-selects table type 1
+
+    EXPECT_THROW(
+            initialize_IVFPQ_precomputed_table(
+                    use_precomputed_table,
+                    &quantizer,
+                    pq,
+                    precomputed_table,
+                    /*by_residual=*/true,
+                    /*verbose=*/false),
+            faiss::FaissException);
+}
+
 // ============================================================
 // SVS fourcc rejection / deserialization safety (Group F: T262015608)
 // ============================================================