Fix OOB read in Panorama::reconstruct via IndexPreTransform (T287094917) (#5562)

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

Lionhead harness `faiss_rmdt_post_read_reconstruct_fuzzer` found a
heap out-of-bounds read at `Panorama.cpp:171`, reported as
`READ of size 5122903332` in `__asan_memcpy`.

Decoding the 232-byte crash input, it is `IxPT -> IxPT -> IxPT -> IxFp`
where every `IndexPreTransform` carries an **empty** transform chain and
the innermost `IndexFlatIPPanorama` has `d = 0x4C565349 = 1280725833`
(so `code_size = d * 4 = 5122903332`, exactly the ASAN read size),
`ntotal = 0`, and an empty `codes` vector. Three independent defects
combine:

1. `IndexFlatPanorama::reconstruct` never bounds-checks `key`. It
overrides `IndexFlat::reconstruct`, which does have
`FAISS_THROW_IF_NOT(key < ntotal)`, and drops the guard. Overriding
`reconstruct_n` likewise removes `IndexFlatCodes::reconstruct_n`'s
`i0 + ni <= ntotal` check. With `ntotal == 0` and empty `codes`,
`reconstruct(0, ...)` memcpy's `code_size` bytes out of a zero-length
buffer. `Panorama::reconstruct` takes an unsized `const uint8_t*`, so it
cannot self-validate -- the caller has to.

2. The `IxFP`/`IxFp` read branch does not range-check `d`. It is the one
index type that does not go through `read_index_header`, so it misses
that function's `FAISS_CHECK_RANGE(d, 0, (1 << 20) + 1)` and
`ntotal >= 0` checks. The `codes.size()` equality check below it does
not catch this, because `ntotal == 0` makes `num_slots == 0` and an
empty `codes` is then "valid".

3. `IndexPreTransform` deserialization skips all consistency checks when
the chain is empty -- the `chain[0].d_in == ixpt->d` /
`chain[nt-1].d_out == index->d` validation sits entirely inside
`if (nt > 0)`. With `nt == 0`, `IndexPreTransform::reconstruct` passes
the caller's buffer (sized from the outer `d`) straight to the
sub-index, so a mismatched sub-index `d` is also a write-side overflow,
for any sub-index type.

Each is fixed:

- `IndexFlatPanorama::reconstruct` / `reconstruct_n` now range-check
  against `ntotal`.
- The `IxFP`/`IxFp` branch applies the same `d` and `ntotal` checks as
  `read_index_header`.
- `read_index` requires `index->d == ixpt->d` when the `IxPT` chain is
  empty.

Two adjacent gaps in the same audit are closed as well:

- `IHfP` (`IndexHNSWFlatPanorama`) never validated `cum_sums.size()`,
  unlike the equivalent `IxFP` check. Both `get_cum_sum()` and
  `HNSW.cpp`'s Panorama search path index it at stride
  `pano.n_levels + 1`, so a short vector is an OOB read at search time.
  It also had no explicit `n_levels > 0` guard.
- `ArrayInvertedListsPanorama::get_single_code` bounds-checked with
  `assert`, which is compiled out in opt builds. It reaches the same
  `Panorama::reconstruct` sink via
  `IndexIVFFlatPanorama::reconstruct_from_offset`.

Reviewed By: mnorris11

Differential Revision: D118495504

fbshipit-source-id: 46bc7a1a6c9335a5589e20221b9478ae5505b9da

faiss/IndexFlat.cpp

@@ -704,10 +704,12 @@ void IndexFlatPanorama::reset() {
 }
 
 void IndexFlatPanorama::reconstruct(idx_t key, float* recons) const {
+    FAISS_THROW_IF_NOT(key >= 0 && key < ntotal);
     pano.reconstruct(key, recons, codes.data());
 }
 
 void IndexFlatPanorama::reconstruct_n(idx_t i, idx_t n, float* recons) const {
+    FAISS_THROW_IF_NOT(n == 0 || (i >= 0 && i + n <= ntotal));
     Index::reconstruct_n(i, n, recons);
 }
 

faiss/impl/index_read.cpp

@@ -1692,6 +1692,10 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         int d;
         size_t n_levels, batch_size;
         READ1(d);
+        // This branch does not go through read_index_header, so the range
+        // check that guards `d` for every other index type has to be
+        // repeated here.
+        FAISS_CHECK_RANGE(d, 0, (1 << 20) + 1);
         READ1(n_levels);
         FAISS_THROW_IF_NOT_FMT(n_levels > 0, "invalid n_levels %zd", n_levels);
         READ1(batch_size);
@@ -1706,6 +1710,10 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
                     d, n_levels, batch_size);
         }
         READ1(idxp->ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                idxp->ntotal >= 0,
+                "invalid ntotal %" PRId64 " read from IxFP index",
+                (int64_t)idxp->ntotal);
         READ1_BOOL(idxp->is_trained);
         READVECTOR(idxp->codes);
         READVECTOR(idxp->cum_sums);
@@ -2384,8 +2392,18 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         // Validate transform chain dimension consistency:
         // chain[0].d_in must equal the outer index d, consecutive
         // transforms must have matching d_out/d_in, and the last
-        // transform's d_out must equal the sub-index d.
-        if (nt > 0) {
+        // transform's d_out must equal the sub-index d. With an empty
+        // chain the sub-index d must equal the outer d directly, because
+        // reconstruct() then hands the caller's buffer -- sized from the
+        // outer d -- straight to the sub-index.
+        if (nt == 0 && ixpt->index) {
+            FAISS_THROW_IF_NOT_FMT(
+                    ixpt->index->d == ixpt->d,
+                    "IndexPreTransform empty chain: sub-index d=%d "
+                    "!= index d=%d",
+                    ixpt->index->d,
+                    ixpt->d);
+        } else if (nt > 0) {
             FAISS_THROW_IF_NOT_FMT(
                     ixpt->chain[0]->d_in == ixpt->d,
                     "IndexPreTransform chain[0] d_in=%d != index d=%d",
@@ -2537,10 +2555,21 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
                     "dynamic_cast to IndexHNSWFlatPanorama failed");
             size_t nlevels;
             READ1(nlevels);
+            FAISS_THROW_IF_NOT_FMT(
+                    nlevels > 0, "invalid IHfP n_levels %zd", nlevels);
             const_cast<size_t&>(idx_panorama->num_panorama_levels) = nlevels;
             const_cast<Panorama&>(idx_panorama->pano) =
                     Panorama(idx_panorama->d * sizeof(float), nlevels, 1);
             READVECTOR(idx_panorama->cum_sums);
+            // Both the search path and get_cum_sum() index cum_sums at
+            // stride pano.n_levels + 1, which may be lower than the
+            // serialized nlevels if set_derived_values() truncated it.
+            FAISS_THROW_IF_NOT(
+                    idx_panorama->cum_sums.size() ==
+                    mul_no_overflow(
+                            (size_t)idx_panorama->ntotal,
+                            idx_panorama->pano.n_levels + 1,
+                            "IndexHNSWFlatPanorama cum_sums"));
         } else if (h == fourcc("IHNc") || h == fourcc("IHc2")) {
             READ1_BOOL(idxhnsw->keep_max_size_level0);
             auto idx_hnsw_cagra = dynamic_cast<IndexHNSWCagra*>(idxhnsw.get());

faiss/invlists/InvertedLists.cpp

@@ -440,8 +440,11 @@ void ArrayInvertedListsPanorama::resize(size_t list_no, size_t new_size) {
 const uint8_t* ArrayInvertedListsPanorama::get_single_code(
         size_t list_no,
         size_t offset) const {
-    assert(list_no < nlist);
-    assert(offset < ids[list_no].size());
+    // Throw rather than assert: Panorama::reconstruct takes an unsized
+    // pointer, so these are the only bounds available and they must hold
+    // in opt builds too.
+    FAISS_THROW_IF_NOT(list_no < nlist);
+    FAISS_THROW_IF_NOT(offset < ids[list_no].size());
 
     uint8_t* recons_buffer = new uint8_t[code_size];
 

tests/test_read_index_deserialize.cpp

@@ -2076,6 +2076,36 @@ TEST(ReadIndexDeserialize, PreTransformChainDimensionMismatch) {
     expect_read_throws_with(buf, "d_out=4");
 }
 
+// Test: IndexPreTransform with an *empty* chain must still agree with its
+// sub-index on d. With no transform, reconstruct() hands the caller's
+// buffer -- sized from the outer d -- straight to the sub-index, so a
+// larger sub-index d is a heap overflow (T287094917).
+TEST(ReadIndexDeserialize, PreTransformEmptyChainDimensionMismatch) {
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IxPT");
+    push_index_header(buf, /*d=*/4, /*ntotal=*/0);
+    push_val<int>(buf, 0); // nt = 0, empty chain
+    push_minimal_flat(buf, /*d=*/8);
+
+    expect_read_throws_with(buf, "IndexPreTransform empty chain");
+}
+
+// A well-formed empty-chain IxPT still deserializes.
+TEST(ReadIndexDeserialize, PreTransformEmptyChainMatchingDimensionAccepted) {
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IxPT");
+    push_index_header(buf, /*d=*/8, /*ntotal=*/0);
+    push_val<int>(buf, 0); // nt = 0, empty chain
+    push_minimal_flat(buf, /*d=*/8);
+
+    VectorIOReader reader;
+    reader.data = buf;
+    EXPECT_NO_THROW({
+        auto idx = read_index_up(&reader);
+        EXPECT_NE(idx, nullptr);
+    });
+}
+
 TEST(ReadIndexDeserialize, HadamardRotationInvalidDout) {
     // HRot format: fourcc("HRot") + seed(int) + d_in(int) + d_out(int) +
     //              is_trained(bool)
@@ -4763,3 +4793,134 @@ TEST(IndexIDMapSafety, SearchRemapsOutOfRangeLabelsToSentinel) {
             idmap.search(1, xq.data(), 1, distances.data(), labels.data()));
     EXPECT_EQ(labels[0], -1);
 }
+
+// -----------------------------------------------------------------------
+// Panorama deserialization and reconstruct bounds (T287094917).
+//
+// "IxFP"/"IxFp" format: fourcc + d(int) + n_levels + batch_size +
+// ntotal(int64) + is_trained(bool) + codes vec + cum_sums vec.
+// Unlike every other index type this branch does not go through
+// read_index_header, so its own range checks are all that guard d and
+// ntotal.
+// -----------------------------------------------------------------------
+static void push_flat_panorama(
+        std::vector<uint8_t>& buf,
+        int d,
+        size_t n_levels,
+        size_t batch_size,
+        int64_t ntotal,
+        size_t codes_bytes,
+        size_t cum_sums_count) {
+    push_fourcc(buf, "IxFP");
+    push_val<int>(buf, d);
+    push_val<size_t>(buf, n_levels);
+    push_val<size_t>(buf, batch_size);
+    push_val<int64_t>(buf, ntotal);
+    push_val<bool>(buf, true);
+    push_val<size_t>(buf, codes_bytes);
+    buf.resize(buf.size() + codes_bytes, 0);
+    push_val<size_t>(buf, cum_sums_count);
+    buf.resize(buf.size() + cum_sums_count * sizeof(float), 0);
+}
+
+// A d large enough to make code_size = d * 4 overflow into gigabytes must
+// be rejected at read time, as it is for every other index type.
+TEST(ReadIndexDeserialize, FlatPanoramaDimensionOutOfRange) {
+    std::vector<uint8_t> buf;
+    push_flat_panorama(
+            buf,
+            /*d=*/1280725833,
+            /*n_levels=*/1,
+            /*batch_size=*/1024,
+            /*ntotal=*/0,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    expect_read_throws_with(buf, "out of range");
+}
+
+TEST(ReadIndexDeserialize, FlatPanoramaNegativeNtotal) {
+    std::vector<uint8_t> buf;
+    push_flat_panorama(
+            buf,
+            /*d=*/8,
+            /*n_levels=*/1,
+            /*batch_size=*/4,
+            /*ntotal=*/-1,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    expect_read_throws_with(buf, "invalid ntotal");
+}
+
+// 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
+// code_size bytes out of a zero-length codes buffer.
+TEST(ReadIndexDeserialize, PreTransformOverEmptyFlatPanoramaReconstruct) {
+    constexpr int d = 8;
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IxPT");
+    push_index_header(buf, d, /*ntotal=*/0);
+    push_val<int>(buf, 0); // nt = 0, empty chain
+    push_flat_panorama(
+            buf,
+            d,
+            /*n_levels=*/1,
+            /*batch_size=*/4,
+            /*ntotal=*/0,
+            /*codes_bytes=*/0,
+            /*cum_sums_count=*/0);
+
+    VectorIOReader reader;
+    reader.data = buf;
+    std::unique_ptr<Index> idx;
+    ASSERT_NO_THROW({ idx = read_index_up(&reader); });
+    ASSERT_NE(idx, nullptr);
+
+    std::vector<float> recons(d, 0.0f);
+    EXPECT_THROW(idx->reconstruct(0, recons.data()), FaissException);
+}
+
+// "IHfP" format: fourcc + index_header + n_levels(size_t) + cum_sums vec +
+// HNSW graph + storage index. cum_sums must hold ntotal * (n_levels + 1)
+// floats; a short vector is an out-of-bounds read on the search path.
+TEST(ReadIndexDeserialize, HNSWFlatPanoramaCumSumsSizeMismatch) {
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IHfP");
+    push_index_header(buf, /*d=*/4, /*ntotal=*/1);
+    push_val<size_t>(buf, size_t(1)); // n_levels
+    push_vector<float>(buf, {0.0f});  // cum_sums: 1 float, needs 2
+
+    expect_read_throws_with(buf, "IndexHNSWFlatPanorama cum_sums");
+}
+
+TEST(ReadIndexDeserialize, HNSWFlatPanoramaZeroLevels) {
+    std::vector<uint8_t> buf;
+    push_fourcc(buf, "IHfP");
+    push_index_header(buf, /*d=*/4, /*ntotal=*/0);
+    push_val<size_t>(buf, size_t(0)); // n_levels = 0
+
+    expect_read_throws_with(buf, "invalid IHfP n_levels");
+}
+
+// IndexFlatPanorama overrides reconstruct()/reconstruct_n() and so does not
+// inherit the range guards in IndexFlat / IndexFlatCodes. Panorama's
+// level-major layout means an out-of-range key reads whole batch strides
+// past the end of `codes`.
+TEST(IndexFlatPanoramaSafety, ReconstructRejectsOutOfRangeKey) {
+    constexpr int d = 8;
+    IndexFlatL2Panorama index(d, /*n_levels_in=*/2, /*batch_size_in=*/4);
+    std::vector<float> recons(2 * d, 0.0f);
+
+    EXPECT_THROW(index.reconstruct(0, recons.data()), FaissException);
+
+    std::vector<float> xb(2 * d, 1.0f);
+    index.add(2, xb.data());
+
+    EXPECT_NO_THROW(index.reconstruct(1, recons.data()));
+    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, 2, recons.data()), FaissException);
+}