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 code vector sizes during index deserialization (#4948)

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

Add validation that the deserialized codes vector size matches
`ntotal * code_size` for the following index types: IndexPQ,
IndexResidualQuantizer, IndexLocalSearchQuantizer,
IndexProductResidualQuantizer, and IndexProductLocalSearchQuantizer.

A mismatch between the codes vector size and the expected size from the
index header indicates a corrupt or maliciously crafted index file. Without
this check, a truncated or inflated codes vector would cause out-of-bounds
memory access during search or reconstruction operations.

Reviewed By: mnorris11

Differential Revision: D96966390

fbshipit-source-id: 37602d0343f0d79a07487c7f54ed45fd09ae45dd
This commit is contained in:
Justin Gibbs
2026-03-19 20:12:51 -07:00
committed by meta-codesync[bot]
parent d8c3d26130
commit ccc3cd619a
2 changed files with 172 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/impl/index_read.cpp
+10
View File
@@ -1110,6 +1110,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductQuantizer(&idxp->pq, f);
idxp->code_size = idxp->pq.code_size;
read_vector(idxp->codes, f);
FAISS_THROW_IF_NOT(
idxp->codes.size() == idxp->ntotal * idxp->code_size);
if (h == fourcc("IxPo") || h == fourcc("IxPq")) {
READ1(idxp->search_type);
READ1(idxp->encode_signs);
@@ -1132,6 +1134,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
}
READ1(idxr->code_size);
read_vector(idxr->codes, f);
FAISS_THROW_IF_NOT(
idxr->codes.size() == idxr->ntotal * idxr->code_size);
idx = std::move(idxr);
} else if (h == fourcc("IxLS")) {
auto idxr = std::make_unique<IndexLocalSearchQuantizer>();
@@ -1139,6 +1143,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_LocalSearchQuantizer(idxr->lsq, f);
READ1(idxr->code_size);
read_vector(idxr->codes, f);
FAISS_THROW_IF_NOT(
idxr->codes.size() == idxr->ntotal * idxr->code_size);
idx = std::move(idxr);
} else if (h == fourcc("IxPR")) {
auto idxpr = std::make_unique<IndexProductResidualQuantizer>();
@@ -1146,6 +1152,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductResidualQuantizer(idxpr->prq, f, io_flags);
READ1(idxpr->code_size);
read_vector(idxpr->codes, f);
FAISS_THROW_IF_NOT(
idxpr->codes.size() == idxpr->ntotal * idxpr->code_size);
idx = std::move(idxpr);
} else if (h == fourcc("IxPL")) {
auto idxpl = std::make_unique<IndexProductLocalSearchQuantizer>();
@@ -1153,6 +1161,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductLocalSearchQuantizer(idxpl->plsq, f);
READ1(idxpl->code_size);
read_vector(idxpl->codes, f);
FAISS_THROW_IF_NOT(
idxpl->codes.size() == idxpl->ntotal * idxpl->code_size);
idx = std::move(idxpl);
} else if (h == fourcc("ImRQ")) {
auto idxr = std::make_unique<ResidualCoarseQuantizer>();
tests/test_read_index_deserialize.cpp
+162
View File
@@ -938,3 +938,165 @@ TEST(ReadIndexDeserialize, SQTrainedSizeMismatch) {
expect_read_throws_with(buf, "ScalarQuantizer trained size");
}
// -----------------------------------------------------------------------
// Test: IndexPQ codes vector size mismatch.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, IndexPQCodesSizeMismatch) {
// "IxPq": fourcc + index_header + PQ + codes + search_type +
// encode_signs + polysemous_ht
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPq");
push_index_header(buf, /*d=*/4, /*ntotal=*/2);
// PQ: d=4, M=2, nbits=8 -> code_size=2, ksub=256
std::vector<float> centroids(4 * 256, 0.0f);
push_pq(buf, /*d=*/4, /*M=*/2, /*nbits=*/8, centroids);
// codes: should be ntotal * code_size = 2 * 2 = 4 bytes
// Provide wrong size (10 bytes)
push_vector<uint8_t>(buf, std::vector<uint8_t>(10, 0));
// IxPq format extras:
push_val<int>(buf, 0); // search_type
push_val<int>(buf, 0); // encode_signs
push_val<int>(buf, 0); // polysemous_ht
expect_read_throws_with(buf, "codes.size()");
}
/// Helper: append a minimal AdditiveQuantizer (d, M, nbits, is_trained,
/// codebooks, search_type=ST_decompress, norm_min, norm_max).
static void push_additive_quantizer(
std::vector<uint8_t>& buf,
size_t d,
size_t M,
const std::vector<size_t>& nbits) {
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
}
/// Helper: append a minimal ResidualQuantizer (AdditiveQuantizer +
/// train_type + max_beam_size). Uses Skip_codebook_tables to avoid
/// compute_codebook_tables on empty codebooks.
static void push_residual_quantizer(
std::vector<uint8_t>& buf,
size_t d,
size_t M,
const std::vector<size_t>& nbits) {
push_additive_quantizer(buf, d, M, nbits);
push_val<int>(buf, 2048); // train_type = Skip_codebook_tables
push_val<int>(buf, 1); // max_beam_size
}
/// Helper: append a minimal LocalSearchQuantizer (AdditiveQuantizer +
/// K + train_iters + encode_ils_iters + train_ils_iters + icm_iters +
/// p + lambd + chunk_size + random_seed + nperts +
/// update_codebooks_with_double).
static void push_local_search_quantizer(
std::vector<uint8_t>& buf,
size_t d,
size_t M,
const std::vector<size_t>& nbits) {
push_additive_quantizer(buf, d, M, nbits);
push_val<size_t>(buf, 256); // K
push_val<size_t>(buf, 25); // train_iters
push_val<size_t>(buf, 8); // encode_ils_iters
push_val<size_t>(buf, 8); // train_ils_iters
push_val<size_t>(buf, 4); // icm_iters
push_val<float>(buf, 0.5f); // p
push_val<float>(buf, 0.0f); // lambd
push_val<size_t>(buf, 0); // chunk_size
push_val<int>(buf, 123); // random_seed
push_val<size_t>(buf, 4); // nperts
push_val<bool>(buf, false); // update_codebooks_with_double
}
// -----------------------------------------------------------------------
// Test: IndexResidualQuantizer codes vector size mismatch.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, IndexResidualQuantizerCodesSizeMismatch) {
// "IxRq": fourcc + index_header + ResidualQuantizer + code_size + codes
std::vector<uint8_t> buf;
push_fourcc(buf, "IxRq");
push_index_header(buf, /*d=*/4, /*ntotal=*/2);
push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
// code_size = 1 (M=1, nbits=8 → 1 byte per code)
push_val<size_t>(buf, 1); // code_size
// codes: should be ntotal * code_size = 2 * 1 = 2 bytes
// Provide wrong size (10 bytes)
push_vector<uint8_t>(buf, std::vector<uint8_t>(10, 0));
expect_read_throws_with(buf, "codes.size()");
}
// -----------------------------------------------------------------------
// Test: IndexLocalSearchQuantizer codes vector size mismatch.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, IndexLocalSearchQuantizerCodesSizeMismatch) {
// "IxLS": fourcc + index_header + LocalSearchQuantizer + code_size + codes
std::vector<uint8_t> buf;
push_fourcc(buf, "IxLS");
push_index_header(buf, /*d=*/4, /*ntotal=*/2);
push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
// code_size = 1 (M=1, nbits=8 → 1 byte per code)
push_val<size_t>(buf, 1); // code_size
// codes: should be ntotal * code_size = 2 * 1 = 2 bytes
// Provide wrong size (10 bytes)
push_vector<uint8_t>(buf, std::vector<uint8_t>(10, 0));
expect_read_throws_with(buf, "codes.size()");
}
// -----------------------------------------------------------------------
// Test: IndexProductResidualQuantizer codes vector size mismatch.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, IndexProductResidualQuantizerCodesSizeMismatch) {
// "IxPR": fourcc + index_header + ProductResidualQuantizer + code_size
// + codes
// ProductResidualQuantizer = ProductAdditiveQuantizer(AdditiveQuantizer
// + nsplits) + nsplits * ResidualQuantizer
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPR");
push_index_header(buf, /*d=*/4, /*ntotal=*/2);
// ProductAdditiveQuantizer: AdditiveQuantizer + nsplits
push_additive_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
push_val<size_t>(buf, 1); // nsplits = 1
// 1 nested ResidualQuantizer (sub-dimension = d / nsplits = 4)
push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
// code_size = 1
push_val<size_t>(buf, 1); // code_size
// codes: should be ntotal * code_size = 2 * 1 = 2 bytes
// Provide wrong size (10 bytes)
push_vector<uint8_t>(buf, std::vector<uint8_t>(10, 0));
expect_read_throws_with(buf, "codes.size()");
}
// -----------------------------------------------------------------------
// Test: IndexProductLocalSearchQuantizer codes vector size mismatch.
// -----------------------------------------------------------------------
TEST(ReadIndexDeserialize, IndexProductLocalSearchQuantizerCodesSizeMismatch) {
// "IxPL": fourcc + index_header + ProductLocalSearchQuantizer + code_size
// + codes
// ProductLocalSearchQuantizer = ProductAdditiveQuantizer(AdditiveQuantizer
// + nsplits) + nsplits * LocalSearchQuantizer
std::vector<uint8_t> buf;
push_fourcc(buf, "IxPL");
push_index_header(buf, /*d=*/4, /*ntotal=*/2);
// ProductAdditiveQuantizer: AdditiveQuantizer + nsplits
push_additive_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
push_val<size_t>(buf, 1); // nsplits = 1
// 1 nested LocalSearchQuantizer (sub-dimension = d / nsplits = 4)
push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8});
// code_size = 1
push_val<size_t>(buf, 1); // code_size
// codes: should be ntotal * code_size = 2 * 1 = 2 bytes
// Provide wrong size (10 bytes)
push_vector<uint8_t>(buf, std::vector<uint8_t>(10, 0));
expect_read_throws_with(buf, "codes.size()");
}