Support limits on index deserialization loops — useful for tests (#4902)

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

When enabled, throw a FAISS exception when index deserialization loops
driven by read data fields (nlist, nsplits, VT chain length, nhash, etc.)
exceed a configured limit.  This can be set by input validation tests to
prevent wasting time on pathological, OOM inducting, inputs.

Default: 0 (no limit).

Reviewed By: mnorris11

Differential Revision: D96016580

fbshipit-source-id: 8a04fdfe5b1d9aecaf1a7156215045be011014e5
This commit is contained in:
Justin Gibbs
2026-03-11 22:06:48 -07:00
committed by meta-codesync[bot]
parent e96ba2d13c
commit 40480fb303
2 changed files with 50 additions and 0 deletions
+39
View File
@@ -78,6 +78,32 @@
namespace faiss {
namespace {
size_t deserialization_loop_limit_ = 0;
} // namespace
size_t get_deserialization_loop_limit() {
return deserialization_loop_limit_;
}
void set_deserialization_loop_limit(size_t value) {
deserialization_loop_limit_ = value;
}
#define FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(val, field_name) \
do { \
auto limit_ = get_deserialization_loop_limit(); \
if (limit_ > 0) { \
FAISS_THROW_IF_NOT_FMT( \
static_cast<size_t>(val) <= limit_, \
"%s=%zd exceeds deserialization_loop_limit" \
" of %zd", \
field_name, \
static_cast<size_t>(val), \
limit_); \
} \
} while (0)
/*************************************************************
* Mmap-ing and viewing facilities
**************************************************************/
@@ -371,6 +397,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
} else if (h == fourcc("ilpn") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
size_t nlist, code_size, n_levels;
READ1(nlist);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nlist, "ilpn nlist");
READ1(code_size);
READ1(n_levels);
auto ailp = std::make_unique<ArrayInvertedListsPanorama>(
@@ -400,6 +427,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
} else if (h == fourcc("ilar") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
auto ails = std::make_unique<ArrayInvertedLists>(0, 0);
READ1(ails->nlist);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ails->nlist, "ilar nlist");
READ1(ails->code_size);
ails->ids.resize(ails->nlist);
ails->codes.resize(ails->nlist);
@@ -430,6 +458,7 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
int h2 = (io_flags & 0xffff0000) | (fourcc("il__") & 0x0000ffff);
size_t nlist, code_size;
READ1(nlist);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nlist, "ilar skip nlist");
READ1(code_size);
std::vector<size_t> sizes(nlist);
read_ArrayInvertedLists_sizes(f, sizes);
@@ -557,6 +586,7 @@ static void read_ProductAdditiveQuantizer(
paq.nsplits > 0,
"invalid ProductAdditiveQuantizer nsplits %zd (must be > 0)",
paq.nsplits);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(paq.nsplits, "nsplits");
}
static void read_ProductResidualQuantizer(
@@ -735,6 +765,8 @@ static void read_HNSW(HNSW& hnsw, IOReader* f) {
static void read_NSG(NSG& nsg, IOReader* f) {
READ1(nsg.ntotal);
READ1(nsg.R);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nsg.ntotal, "nsg.ntotal");
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(nsg.R, "nsg.R");
FAISS_THROW_IF_NOT_FMT(nsg.R > 0, "invalid NSG R %d (must be > 0)", nsg.R);
READ1(nsg.L);
READ1(nsg.C);
@@ -858,6 +890,7 @@ void read_ivf_header(
std::vector<std::vector<idx_t>>* ids) {
read_index_header(*ivf, f);
READ1(ivf->nlist);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ivf->nlist, "nlist");
READ1(ivf->nprobe);
ivf->quantizer = read_index(f);
ivf->own_fields = true;
@@ -1368,6 +1401,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
nt >= 0,
"invalid VectorTransform chain length %d (must be >= 0)",
nt);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(
nt, "VectorTransform chain length");
for (int i = 0; i < nt; i++) {
ixpt->chain.push_back(read_VectorTransform(f));
}
@@ -1898,6 +1933,7 @@ static void read_binary_ivf_header(
std::vector<std::vector<idx_t>>* ids = nullptr) {
read_index_binary_header(ivf, f);
READ1(ivf.nlist);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(ivf.nlist, "nlist");
READ1(ivf.nprobe);
ivf.quantizer = read_index_binary(f);
ivf.own_fields = true;
@@ -1915,6 +1951,7 @@ static void read_binary_hash_invlists(
IOReader* f) {
size_t sz;
READ1(sz);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(sz, "binary hash invlists sz");
int il_nbit = 0;
READ1(il_nbit);
FAISS_THROW_IF_NOT_FMT(
@@ -1965,6 +2002,7 @@ static void read_binary_multi_hash_map(
size_t sz;
READ1(id_bits);
READ1(sz);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(sz, "multi hash map sz");
std::vector<uint8_t> buf;
READVECTOR(buf);
size_t nbit = add_no_overflow(
@@ -2063,6 +2101,7 @@ std::unique_ptr<IndexBinary> read_index_binary_up(IOReader* f, int io_flags) {
idxmh->nhash > 0,
"invalid IndexBinaryMultiHash nhash %d (must be > 0)",
idxmh->nhash);
FAISS_CHECK_DESERIALIZATION_LOOP_LIMIT(idxmh->nhash, "nhash");
READ1(idxmh->nflip);
idxmh->maps.resize(idxmh->nhash);
for (int i = 0; i < idxmh->nhash; i++) {
+11
View File
@@ -113,6 +113,17 @@ std::unique_ptr<InvertedLists> read_InvertedLists_up(
IOReader* reader,
int io_flags = 0);
// Returns the current deserialization loop limit.
// When nonzero, deserialization rejects loop-driving fields (nlist,
// nsplits, VT chain length, nhash, etc.) that exceed this value.
// Default: 0 (no limit).
size_t get_deserialization_loop_limit();
// Sets the deserialization loop limit.
// NOT thread-safe: set before any concurrent deserialization calls
// and do not modify while deserialization is in progress on other threads.
void set_deserialization_loop_limit(size_t value);
} // namespace faiss
#endif