From 87bc9b08722d3a5ec01a91f5b16236fc56caf3bc Mon Sep 17 00:00:00 2001 From: Justin Gibbs Date: Wed, 1 Apr 2026 15:09:00 -0700 Subject: [PATCH] Validate IndexFastScan quantizer state during deserialization (#5009) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5009 Add deserialization-time validation for all three FastScan index types (IndexPQFastScan, IndexAdditiveQuantizerFastScan, IndexRaBitQFastScan) in `index_read.cpp`. A crafted or corrupted index file could previously supply invalid quantizer parameters that would cause crashes or undefined behavior when the index was later used for search. The following checks are added at deserialization time for each FastScan variant: 1. **M > 0 and ksub > 0**: These quantizer dimensions are used as divisors and loop bounds in `compute_float_LUT` and search kernels. Zero values would cause division-by-zero or empty allocations followed by out-of-bounds access. 2. **bbs > 0 and bbs % 32 == 0**: The block batch size (`bbs`) is used as a divisor (e.g. `ntotal2 / bbs`) in search and is assumed to be aligned to 32 by the SIMD packing logic. A zero value causes division-by-zero; a misaligned value causes incorrect memory access patterns. 3. **ksub * M and ksub * M2 overflow**: These products are used to size lookup tables and scoring buffers. A corrupted file could set M2 to a huge value, causing `size_t` overflow and subsequent under-sized allocations with out-of-bounds writes. The existing `mul_no_overflow` helper is used to detect this. Reviewed By: mnorris11 Differential Revision: D98787875 fbshipit-source-id: 76656f1c1bea4afe2d268dbd46af5a6b3a7b5932 --- faiss/impl/index_read.cpp | 52 ++++++++ tests/test_read_index_deserialize.cpp | 180 ++++++++++++++++++++++++++ 2 files changed, 232 insertions(+) diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index bfaeecf21..44cb75371 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -1480,6 +1480,27 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { READ1(idxaqfs->max_train_points); READVECTOR(idxaqfs->codes); + + FAISS_THROW_IF_NOT_FMT( + idxaqfs->M > 0 && idxaqfs->ksub > 0, + "IndexAdditiveQuantizerFastScan: invalid quantizer state " + "(M=%zd, ksub=%zd, must be > 0)", + idxaqfs->M, + idxaqfs->ksub); + FAISS_THROW_IF_NOT_FMT( + idxaqfs->bbs > 0 && idxaqfs->bbs % 32 == 0, + "IndexAdditiveQuantizerFastScan: invalid bbs=%d " + "(must be > 0 and a multiple of 32)", + idxaqfs->bbs); + mul_no_overflow( + idxaqfs->ksub, + idxaqfs->M, + "IndexAdditiveQuantizerFastScan ksub * M"); + mul_no_overflow( + idxaqfs->ksub, + idxaqfs->M2, + "IndexAdditiveQuantizerFastScan ksub * M2"); + idx = std::move(idxaqfs); } else if ( h == fourcc("IVLf") || h == fourcc("IVRf") || h == fourcc("NPLf") || @@ -1968,6 +1989,21 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxpqfs->ksub = (1 << pq.nbits); idxpqfs->code_size = pq.code_size; + FAISS_THROW_IF_NOT_FMT( + idxpqfs->M > 0 && idxpqfs->ksub > 0, + "IndexPQFastScan: invalid quantizer state " + "(M=%zd, ksub=%zd, must be > 0)", + idxpqfs->M, + idxpqfs->ksub); + FAISS_THROW_IF_NOT_FMT( + idxpqfs->bbs > 0 && idxpqfs->bbs % 32 == 0, + "IndexPQFastScan: invalid bbs=%d " + "(must be > 0 and a multiple of 32)", + idxpqfs->bbs); + mul_no_overflow(idxpqfs->ksub, idxpqfs->M, "IndexPQFastScan ksub * M"); + mul_no_overflow( + idxpqfs->ksub, idxpqfs->M2, "IndexPQFastScan ksub * M2"); + idx = std::move(idxpqfs); } else if (h == fourcc("IwPf")) { @@ -2038,6 +2074,22 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxqfs->nbits = nbits_fastscan; idxqfs->ksub = (1 << nbits_fastscan); + FAISS_THROW_IF_NOT_FMT( + idxqfs->M > 0 && idxqfs->ksub > 0, + "IndexRaBitQFastScan: invalid quantizer state " + "(M=%zd, ksub=%zd, must be > 0)", + idxqfs->M, + idxqfs->ksub); + FAISS_THROW_IF_NOT_FMT( + idxqfs->bbs > 0 && idxqfs->bbs % 32 == 0, + "IndexRaBitQFastScan: invalid bbs=%d " + "(must be > 0 and a multiple of 32)", + idxqfs->bbs); + mul_no_overflow( + idxqfs->ksub, idxqfs->M, "IndexRaBitQFastScan ksub * M"); + mul_no_overflow( + idxqfs->ksub, idxqfs->M2, "IndexRaBitQFastScan ksub * M2"); + READVECTOR(idxqfs->codes); if (is_legacy) { diff --git a/tests/test_read_index_deserialize.cpp b/tests/test_read_index_deserialize.cpp index b1e9bfcfd..22ad22824 100644 --- a/tests/test_read_index_deserialize.cpp +++ b/tests/test_read_index_deserialize.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -20,6 +21,8 @@ #include #include #include + +#include #include #include #include @@ -2318,6 +2321,183 @@ TEST(ReadIndexDeserialize, IwIQVtDoutMismatch) { EXPECT_THROW(read_index(&reader), FaissException); } +/// Helper: append a minimal valid serialized IndexPQFastScan ("IPfs"). +/// Caller can override bbs and M2 to inject invalid values. +static std::vector build_IndexPQFastScan_buf( + int bbs = 32, + size_t M2 = 2) { + // PQ: d=4, M=2, nbits=4 → ksub=16, centroids size = d*ksub = 64 floats + std::vector buf; + push_fourcc(buf, "IPfs"); + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + push_pq(buf, /*d=*/4, /*M=*/2, /*nbits=*/4, std::vector(64, 0.0f)); + push_val(buf, 0); // implem + push_val(buf, bbs); // bbs + push_val(buf, 0); // qbs + push_val(buf, 0); // ntotal2 + push_val(buf, M2); // M2 + push_vector(buf, {}); // codes + return buf; +} + +/// Helper: append a minimal valid serialized +/// IndexResidualQuantizerFastScan ("IRfs"). +/// Writes AdditiveQuantizer + ResidualQuantizer fields, then FastScan fields. +/// Caller can override FastScan M, ksub, bbs to inject invalid values. +static std::vector build_AQFastScan_buf( + size_t fastscan_M = 3, + size_t fastscan_ksub = 16, + int bbs = 32) { + std::vector buf; + push_fourcc(buf, "IRfs"); + push_index_header(buf, /*d=*/4, /*ntotal=*/0); + + // AdditiveQuantizer fields: + push_val(buf, 4); // d + push_val(buf, 1); // M (AQ M, not FastScan M) + push_vector(buf, {4}); // nbits = [4] + push_val(buf, false); // is_trained + push_vector(buf, std::vector(64, 0.0f)); // codebooks (d*16) + push_val(buf, 0); // search_type = ST_decompress + push_val(buf, 0.0f); // norm_min + push_val(buf, 1.0f); // norm_max + + // ResidualQuantizer extras: + push_val(buf, 2048); // train_type = Skip_codebook_tables + push_val(buf, 1); // max_beam_size + + // FastScan fields (IndexAdditiveQuantizerFastScan): + push_val(buf, 0); // implem + push_val(buf, bbs); // bbs + push_val(buf, 0); // qbs + push_val(buf, fastscan_M); // M + push_val(buf, 4); // nbits + push_val(buf, fastscan_ksub); // ksub + push_val(buf, 2); // code_size + push_val(buf, 0); // ntotal2 + push_val(buf, fastscan_M + (fastscan_M % 2)); // M2 (rounded up) + push_val(buf, true); // rescale_norm + push_val(buf, 1); // norm_scale + push_val(buf, 48); // max_train_points + push_vector(buf, {}); // codes + + return buf; +} + +/// Helper: serialize a valid IndexRaBitQFastScan and return the raw bytes. +/// Then locate the bbs field and patch it to the given value. +static std::vector build_RaBitQFastScan_buf(int bbs) { + IndexRaBitQFastScan idx(4, METRIC_L2, 32, 1); + + // Serialize the valid index + VectorIOWriter writer; + write_index(&idx, &writer); + + auto buf = std::move(writer.data); + + // Locate bbs (int, value 32) in the serialized data. + // Format: ... qb (uint8_t=8) then bbs (int=32). + // Search for qb=8 (1 byte) followed by bbs=32 (4 bytes). + uint8_t target_qb = 8; + int target_bbs = 32; + bool patched = false; + for (size_t i = 0; i + sizeof(uint8_t) + sizeof(int) <= buf.size(); i++) { + uint8_t val_qb; + int val_bbs; + memcpy(&val_qb, &buf[i], sizeof(uint8_t)); + memcpy(&val_bbs, &buf[i + sizeof(uint8_t)], sizeof(int)); + if (val_qb == target_qb && val_bbs == target_bbs) { + memcpy(&buf[i + sizeof(uint8_t)], &bbs, sizeof(int)); + patched = true; + break; + } + } + EXPECT_TRUE(patched) << "Could not find bbs field in serialized data"; + return buf; +} + +// ----------------------------------------------------------------------- +// IndexPQFastScan deserialization: bbs=0 must be rejected. +// Without this check, search() would divide by zero in ntotal2/bbs. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexPQFastScanBbsZero) { + auto buf = build_IndexPQFastScan_buf(/*bbs=*/0); + expect_read_throws_with(buf, "invalid bbs"); +} + +// ----------------------------------------------------------------------- +// IndexPQFastScan deserialization: bbs not a multiple of 32. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexPQFastScanBbsNotAligned) { + auto buf = build_IndexPQFastScan_buf(/*bbs=*/33); + expect_read_throws_with(buf, "invalid bbs"); +} + +// ----------------------------------------------------------------------- +// IndexPQFastScan deserialization: ksub * M2 overflow. +// M2 is read directly from the file and could be corrupted to a huge +// value that causes ksub * M2 to overflow size_t. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexPQFastScanKsubM2Overflow) { + auto buf = build_IndexPQFastScan_buf( + /*bbs=*/32, /*M2=*/std::numeric_limits::max()); + expect_read_throws_with(buf, "overflow"); +} + +// ----------------------------------------------------------------------- +// IndexAdditiveQuantizerFastScan deserialization: M=0 must be rejected. +// Without this check, search() would crash in compute_float_LUT. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, AQFastScanMZero) { + auto buf = build_AQFastScan_buf(/*fastscan_M=*/0); + expect_read_throws_with(buf, "invalid quantizer state"); +} + +// ----------------------------------------------------------------------- +// IndexAdditiveQuantizerFastScan deserialization: ksub=0 must be rejected. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, AQFastScanKsubZero) { + auto buf = build_AQFastScan_buf( + /*fastscan_M=*/3, /*fastscan_ksub=*/0); + expect_read_throws_with(buf, "invalid quantizer state"); +} + +// ----------------------------------------------------------------------- +// IndexAdditiveQuantizerFastScan deserialization: bbs=0 must be rejected. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, AQFastScanBbsZero) { + auto buf = build_AQFastScan_buf( + /*fastscan_M=*/3, /*fastscan_ksub=*/16, /*bbs=*/0); + expect_read_throws_with(buf, "invalid bbs"); +} + +// ----------------------------------------------------------------------- +// IndexAdditiveQuantizerFastScan deserialization: ksub * M overflow. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, AQFastScanKsubMOverflow) { + auto buf = build_AQFastScan_buf( + /*fastscan_M=*/std::numeric_limits::max(), + /*fastscan_ksub=*/16, + /*bbs=*/32); + expect_read_throws_with(buf, "overflow"); +} + +// ----------------------------------------------------------------------- +// IndexRaBitQFastScan deserialization: bbs=0 must be rejected. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, RaBitQFastScanBbsZero) { + auto buf = build_RaBitQFastScan_buf(/*bbs=*/0); + expect_read_throws_with(buf, "invalid bbs"); +} + +// ----------------------------------------------------------------------- +// IndexRaBitQFastScan deserialization: bbs not a multiple of 32. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, RaBitQFastScanBbsNotAligned) { + auto buf = build_RaBitQFastScan_buf(/*bbs=*/17); + expect_read_throws_with(buf, "invalid bbs"); +} + // ----------------------------------------------------------------------- // Test: read_AdditiveQuantizer rejects codebooks that are too small for // d * total_codebook_size. Previously this was only caught at