learningcircuit--local-deep-research/tests/research_library/services/test_faiss_safe_load.py

"""Tests for the restricted-unpickler FAISS loader (faiss_safe_load).

Covers:
- A real ``FAISS.save_local`` index round-trips through ``safe_load_faiss``
  (proves the restricted unpickler accepts legitimate docstore payloads).
- A malicious pickle that would execute code under a plain ``pickle.load`` is
  rejected before execution (proves the RCE is actually closed, including the
  realistic ".pkl-only swap" where the .faiss is left byte-identical).
"""

import datetime
import io
import pickle
from pathlib import Path

import pytest
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

from local_deep_research.research_library.services.faiss_safe_load import (
    _RestrictedFaissUnpickler,
    safe_load_faiss,
)

try:  # location varies across langchain versions
    from langchain_community.embeddings import FakeEmbeddings
except Exception:  # pragma: no cover
    from langchain_core.embeddings import FakeEmbeddings


def _build_and_save(
    dir_path: Path, index_name: str = "index"
) -> FakeEmbeddings:
    """Create a small real FAISS index on disk and return its embeddings."""
    emb = FakeEmbeddings(size=16)
    texts = ["alpha chunk", "beta chunk", "gamma chunk"]
    # Metadata mirrors what the RAG service attaches: plain scalars + lists.
    metas = [
        {"source": "a.pdf", "page": 1, "word_count": 2, "authors": None},
        {"source": "b.pdf", "page": 2, "tags": ["x", "y"]},
        {"source": "c.pdf", "score": 0.9, "nested": {"k": 1}},
    ]
    vs = FAISS.from_texts(texts, emb, metadatas=metas)
    vs.save_local(str(dir_path), index_name=index_name)
    return emb


class TestSafeLoadRoundTrip:
    def test_loads_real_index(self, tmp_path):
        emb = _build_and_save(tmp_path)
        vs = safe_load_faiss(
            str(tmp_path), emb, index_name="index", normalize_L2=True
        )
        contents = {d.page_content for d in vs.docstore._dict.values()}
        assert contents == {"alpha chunk", "beta chunk", "gamma chunk"}
        assert len(vs.index_to_docstore_id) == 3
        # Metadata (scalars, list, nested dict) survives the restricted load.
        metas = [d.metadata for d in vs.docstore._dict.values()]
        assert any(m.get("tags") == ["x", "y"] for m in metas)
        # normalize_L2 is threaded into the FAISS constructor by hand — guard
        # it, since a silent drop would corrupt cosine-distance search.
        assert vs._normalize_L2 is True

    def test_normalize_l2_defaults_false(self, tmp_path):
        emb = _build_and_save(tmp_path)
        vs = safe_load_faiss(str(tmp_path), emb, index_name="index")
        assert vs._normalize_L2 is False

    def test_search_works_after_safe_load(self, tmp_path):
        emb = _build_and_save(tmp_path)
        vs = safe_load_faiss(str(tmp_path), emb, index_name="index")
        results = vs.similarity_search("alpha", k=1)
        assert len(results) == 1


# Module-level helpers so pickle can reference them by qualified name. The
# payload writes a marker file instead of doing real harm, so the test can
# prove (a) it WOULD run under plain pickle and (b) it does NOT under the
# restricted unpickler.
def _write_marker(path: str):
    Path(path).write_text("EXECUTED")
    return InMemoryDocstore({})


class _Evil:
    def __init__(self, marker: str):
        self._marker = marker

    def __reduce__(self):
        return (_write_marker, (self._marker,))


class TestSafeLoadBlocksMalicious:
    def test_restricted_unpickler_blocks_and_does_not_execute(self, tmp_path):
        marker = tmp_path / "marker.txt"
        blob = pickle.dumps((_Evil(str(marker)), {0: "id"}))

        # Sanity: a plain pickle.load DOES execute the payload. This
        # unsafe call is the point of the test — it demonstrates the very
        # RCE that safe_load_faiss prevents, on a benign marker payload.
        pickle.loads(blob)  # noqa: S301
        assert marker.read_text() == "EXECUTED"
        marker.unlink()

        # The restricted unpickler refuses and never runs the payload.
        with pytest.raises(pickle.UnpicklingError):
            _RestrictedFaissUnpickler(io.BytesIO(blob)).load()
        assert not marker.exists()

    def test_safe_load_faiss_rejects_pkl_only_swap(self, tmp_path):
        """The realistic attack: leave .faiss byte-identical, swap the .pkl."""
        emb = _build_and_save(tmp_path)  # valid .faiss + .pkl
        marker = tmp_path / "marker.txt"
        faiss_before = (tmp_path / "index.faiss").read_bytes()

        (tmp_path / "index.pkl").write_bytes(
            pickle.dumps((_Evil(str(marker)), {0: "id"}))
        )

        with pytest.raises(pickle.UnpicklingError):
            safe_load_faiss(str(tmp_path), emb, index_name="index")
        assert not marker.exists()
        # We did not touch the .faiss (mirrors the attack precondition).
        assert (tmp_path / "index.faiss").read_bytes() == faiss_before

    def test_blocks_copyreg_extension_opcode_even_when_cache_warm(
        self, tmp_path
    ):
        """copyreg EXT1/EXT2/EXT4 opcodes resolve a callable via the extension
        registry. ``find_class`` alone is NOT enough: ``get_extension``
        short-circuits on the process-global ``copyreg._extension_cache`` and
        skips ``find_class`` once a code has been resolved. So this test WARMS
        the cache first (a plain ``pickle.loads`` that runs the payload), which
        is exactly the state in which a ``find_class``-only guard would be
        bypassed — then asserts the restricted unpickler still refuses (proving
        the ``get_extension`` override, not ``find_class``, does the blocking).
        """
        import copyreg
        import pickletools

        marker = tmp_path / "marker.txt"
        code = 0xC0DE
        copyreg.add_extension(
            _write_marker.__module__, _write_marker.__name__, code
        )
        try:
            blob = pickle.dumps((_Evil(str(marker)), {0: "id"}), protocol=2)
            # Guard: the payload must actually use an EXT opcode, else this
            # test would pass for the wrong reason.
            op_names = [op.name for op, _, _ in pickletools.genops(blob)]
            assert any(n.startswith("EXT") for n in op_names), op_names

            # Warm copyreg._extension_cache for this code. A find_class-only
            # restricted unpickler would now be bypassed on the EXT path.
            pickle.loads(blob)  # noqa: S301
            assert marker.exists()  # confirms the cache is now warm
            marker.unlink()

            with pytest.raises(pickle.UnpicklingError):
                _RestrictedFaissUnpickler(io.BytesIO(blob)).load()
            assert not marker.exists()
        finally:
            copyreg.remove_extension(
                _write_marker.__module__, _write_marker.__name__, code
            )
            # Drop the warmed cache entry so it can't leak into other tests.
            copyreg._extension_cache.pop(code, None)

    def test_blocks_disallowed_class_buried_in_metadata(self):
        """The allow-list applies to EVERY global the pickle resolves, not
        just the top-level docstore class. A disallowed class hidden inside a
        Document's metadata (here a benign ``datetime``) must still be
        refused — otherwise an attacker could smuggle a gadget through a
        nested value while the outer objects look legitimate.
        """
        doc = Document(
            page_content="x",
            metadata={"created": datetime.datetime(2026, 1, 1)},
        )
        blob = pickle.dumps((InMemoryDocstore({"id1": doc}), {0: "id1"}))

        with pytest.raises(pickle.UnpicklingError, match="datetime"):
            _RestrictedFaissUnpickler(io.BytesIO(blob)).load()