Make Panorama bounds checks overflow-safe (#5616)

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

Follow up on D118495504 by closing two arithmetic gaps in the Panorama
bounds validation. Avoid signed overflow and reject negative counts in
`IndexFlatPanorama::reconstruct_n`, and compute the serialized
`IndexFlatPanorama` batch count without overflowing
`ntotal + batch_size - 1`.

Add regression coverage for negative reconstruction counts and a maximal
serialized batch size that previously admitted empty storage for a
non-empty index.

Reviewed By: mnorris11

Differential Revision: D119272683

fbshipit-source-id: 7ea841bd4bd90d0358b78797441fa3cc815db46a

faiss/IndexFlat.cpp

@@ -711,7 +711,7 @@ void IndexFlatPanorama::reconstruct(idx_t key, float* recons) const {
 }
 
 void IndexFlatPanorama::reconstruct_n(idx_t i, idx_t n, float* recons) const {
-    FAISS_THROW_IF_NOT(n == 0 || (i >= 0 && i + n <= ntotal));
+    FAISS_THROW_IF_NOT(i >= 0 && i <= ntotal && n >= 0 && n <= ntotal - i);
     Index::reconstruct_n(i, n, recons);
 }
 

faiss/impl/index_read.cpp

@@ -1715,21 +1715,29 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         READ1_BOOL(idxp->is_trained);
         READVECTOR(idxp->codes);
         READVECTOR(idxp->cum_sums);
-        size_t num_slots = mul_no_overflow(
-                ((size_t)idxp->ntotal + idxp->batch_size - 1) /
-                        idxp->batch_size,
-                idxp->batch_size,
+        const size_t ntotal = (size_t)idxp->ntotal;
+        const size_t num_batches =
+                ntotal / batch_size + (ntotal % batch_size != 0);
+        const size_t num_slots = mul_no_overflow(
+                num_batches,
+                batch_size,
                 "IndexFlatPanorama num_batches*batch_size");
-        FAISS_THROW_IF_NOT(
-                idxp->codes.size() ==
-                mul_no_overflow(
-                        num_slots, idxp->code_size, "IndexFlatPanorama codes"));
-        FAISS_THROW_IF_NOT(
-                idxp->cum_sums.size() ==
-                mul_no_overflow(
-                        num_slots,
-                        idxp->pano.n_levels + 1,
-                        "IndexFlatPanorama cum_sums"));
+        const size_t expected_codes_size = mul_no_overflow(
+                num_slots, idxp->code_size, "IndexFlatPanorama codes");
+        FAISS_THROW_IF_NOT_FMT(
+                idxp->codes.size() == expected_codes_size,
+                "IndexFlatPanorama codes size mismatch: got %zu, expected %zu",
+                idxp->codes.size(),
+                expected_codes_size);
+        const size_t expected_cum_sums_size = mul_no_overflow(
+                num_slots,
+                idxp->pano.n_levels + 1,
+                "IndexFlatPanorama cum_sums");
+        FAISS_THROW_IF_NOT_FMT(
+                idxp->cum_sums.size() == expected_cum_sums_size,
+                "IndexFlatPanorama cum_sums size mismatch: got %zu, expected %zu",
+                idxp->cum_sums.size(),
+                expected_cum_sums_size);
         idxp->verbose = false;
         idx = std::move(idxp);
     } else if (

tests/test_read_index_deserialize.cpp

@@ -10,6 +10,7 @@
 #include <cstdint>
 #include <cstring>
 #include <filesystem>
+#include <limits>
 #include <numeric>
 #include <random>
 #include <string>
@@ -4865,6 +4866,57 @@ TEST(ReadIndexDeserialize, FlatPanoramaNegativeNtotal) {
     expect_read_throws_with(buf, "invalid ntotal");
 }
 
+TEST(ReadIndexDeserialize, FlatPanoramaZeroBatchSizeRejected) {
+    std::vector<uint8_t> buf;
+    push_flat_panorama(
+            buf,
+            /*d=*/8,
+            /*n_levels=*/1,
+            /*batch_size=*/0,
+            /*ntotal=*/2,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    expect_read_throws_with(buf, "invalid IxFP batch_size");
+}
+
+TEST(ReadIndexDeserialize, FlatPanoramaCodesSizeOverflowRejected) {
+    std::vector<uint8_t> buf;
+    push_flat_panorama(
+            buf,
+            /*d=*/8,
+            /*n_levels=*/1,
+            /*batch_size=*/(std::numeric_limits<size_t>::max)(),
+            /*ntotal=*/2,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    // With batch_size=SIZE_MAX and ntotal=2, num_batches computes to 1 and
+    // num_slots=1*SIZE_MAX does not overflow. The throw actually originates
+    // from the later num_slots * code_size multiplication in mul_no_overflow,
+    // which emits "integer overflow in IndexFlatPanorama codes: ...". Match
+    // on that full prefix so this test cannot be silently satisfied by the
+    // sibling "IndexFlatPanorama codes size mismatch" throw.
+    expect_read_throws_with(buf, "integer overflow in IndexFlatPanorama codes");
+}
+
+TEST(ReadIndexDeserialize, FlatPanoramaCodesSizeMismatch) {
+    std::vector<uint8_t> buf;
+    push_flat_panorama(
+            buf,
+            /*d=*/8,
+            /*n_levels=*/1,
+            /*batch_size=*/4,
+            /*ntotal=*/2,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    // Match on the explicit FAISS_THROW_IF_NOT_FMT message so the test is
+    // pinned to the intended error rather than a stringified condition
+    // expression that would break on trivial refactors.
+    expect_read_throws_with(buf, "IndexFlatPanorama codes size mismatch");
+}
+
 // End-to-end shape of the fuzzer crash: an IndexPreTransform with an empty
 // chain wrapping an empty IndexFlatPanorama. Everything deserializes, and
 // reconstruct() must then reject the out-of-range key instead of memcpy'ing
@@ -4934,7 +4986,15 @@ TEST(IndexFlatPanoramaSafety, ReconstructRejectsOutOfRangeKey) {
     EXPECT_THROW(index.reconstruct(2, recons.data()), FaissException);
     EXPECT_THROW(index.reconstruct(-1, recons.data()), FaissException);
     EXPECT_NO_THROW(index.reconstruct_n(0, 2, recons.data()));
+    EXPECT_THROW(index.reconstruct_n(-1, 1, recons.data()), FaissException);
+    EXPECT_THROW(index.reconstruct_n(0, -1, recons.data()), FaissException);
+    EXPECT_THROW(
+            index.reconstruct_n(
+                    1, (std::numeric_limits<idx_t>::max)(), recons.data()),
+            FaissException);
     EXPECT_THROW(index.reconstruct_n(1, 2, recons.data()), FaissException);
+    EXPECT_THROW(index.reconstruct_n(-1, 0, recons.data()), FaissException);
+    EXPECT_THROW(index.reconstruct_n(3, 0, recons.data()), FaissException);
 }
 
 namespace {