mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
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:
committed by
meta-codesync[bot]
parent
e96ba2d13c
commit
40480fb303
@@ -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++) {
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user