facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Validate codebooks size at deserialization time in read_AdditiveQuantizer (#5008)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5008

Add deserialization-time validation in `index_read.cpp` for
AdditiveQuantizer-based indices, catching corrupt or malicious index
files early rather than crashing during search.

1. **Codebooks divisibility check in `read_AdditiveQuantizer`**:
   `codebooks.size()` must be a positive multiple of
   `total_codebook_size`. This catches obvious corruption without
   needing the effective dimension.

2. **`validate_codebooks_size` helper**:
   Checks `codebooks.size() >= effective_d * total_codebook_size`.
   Called from `read_ResidualQuantizer` and `read_LocalSearchQuantizer`
   with `d`, and from `read_ProductAdditiveQuantizer` with
   `d / nsplits`.

3. **`d > 0` validation**:
   Reject zero dimension in `read_AdditiveQuantizer`.

4. **ProductAdditiveQuantizer checks**:
   - `d % nsplits == 0` in `read_ProductAdditiveQuantizer`
   - Sub-quantizer `d == d_sub` in `read_ProductResidualQuantizer`
     and `read_ProductLocalSearchQuantizer`

Existing tests updated to write correctly-sized codebooks so they
continue to reach their intended validation points.

Reviewed By: mnorris11

Differential Revision: D98767295

fbshipit-source-id: 95d8debbbb0f3822bae087a546a573e5c2c3c9dd
This commit is contained in:
Justin Gibbs
2026-04-01 15:09:00 -07:00
committed by meta-codesync[bot]
parent d81584d548
commit 4c391564a6
2 changed files with 286 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/impl/index_read.cpp
+65 -2
View File
@@ -610,8 +610,11 @@ void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
static void read_ResidualQuantizer_old(ResidualQuantizer& rq, IOReader* f) {
READ1(rq.d);
FAISS_CHECK_RANGE(rq.d, 0, (1 << 20) + 1);
FAISS_THROW_IF_NOT_FMT(
rq.d > 0, "invalid AdditiveQuantizer d %zd, must be > 0", rq.d);
READ1(rq.M);
FAISS_THROW_IF_NOT_FMT(
rq.M > 0, "invalid AdditiveQuantizer M %zd, must be > 0", rq.M);
READVECTOR(rq.nbits);
FAISS_THROW_IF_NOT_FMT(
rq.nbits.size() == rq.M,
@@ -630,7 +633,8 @@ static void read_ResidualQuantizer_old(ResidualQuantizer& rq, IOReader* f) {
static void read_AdditiveQuantizer(AdditiveQuantizer& aq, IOReader* f) {
READ1(aq.d);
FAISS_CHECK_RANGE(aq.d, 0, (1 << 20) + 1);
FAISS_THROW_IF_NOT_FMT(
aq.d > 0, "invalid AdditiveQuantizer d %zd, must be > 0", aq.d);
READ1(aq.M);
FAISS_THROW_IF_NOT_FMT(
aq.M > 0, "invalid AdditiveQuantizer M %zd, must be > 0", aq.M);
@@ -660,6 +664,37 @@ static void read_AdditiveQuantizer(AdditiveQuantizer& aq, IOReader* f) {
}
aq.set_derived_values();
// Sanity-check codebooks size without knowing the effective dimension.
// codebooks stores effective_d * total_codebook_size floats, so its
// size must be a positive multiple of total_codebook_size.
if (aq.total_codebook_size > 0) {
FAISS_THROW_IF_NOT_FMT(
aq.codebooks.size() >= aq.total_codebook_size &&
aq.codebooks.size() % aq.total_codebook_size == 0,
"AdditiveQuantizer codebooks size %zd is not a positive "
"multiple of total_codebook_size %zd",
aq.codebooks.size(),
aq.total_codebook_size);
}
}
// Validate that the codebooks vector is large enough for the given
// effective dimension. For a standalone AdditiveQuantizer the effective
// dimension equals aq.d. For a ProductAdditiveQuantizer the codebooks
// are sized for d_sub = d / nsplits, so callers pass that instead.
static void validate_codebooks_size(
const AdditiveQuantizer& aq,
size_t effective_d) {
size_t required = mul_no_overflow(
effective_d, aq.total_codebook_size, "codebooks validation");
FAISS_THROW_IF_NOT_FMT(
aq.codebooks.size() >= required,
"AdditiveQuantizer codebooks size %zd too small for "
"d=%zd * total_codebook_size=%zd",
aq.codebooks.size(),
effective_d,
aq.total_codebook_size);
}
static void read_ResidualQuantizer(
@@ -667,6 +702,7 @@ static void read_ResidualQuantizer(
IOReader* f,
int io_flags) {
read_AdditiveQuantizer(rq, f);
validate_codebooks_size(rq, rq.d);
READ1(rq.train_type);
READ1(rq.max_beam_size);
FAISS_THROW_IF_NOT_FMT(
@@ -699,6 +735,7 @@ static void read_ResidualQuantizer(
static void read_LocalSearchQuantizer(LocalSearchQuantizer& lsq, IOReader* f) {
read_AdditiveQuantizer(lsq, f);
validate_codebooks_size(lsq, lsq.d);
READ1(lsq.K);
READ1(lsq.train_iters);
READ1(lsq.encode_ils_iters);
@@ -722,6 +759,12 @@ static void read_ProductAdditiveQuantizer(
"invalid ProductAdditiveQuantizer nsplits %zd (must be > 0)",
paq.nsplits);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(paq.nsplits, "nsplits");
FAISS_THROW_IF_NOT_FMT(
paq.d % paq.nsplits == 0,
"ProductAdditiveQuantizer d=%zd not divisible by nsplits=%zd",
paq.d,
paq.nsplits);
validate_codebooks_size(paq, paq.d / paq.nsplits);
}
static void read_ProductResidualQuantizer(
@@ -730,9 +773,19 @@ static void read_ProductResidualQuantizer(
int io_flags) {
read_ProductAdditiveQuantizer(prq, f);
size_t d_sub = prq.d / prq.nsplits;
for (size_t i = 0; i < prq.nsplits; i++) {
auto rq = std::make_unique<ResidualQuantizer>();
read_ResidualQuantizer(*rq, f, io_flags);
FAISS_THROW_IF_NOT_FMT(
rq->d == d_sub,
"ProductResidualQuantizer sub-quantizer %zd has d=%zd, "
"expected d_sub=%zd (d=%zd / nsplits=%zd)",
i,
rq->d,
d_sub,
prq.d,
prq.nsplits);
prq.quantizers.push_back(rq.release());
}
}
@@ -742,9 +795,19 @@ static void read_ProductLocalSearchQuantizer(
IOReader* f) {
read_ProductAdditiveQuantizer(plsq, f);
size_t d_sub = plsq.d / plsq.nsplits;
for (size_t i = 0; i < plsq.nsplits; i++) {
auto lsq = std::make_unique<LocalSearchQuantizer>();
read_LocalSearchQuantizer(*lsq, f);
FAISS_THROW_IF_NOT_FMT(
lsq->d == d_sub,
"ProductLocalSearchQuantizer sub-quantizer %zd has d=%zd, "
"expected d_sub=%zd (d=%zd / nsplits=%zd)",
i,
lsq->d,
d_sub,
plsq.d,
plsq.nsplits);
plsq.quantizers.push_back(lsq.release());
}
}
tests/test_read_index_deserialize.cpp
+221 -10
View File
@@ -611,10 +611,11 @@ TEST(ReadIndexDeserialize, ProductAdditiveQuantizerZeroNsplits) {
push_val<size_t>(buf, 1); // M
push_vector<size_t>(buf, {8}); // nbits (1 element matching M=1)
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, {}); // codebooks (empty)
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
// codebooks: d * total_codebook_size = 4 * 256 = 1024 floats
push_vector<float>(buf, std::vector<float>(4 * 256, 0.0f));
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
// ProductAdditiveQuantizer field:
push_val<size_t>(buf, 0); // nsplits = 0 (invalid)
@@ -987,19 +988,37 @@ TEST(ReadIndexDeserialize, IndexPQCodesSizeMismatch) {
/// Helper: append a minimal AdditiveQuantizer (d, M, nbits, is_trained,
/// codebooks, search_type=ST_decompress, norm_min, norm_max).
/// Writes zero-filled codebooks of size codebook_d * total_codebook_size so
/// that the deserialization codebooks-size validation passes.
/// For standalone quantizers codebook_d == d. For ProductAdditiveQuantizer
/// codebook_d == d / nsplits. Pass codebook_d = 0 (default) to use d.
static void push_additive_quantizer(
std::vector<uint8_t>& buf,
size_t d,
size_t M,
const std::vector<size_t>& nbits) {
const std::vector<size_t>& nbits,
size_t codebook_d = 0) {
if (codebook_d == 0) {
codebook_d = d;
}
// Compute total_codebook_size = sum(2^nbits[i]).
size_t total_codebook_size = 0;
for (auto nb : nbits) {
total_codebook_size += size_t{1} << nb;
}
push_val<size_t>(buf, d);
push_val<size_t>(buf, M);
push_vector<size_t>(buf, nbits);
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, {}); // codebooks (empty)
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
push_val<bool>(buf, true); // is_trained
push_vector<float>(
buf,
std::vector<float>(
codebook_d * total_codebook_size, 0.0f)); // codebooks
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
}
/// Helper: append a minimal ResidualQuantizer (AdditiveQuantizer +
@@ -2298,3 +2317,195 @@ TEST(ReadIndexDeserialize, IwIQVtDoutMismatch) {
reader.data = corrupted;
EXPECT_THROW(read_index(&reader), FaissException);
}
// -----------------------------------------------------------------------
// Test: read_AdditiveQuantizer rejects codebooks that are too small for
// d * total_codebook_size. Previously this was only caught at
// compute_LUT() time; now deserialization itself must throw.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, AdditiveQuantizerCodebooksTooSmall) {
// Build an IndexResidualQuantizer ("IxRq") with d=4, M=1, nbits={8}.
// Write AdditiveQuantizer fields manually with empty codebooks so that
// the deserialization check fires (codebooks.size()=0 < d*256=1024).
std::vector<uint8_t> buf;
push_fourcc(buf, "IxRq");
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
// AdditiveQuantizer fields (inline, not via push_additive_quantizer
// which now writes correctly-sized codebooks):
push_val<size_t>(buf, 4); // d
push_val<size_t>(buf, 1); // M
push_vector<size_t>(buf, {8}); // nbits
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, {}); // codebooks (empty — triggers check)
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
// ResidualQuantizer fields:
push_val<int>(buf, 2048); // train_type = Skip_codebook_tables
push_val<int>(buf, 1); // max_beam_size
// IndexResidualQuantizer fields:
push_val<size_t>(buf, 1); // code_size
push_vector<uint8_t>(buf, {}); // codes (empty, matches ntotal=0)
expect_read_throws_with(buf, "not a positive multiple");
}
// -----------------------------------------------------------------------
// Test: AdditiveQuantizer with d=0 is rejected during deserialization.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, AdditiveQuantizerDZero) {
std::vector<uint8_t> buf;
push_fourcc(buf, "IxRq");
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
// AdditiveQuantizer fields with d=0:
push_val<size_t>(buf, 0); // d = 0 (invalid)
push_val<size_t>(buf, 1); // M
push_vector<size_t>(buf, {8}); // nbits
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, {}); // codebooks
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
expect_read_throws_with(buf, "invalid AdditiveQuantizer d");
}
// -----------------------------------------------------------------------
// Test: AdditiveQuantizer with codebooks size not a multiple of
// total_codebook_size is rejected.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, AdditiveQuantizerCodebooksNotMultiple) {
// d=4, M=1, nbits={8} → total_codebook_size=256.
// codebooks should be a multiple of 256 floats. Provide 257 floats.
std::vector<uint8_t> buf;
push_fourcc(buf, "IxRq");
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
push_val<size_t>(buf, 4); // d
push_val<size_t>(buf, 1); // M
push_vector<size_t>(buf, {8}); // nbits
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, std::vector<float>(257, 0.0f)); // not a multiple
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
expect_read_throws_with(buf, "not a positive multiple");
}
// -----------------------------------------------------------------------
// Test: ProductAdditiveQuantizer with d not divisible by nsplits throws.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, ProductAdditiveQuantizerDNotDivisible) {
// d=5, nsplits=2 → d % nsplits != 0
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPR");
push_index_header(buf, /*d=*/5, /*ntotal=*/0);
// AdditiveQuantizer fields: d=5, M=2, nbits={4,4}
// For a PAQ with nsplits=2, codebook_d should be d/nsplits, but since
// d is not divisible we can just write codebooks sized for d
// (the divisibility check fires before the codebooks size check).
push_val<size_t>(buf, 5); // d
push_val<size_t>(buf, 2); // M
push_vector<size_t>(buf, {4, 4}); // nbits
push_val<bool>(buf, true); // is_trained
push_vector<float>(buf, std::vector<float>(5 * 32, 0.0f)); // codebooks
push_val<int>(buf, 0); // search_type = ST_decompress
push_val<float>(buf, 0.0f); // norm_min
push_val<float>(buf, 1.0f); // norm_max
// ProductAdditiveQuantizer:
push_val<size_t>(buf, 2); // nsplits = 2
expect_read_throws_with(buf, "not divisible by nsplits");
}
// -----------------------------------------------------------------------
// Test: ProductResidualQuantizer sub-quantizer d mismatch throws.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, ProductResidualQuantizerSubDMismatch) {
// d=8, nsplits=2 → d_sub should be 4, but sub-quantizer has d=6
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPR");
push_index_header(buf, /*d=*/8, /*ntotal=*/0);
// ProductAdditiveQuantizer: AQ(d=8) + nsplits=2 → codebook_d = 4
push_additive_quantizer(
buf, /*d=*/8, /*M=*/2, /*nbits=*/{4, 4}, /*codebook_d=*/4);
push_val<size_t>(buf, 2); // nsplits = 2
// First sub-quantizer: d=6 instead of expected d_sub=4
push_residual_quantizer(buf, /*d=*/6, /*M=*/1, /*nbits=*/{4});
// (second sub-quantizer not needed — should throw on first)
expect_read_throws_with(buf, "sub-quantizer");
}
// -----------------------------------------------------------------------
// Test: ProductLocalSearchQuantizer sub-quantizer d mismatch throws.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, ProductLocalSearchQuantizerSubDMismatch) {
// d=8, nsplits=2 → d_sub should be 4, but sub-quantizer has d=6
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPL");
push_index_header(buf, /*d=*/8, /*ntotal=*/0);
// ProductAdditiveQuantizer: AQ(d=8) + nsplits=2 → codebook_d = 4
push_additive_quantizer(
buf, /*d=*/8, /*M=*/2, /*nbits=*/{4, 4}, /*codebook_d=*/4);
push_val<size_t>(buf, 2); // nsplits = 2
// First sub-quantizer: d=6 instead of expected d_sub=4
push_local_search_quantizer(buf, /*d=*/6, /*M=*/1, /*nbits=*/{4});
// (second sub-quantizer not needed — should throw on first)
expect_read_throws_with(buf, "sub-quantizer");
}
// -----------------------------------------------------------------------
// Test: ProductResidualQuantizer with nsplits > 1 round-trips correctly
// (codebooks sized for d_sub, not d).
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, ProductResidualQuantizerMultiSplitValid) {
// d=8, nsplits=2, each sub-RQ has d_sub=4, M=1, nbits={4}
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPR");
push_index_header(buf, /*d=*/8, /*ntotal=*/0);
// ProductAdditiveQuantizer: codebook_d = d/nsplits = 4
push_additive_quantizer(
buf, /*d=*/8, /*M=*/2, /*nbits=*/{4, 4}, /*codebook_d=*/4);
push_val<size_t>(buf, 2); // nsplits = 2
// Two sub-RQs with d_sub=4
push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{4});
push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{4});
// code_size = 1 (M=2, nbits={4,4} → 8 bits → 1 byte)
push_val<size_t>(buf, 1); // code_size
push_vector<uint8_t>(buf, {}); // codes (empty, ntotal=0)
VectorIOReader reader;
reader.data = buf;
auto idx = read_index_up(&reader);
EXPECT_EQ(idx->ntotal, 0);
EXPECT_EQ(idx->d, 8);
}
// -----------------------------------------------------------------------
// Test: ProductLocalSearchQuantizer with nsplits > 1 round-trips correctly
// (codebooks sized for d_sub, not d).
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, ProductLocalSearchQuantizerMultiSplitValid) {
// d=8, nsplits=2, each sub-LSQ has d_sub=4, M=1, nbits={4}
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPL");
push_index_header(buf, /*d=*/8, /*ntotal=*/0);
// ProductAdditiveQuantizer: codebook_d = d/nsplits = 4
push_additive_quantizer(
buf, /*d=*/8, /*M=*/2, /*nbits=*/{4, 4}, /*codebook_d=*/4);
push_val<size_t>(buf, 2); // nsplits = 2
// Two sub-LSQs with d_sub=4
push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{4});
push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{4});
// code_size = 1
push_val<size_t>(buf, 1); // code_size
push_vector<uint8_t>(buf, {}); // codes (empty, ntotal=0)
VectorIOReader reader;
reader.data = buf;
auto idx = read_index_up(&reader);
EXPECT_EQ(idx->ntotal, 0);
EXPECT_EQ(idx->d, 8);
}