Coverage for src/local_deep_research/research_library/services/faiss_safe_load.py: 100%

1"""Safe loading of on-disk FAISS indexes. 

2 

3LangChain's ``FAISS.load_local(..., allow_dangerous_deserialization=True)`` 

4deserializes the ``<index>.pkl`` docstore companion with a raw 

5``pickle.load`` (see ``langchain_community/vectorstores/faiss.py``). Pickle 

6executes arbitrary code during deserialization, so a tampered ``.pkl`` 

7dropped into the index directory yields arbitrary code execution when the 

8index is next loaded. 

9 

10The ``.faiss`` and ``.pkl`` files are loaded independently with no 

11cross-check, so a content checksum over only the ``.faiss`` file (which is 

12what the file-integrity system records) provides **no** protection against a 

13``.pkl``-only swap. Rather than gate the dangerous load behind a checksum, 

14this module removes the dangerous deserialization entirely: it replicates 

15``load_local`` but unpickles the docstore through a restricted unpickler that 

16only permits the two classes a legitimate FAISS docstore contains 

17(:class:`InMemoryDocstore` and :class:`Document`). Any other global 

18(``os.system``, ``subprocess.Popen``, a crafted ``__reduce__`` payload, …) 

19raises :class:`pickle.UnpicklingError` before it can execute. 

20 

21The payload is the 2-tuple ``(docstore, index_to_docstore_id)`` written by 

22``FAISS.save_local``; the metadata it carries is plain JSON-ish scalars 

23(str/int/float/bool/None/list/dict), none of which require a class global to 

24unpickle. If a future LangChain/Pydantic version changes the on-disk format 

25to need additional classes, loading fails closed (a clear error) rather than 

26silently — and the round-trip test in 

27``tests/vector_stores/test_faiss_safe_load_security.py`` will catch it so 

28the allow-list can be widened deliberately. 

29""" 

30 

31import pickle 

32from typing import Any 

33 

34# (module, qualname) pairs that a legitimate FAISS docstore pickle resolves 

35# via ``Unpickler.find_class``. Determined empirically by round-tripping a 

36# real ``save_local`` payload with varied metadata; only these two appear. 

37_ALLOWED_GLOBALS = frozenset( 

38 { 

39 ("langchain_community.docstore.in_memory", "InMemoryDocstore"), 

40 ("langchain_core.documents.base", "Document"), 

41 } 

42) 

43 

44 

45# Subclass the PURE-PYTHON unpickler (``pickle._Unpickler``), NOT the default 

46# C ``pickle.Unpickler``. We need to override ``get_extension`` to refuse the 

47# copyreg extension opcodes (EXT1/EXT2/EXT4), and only the pure-Python unpickler 

48# routes those opcodes through an overridable ``get_extension`` method — the C 

49# unpickler resolves them internally. (Why ``get_extension`` and not just 

50# ``find_class``: both unpicklers normally route EXT through ``find_class`` on a 

51# cold cache, but ``get_extension`` short-circuits on a warm process-global 

52# ``copyreg._extension_cache`` and skips ``find_class`` entirely. Refusing 

53# ``get_extension`` outright closes that path regardless of cache/registry 

54# state, instead of relying on the extension registry happening to be empty.) 

55# ``_Unpickler`` has been the stable pure-Python name since Python 3.0; if a 

56# future runtime drops it, this fails loudly at import (caught by tests) rather 

57# than silently falling back to the C unpickler. 

58class _RestrictedFaissUnpickler(pickle._Unpickler): 

59 """Unpickler that only resolves the classes a FAISS docstore contains. 

60 

61 Defense is by construction: a malicious pickle can never obtain an 

62 attacker-chosen callable to invoke, because every opcode that resolves a 

63 global or callable is refused or allow-listed: 

64 

65 - ``find_class`` (GLOBAL/STACK_GLOBAL/INST/OBJ, then REDUCE/NEWOBJ/BUILD) 

66 allow-lists only ``InMemoryDocstore`` and ``Document`` — both inert. 

67 - ``get_extension`` (copyreg EXT1/EXT2/EXT4) is refused outright; a 

68 legitimate FAISS docstore pickle never uses copyreg extensions. 

69 - ``persistent_load`` (PERSID/BINPERSID) is refused outright; ``save_local`` 

70 never emits persistent ids. 

71 """ 

72 

73 def find_class(self, module: str, name: str) -> Any: 

74 if (module, name) in _ALLOWED_GLOBALS: 

75 return super().find_class(module, name) 

76 raise pickle.UnpicklingError( 

77 f"Refusing to unpickle disallowed object '{module}.{name}' from " 

78 f"FAISS docstore (possible tampering)." 

79 ) 

80 

81 def get_extension(self, code: int) -> Any: 

82 raise pickle.UnpicklingError( 

83 f"Refusing to resolve copyreg extension code {code} from FAISS " 

84 f"docstore (possible tampering)." 

85 ) 

86 

87 def persistent_load(self, pid: Any) -> Any: 

88 raise pickle.UnpicklingError( 

89 "Refusing to resolve a persistent id from FAISS docstore " 

90 "(possible tampering)." 

91 ) 

92 

93 

94def load_index_to_docstore_id(pkl_path: str) -> dict: 

95 """Read ONLY the ``index_to_docstore_id`` (FAISS position -> docstore key) 

96 map from a legacy ``<index>.pkl``, discarding the docstore. 

97 

98 Uses :class:`_RestrictedFaissUnpickler`, so a tampered 

99 ``.pkl`` cannot execute code on read. The returned dict contains only ints 

100 (FAISS positions) and the original string ids (uuids) — **no document text 

101 or metadata** — which is exactly what the one-time index-format migration 

102 needs to preserve when it strips the plaintext docstore from disk. 

103 

104 The docstore (which holds the chunk text) IS deserialized into memory 

105 transiently by ``pickle`` here, but it is never written back out — only the 

106 id map is returned. 

107 """ 

108 with open(pkl_path, "rb") as f: 

109 _docstore, index_to_docstore_id = _RestrictedFaissUnpickler(f).load() 

110 return index_to_docstore_id