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