diff --git a/src/memos/mem_reader/read_skill_memory/upload_skill_memory.py b/src/memos/mem_reader/read_skill_memory/upload_skill_memory.py index fef6b60f2..cfcf2d9d9 100644 --- a/src/memos/mem_reader/read_skill_memory/upload_skill_memory.py +++ b/src/memos/mem_reader/read_skill_memory/upload_skill_memory.py @@ -1,9 +1,10 @@ +import os import re import shutil import tempfile import zipfile -from pathlib import Path +from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any from urllib.parse import urlparse from uuid import uuid4 @@ -91,6 +92,63 @@ def _download_zip(url: str, tmp_dir: Path) -> Path: return zip_path +# Unix mode bits for a symlink (S_IFLNK). +_SYMLINK_MODE = 0o120000 + + +def _safe_extract_zip(zf: zipfile.ZipFile, extract_dir: Path) -> None: + """Extract ``zf`` into ``extract_dir`` after validating every entry. + + Defense against zip-slip / path-traversal (CWE-22) and symlink escape. + Any hostile entry causes the entire extraction to be rejected with a + ``ValueError`` — we do not silently skip, because the caller path is + reachable from the unauthenticated ``/product`` surface and a partial + extraction produces a subtly-wrong SkillMemory that hides the attack. + + Rules per entry: + * name must not be an absolute path (POSIX or NT style) + * name must not resolve outside ``extract_dir`` after normalization + (catches ``..`` traversal and Windows drive prefixes) + * external_attr must not encode a symlink (S_IFLNK) + + Called by :func:`_extract_and_parse_skill_zip` in place of + ``zf.extractall(extract_dir)``. + """ + base = extract_dir.resolve() + base.mkdir(parents=True, exist_ok=True) + + for info in zf.infolist(): + name = info.filename + + # Symlink entry — reject before any I/O. + unix_mode = (info.external_attr >> 16) & 0o170000 + if unix_mode == _SYMLINK_MODE: + raise ValueError(f"Refusing to extract symlink entry from skill zip: {name!r}") + + # Absolute-path entry — reject. os.path.isabs on POSIX only + # recognises leading "/", so we also cross-check with pure-path + # variants for both POSIX ("/foo") and NT ("C:\\foo", "\\\\srv\\s") + # styles regardless of the host OS. The downstream ``base not in + # candidate.parents`` guard would catch drive-relative paths, but + # a fully platform-independent pre-filter here means the guard + # cannot be silently defeated if future code paths relax it. + if ( + not name + or os.path.isabs(name) + or PurePosixPath(name).is_absolute() + or PureWindowsPath(name).is_absolute() + ): + raise ValueError(f"Refusing to extract absolute-path entry from skill zip: {name!r}") + + # Compute the would-be destination and check containment. + candidate = (base / name).resolve() + if candidate != base and base not in candidate.parents: + raise ValueError(f"Refusing to extract entry outside extract_dir: {name!r}") + + # All entries validated — extract. + zf.extractall(base) + + def _extract_and_parse_skill_zip(zip_path: Path) -> dict[str, Any]: """ Extract a skill zip and parse SKILL.md + directory contents into a skill_memory dict. @@ -102,7 +160,7 @@ def _extract_and_parse_skill_zip(zip_path: Path) -> dict[str, Any]: # Step 1: extract & locate SKILL.md extract_dir = zip_path.parent / zip_path.stem with zipfile.ZipFile(zip_path, "r") as zf: - zf.extractall(extract_dir) + _safe_extract_zip(zf, extract_dir) skill_md_path = None for candidate in extract_dir.rglob("SKILL.md"): diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..aa0fd4f7d 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -8,11 +8,55 @@ from memos.configs.memory import KVCacheMemoryConfig from memos.dependency import require_python_package from memos.llms.factory import LLMFactory +from memos.log import get_logger from memos.memories.activation.base import BaseActMemory from memos.memories.activation.item import KVCacheItem +from memos.memories.activation.safe_unpickler import ( + _BASE_ALLOWED_CLASSES, + _BaseSafeUnpickler, +) from memos.memories.textual.item import TextualMemoryItem +logger = get_logger(__name__) + + +# Extra classes KVCacheMemory.dump() produces on top of the shared base +# allowlist: +# * KVCacheItem instances +# * DynamicCache (transformers) with tensor state (torch rebuild helpers) +# +# ``torch.storage._load_from_bytes`` is *deliberately* NOT in this list — +# it is a known pickle-allowlist bypass vector: it internally calls +# ``pickle.loads`` on its bytes argument with the standard (unrestricted) +# unpickler, which would defeat the ``find_class`` guard for any nested +# payload. See the discussion on PR #2204. +_KV_ALLOWED_CLASSES: frozenset[tuple[str, str]] = _BASE_ALLOWED_CLASSES | frozenset( + { + ("transformers.cache_utils", "DynamicCache"), + ("transformers.cache_utils", "DynamicLayer"), + ("memos.memories.activation.item", "KVCacheItem"), + ("memos.memories.activation.item", "KVCacheRecords"), + # torch tensor rebuilders — required to re-materialize DynamicCache + ("torch._utils", "_rebuild_tensor_v2"), + ("torch._utils", "_rebuild_tensor"), + ("torch._utils", "_rebuild_parameter"), + ("torch._utils", "_rebuild_qtensor"), + ("torch._utils", "_rebuild_meta_tensor_no_storage"), + ("torch", "Tensor"), + ("torch", "device"), + ("torch", "dtype"), + ("torch", "Size"), + } +) + + +class _SafeUnpickler(_BaseSafeUnpickler): + """Restricted pickle.Unpickler for the KV activation cache.""" + + _allowed_classes = _KV_ALLOWED_CLASSES + + class KVCacheMemory(BaseActMemory): """ Key-Value Cache Memory for activation memories. @@ -155,7 +199,10 @@ def load(self, dir: str) -> None: torch.serialization.add_safe_globals([DynamicCache, KVCacheItem]) with open(file_path, "rb") as f: - data = pickle.load(f) + # Restricted unpickler — rejects any class outside the + # KVCache allowlist before invoking any reduce callable. + # See _KV_ALLOWED_CLASSES / _SafeUnpickler above. + data = _SafeUnpickler(f).load() if isinstance(data, dict): # Load memories, handle both old and new formats @@ -176,8 +223,24 @@ def load(self, dir: str) -> None: # Reset to empty if data format is unexpected self.kv_cache_memories = {} - except (EOFError, pickle.UnpicklingError, Exception): - # If loading fails, start with empty memories + except pickle.UnpicklingError as e: + # The safe unpickler refused a class — likely a hostile cache + # file. Log at WARN level so operators see it, then reset. + logger.warning( + "[KVCacheMemory] Refused to load activation cache (%s); resetting.", + e, + ) + self.kv_cache_memories = {} + except (EOFError, OSError, ValueError) as e: + # Truncated file, I/O failure, or malformed pickle stream — + # log so operators can investigate silent data-loss cases, + # then reset. Do NOT catch bare Exception here: unexpected + # errors (MemoryError, AttributeError, etc.) should surface. + logger.warning( + "[KVCacheMemory] Failed to load activation cache (%s: %s); resetting.", + type(e).__name__, + e, + ) self.kv_cache_memories = {} def dump(self, dir: str) -> None: diff --git a/src/memos/memories/activation/safe_unpickler.py b/src/memos/memories/activation/safe_unpickler.py new file mode 100644 index 000000000..c50f64dd5 --- /dev/null +++ b/src/memos/memories/activation/safe_unpickler.py @@ -0,0 +1,61 @@ +"""Shared restricted-unpickler infrastructure for activation caches. + +Issue #2203: raw ``pickle.load`` on an activation-cache file is a CWE-502 +sink. Both :mod:`memos.memories.activation.kv` and +:mod:`memos.memories.activation.vllmkv` need a restricted unpickler that +enforces a per-cache allowlist at ``find_class`` time (before any reduce +callable runs). + +To avoid drift between the two caches, the common allowlist and the base +``_SafeUnpickler`` class live here. Each cache module extends the base +allowlist with its own extra entries and passes the merged set to a +subclass of :class:`_BaseSafeUnpickler`. +""" + +from __future__ import annotations + +import pickle + + +# Classes that both KV and vLLM caches always need: the ``dict`` wrapper +# around ``kv_cache_memories``, common stdlib containers, and +# ``datetime.*`` for metadata timestamps. +_BASE_ALLOWED_CLASSES: frozenset[tuple[str, str]] = frozenset( + { + ("builtins", "dict"), + ("builtins", "list"), + ("builtins", "tuple"), + ("builtins", "set"), + ("builtins", "frozenset"), + ("builtins", "str"), + ("builtins", "int"), + ("builtins", "float"), + ("builtins", "bool"), + ("builtins", "bytes"), + ("collections", "OrderedDict"), + ("datetime", "datetime"), + ("datetime", "date"), + ("datetime", "time"), + ("datetime", "timedelta"), + ("datetime", "timezone"), + } +) + + +class _BaseSafeUnpickler(pickle.Unpickler): + """Restricted :class:`pickle.Unpickler` with a class-attribute allowlist. + + Subclasses must define ``_allowed_classes: frozenset[tuple[str, str]]``. + Any class outside the allowlist is rejected at ``find_class`` time, + before the class is imported and before any reduce callable is + invoked. This blocks the CWE-502 sink documented in issue #2203. + """ + + _allowed_classes: frozenset[tuple[str, str]] = frozenset() + + def find_class(self, module: str, name: str): # type: ignore[override] + if (module, name) not in self._allowed_classes: + raise pickle.UnpicklingError( + f"Refusing to load class {module}.{name} from activation cache" + ) + return super().find_class(module, name) diff --git a/src/memos/memories/activation/vllmkv.py b/src/memos/memories/activation/vllmkv.py index 4b74115f4..452d9b02e 100644 --- a/src/memos/memories/activation/vllmkv.py +++ b/src/memos/memories/activation/vllmkv.py @@ -6,11 +6,36 @@ from memos.configs.memory import KVCacheMemoryConfig from memos.dependency import require_python_package from memos.llms.factory import LLMFactory +from memos.log import get_logger from memos.memories.activation.base import BaseActMemory from memos.memories.activation.item import VLLMKVCacheItem +from memos.memories.activation.safe_unpickler import ( + _BASE_ALLOWED_CLASSES, + _BaseSafeUnpickler, +) from memos.memories.textual.item import TextualMemoryItem +logger = get_logger(__name__) + + +# Extra classes VLLMKVCacheMemory.dump() produces on top of the shared +# base allowlist. VLLMKVCacheItem.memory is a plain str (the preloaded +# prompt), so no torch tensor rebuilders are needed here. +_VLLM_ALLOWED_CLASSES: frozenset[tuple[str, str]] = _BASE_ALLOWED_CLASSES | frozenset( + { + ("memos.memories.activation.item", "VLLMKVCacheItem"), + ("memos.memories.activation.item", "KVCacheRecords"), + } +) + + +class _SafeUnpickler(_BaseSafeUnpickler): + """Restricted pickle.Unpickler for the vLLM activation cache.""" + + _allowed_classes = _VLLM_ALLOWED_CLASSES + + class VLLMKVCacheMemory(BaseActMemory): """ VLLM Key-Value Cache Memory for activation memories. @@ -155,13 +180,14 @@ def load(self, dir: str) -> None: return try: - # Allow loading VLLMKVCacheItem types - import torch - - torch.serialization.add_safe_globals([VLLMKVCacheItem]) + # torch.serialization.add_safe_globals is not needed here: + # _SafeUnpickler enforces _VLLM_ALLOWED_CLASSES directly and + # does not delegate to torch's safe-globals registry. with open(file_path, "rb") as f: - data = pickle.load(f) + # Restricted unpickler — rejects any class outside the + # vLLM allowlist before invoking any reduce callable. + data = _SafeUnpickler(f).load() if isinstance(data, dict): # Load memories, handle both old and new formats @@ -182,8 +208,23 @@ def load(self, dir: str) -> None: # Reset to empty if data format is unexpected self.kv_cache_memories = {} - except (EOFError, pickle.UnpicklingError, Exception): - # If loading fails, start with empty memories + except pickle.UnpicklingError as e: + # Safe unpickler refused a class — likely a hostile cache. + logger.warning( + "[VLLMKVCacheMemory] Refused to load activation cache (%s); resetting.", + e, + ) + self.kv_cache_memories = {} + except (EOFError, OSError, ValueError) as e: + # Truncated file, I/O failure, or malformed pickle stream — + # log so operators can investigate silent data-loss cases. + # Do NOT catch bare Exception here: unexpected errors + # (MemoryError, AttributeError, etc.) should surface. + logger.warning( + "[VLLMKVCacheMemory] Failed to load activation cache (%s: %s); resetting.", + type(e).__name__, + e, + ) self.kv_cache_memories = {} def dump(self, dir: str) -> None: diff --git a/tests/mem_reader/read_skill_memory/__init__.py b/tests/mem_reader/read_skill_memory/__init__.py new file mode 100644 index 000000000..d565db704 --- /dev/null +++ b/tests/mem_reader/read_skill_memory/__init__.py @@ -0,0 +1 @@ +"""Marker so pytest treats this directory as a package.""" diff --git a/tests/mem_reader/read_skill_memory/test_safe_extract.py b/tests/mem_reader/read_skill_memory/test_safe_extract.py new file mode 100644 index 000000000..64bf353b2 --- /dev/null +++ b/tests/mem_reader/read_skill_memory/test_safe_extract.py @@ -0,0 +1,207 @@ +"""Regression tests for zip-slip / path-traversal hardening in +``_extract_and_parse_skill_zip``. + +Issue #2203: ``zipfile.ZipFile.extractall`` was called without per-entry +path validation. A crafted zip could write files outside the requested +extraction directory. These tests build hostile zips in-memory and assert +that the safe-extract helper rejects them with ``ValueError``. +""" + +from __future__ import annotations + +import zipfile + +from typing import TYPE_CHECKING + +import pytest + +from memos.mem_reader.read_skill_memory.upload_skill_memory import ( + _extract_and_parse_skill_zip, + _safe_extract_zip, +) + + +if TYPE_CHECKING: + from pathlib import Path + + +def _write_zip(tmp_path: Path, name: str, entries: list[tuple[str, bytes]]) -> Path: + """Write a zip file with the given (arcname, data) entries.""" + zip_path = tmp_path / name + with zipfile.ZipFile(zip_path, "w") as zf: + for arcname, data in entries: + zf.writestr(arcname, data) + return zip_path + + +def _write_zip_with_symlink(tmp_path: Path, name: str, link_name: str, target: str) -> Path: + """Write a zip file containing a symlink entry pointing to ``target``.""" + zip_path = tmp_path / name + # Build ZipInfo with the symlink mode (0o120000 in the high bits of + # external_attr, ignore the trailing perm bits). + with zipfile.ZipFile(zip_path, "w") as zf: + info = zipfile.ZipInfo(link_name) + # UPPER 16 bits = unix mode; 0o120000 = symlink. + info.external_attr = (0o120777 & 0xFFFF) << 16 + zf.writestr(info, target) + return zip_path + + +# --------------------------------------------------------------------------- +# _safe_extract_zip — direct unit tests +# --------------------------------------------------------------------------- + + +def test_safe_extract_rejects_parent_traversal(tmp_path: Path) -> None: + """An entry with `../` should raise ValueError, no file created.""" + zip_path = _write_zip( + tmp_path, + "malicious.zip", + [("../escaped.txt", b"poc")], + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf, pytest.raises(ValueError): + _safe_extract_zip(zf, extract_dir) + assert not (tmp_path / "escaped.txt").exists() + + +def test_safe_extract_rejects_deep_traversal(tmp_path: Path) -> None: + """Multiple `..` segments should also be rejected.""" + zip_path = _write_zip( + tmp_path, + "malicious.zip", + [("../../../../../etc/escaped_canary", b"poc")], + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf, pytest.raises(ValueError): + _safe_extract_zip(zf, extract_dir) + + +def test_safe_extract_rejects_absolute_path(tmp_path: Path) -> None: + """An absolute-path entry should be rejected.""" + # Use the tmp_path/canary as the "absolute" target so we can assert it + # was not created (rather than picking /tmp/... which might already + # exist as a file the test runner cannot write to). + canary = tmp_path / "canary_should_not_exist.txt" + zip_path = _write_zip( + tmp_path, + "malicious.zip", + # We fabricate an absolute name manually with a ZipInfo so the + # zipfile library does not sanitize it in transit. + [], + ) + with zipfile.ZipFile(zip_path, "a") as zf: + info = zipfile.ZipInfo(str(canary)) + zf.writestr(info, b"poc") + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf, pytest.raises(ValueError): + _safe_extract_zip(zf, extract_dir) + assert not canary.exists() + + +def test_safe_extract_rejects_symlink_entry(tmp_path: Path) -> None: + """A symlink zip entry should be rejected outright.""" + zip_path = _write_zip_with_symlink( + tmp_path, + "malicious.zip", + link_name="link_to_shadow", + target="/etc/shadow", + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf, pytest.raises(ValueError): + _safe_extract_zip(zf, extract_dir) + + +def test_safe_extract_accepts_legitimate_zip(tmp_path: Path) -> None: + """A well-formed zip must still extract successfully.""" + zip_path = _write_zip( + tmp_path, + "legit.zip", + [ + ("SKILL.md", b"---\nname: demo\n---\n# Trigger\ndo the thing\n"), + ("scripts/hello.py", b"print('hi')\n"), + ("reference/notes.md", b"see also\n"), + ], + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf: + _safe_extract_zip(zf, extract_dir) + + assert ( + extract_dir / "SKILL.md" + ).read_text() == "---\nname: demo\n---\n# Trigger\ndo the thing\n" + assert (extract_dir / "scripts" / "hello.py").exists() + assert (extract_dir / "reference" / "notes.md").exists() + + +# --------------------------------------------------------------------------- +# _extract_and_parse_skill_zip — end-to-end guard +# --------------------------------------------------------------------------- + + +def test_extract_and_parse_rejects_malicious_zip(tmp_path: Path) -> None: + """The high-level parser must propagate the safe-extract rejection.""" + zip_path = _write_zip( + tmp_path, + "malicious.zip", + [ + ("SKILL.md", b"---\nname: demo\n---\n# Trigger\nx\n"), + ("../escaped.txt", b"poc"), + ], + ) + with pytest.raises(ValueError): + _extract_and_parse_skill_zip(zip_path) + assert not (tmp_path / "escaped.txt").exists() + + +def test_extract_and_parse_accepts_legitimate_zip(tmp_path: Path) -> None: + """Well-formed skill zip must parse into a skill_memory dict.""" + zip_path = _write_zip( + tmp_path, + "legit.zip", + [ + ( + "SKILL.md", + b"---\nname: demo-skill\ndescription: a demo\n---\n" + b"# Trigger\nsome trigger\n\n# Procedure\nsteps here\n", + ), + ], + ) + skill = _extract_and_parse_skill_zip(zip_path) + assert skill["name"] == "demo-skill" + assert skill["description"] == "a demo" + assert skill["procedure"].startswith("steps") + + +def test_safe_extract_when_symlink_ext_attr_is_ignored_on_windows(tmp_path: Path) -> None: + """Defensive: even on platforms where symlink bits are ignored during + extractall, the pre-check must still reject.""" + # Same shape as test_safe_extract_rejects_symlink_entry but re-asserts + # nothing landed in extract_dir if the check fires. + zip_path = _write_zip_with_symlink( + tmp_path, + "malicious.zip", + link_name="link_to_target", + target=str(tmp_path / "nonexistent"), + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf, pytest.raises(ValueError): + _safe_extract_zip(zf, extract_dir) + # extract_dir should not contain the symlink entry + if extract_dir.exists(): + assert not any(extract_dir.iterdir()) + # Also assert we did not follow the target + assert not (tmp_path / "nonexistent").exists() + + +def test_safe_extract_permits_dotdot_inside_name(tmp_path: Path) -> None: + """Names like `foo..bar` (no path separator) must still work.""" + zip_path = _write_zip( + tmp_path, + "legit.zip", + [("file..name.txt", b"ok")], + ) + extract_dir = tmp_path / "sandbox" + with zipfile.ZipFile(zip_path, "r") as zf: + _safe_extract_zip(zf, extract_dir) + assert (extract_dir / "file..name.txt").read_bytes() == b"ok" diff --git a/tests/memories/activation/test_safe_pickle_load.py b/tests/memories/activation/test_safe_pickle_load.py new file mode 100644 index 000000000..edc922d21 --- /dev/null +++ b/tests/memories/activation/test_safe_pickle_load.py @@ -0,0 +1,168 @@ +"""Regression tests for the restricted-unpickler hardening in +``KVCacheMemory.load`` and ``VLLMKVCacheMemory.load``. + +Issue #2203: raw ``pickle.load`` on the activation-cache file is an RCE +sink. These tests craft pickle streams whose ``__reduce__`` names a +disallowed class (``os.system``) and assert the unpickler rejects them +without executing the payload. +""" + +from __future__ import annotations + +import os +import pickle +import shlex + +from typing import TYPE_CHECKING +from unittest.mock import MagicMock + +import pytest + + +if TYPE_CHECKING: + from pathlib import Path + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _OsSystemPayload: + """__reduce__ returns (os.system, ("touch ",)). Loading this + via a raw ``pickle.load`` writes the canary; loading via SafeUnpickler + must raise before ``os.system`` is invoked.""" + + def __init__(self, canary_path: str) -> None: + self.canary_path = canary_path + + def __reduce__(self): # type: ignore[override] + # Shell-quote the path: pytest's ``tmp_path`` can contain spaces + # on CI configurations (e.g. ``/tmp/pytest-of-user name/...``), + # which would silently split into multiple ``touch`` arguments + # and fail — making the security assertion below indistinguishable + # from a genuine SafeUnpickler regression. + return (os.system, (f"touch {shlex.quote(self.canary_path)}",)) + + +def _write_hostile_pickle(target: Path, canary_path: str) -> None: + with open(target, "wb") as f: + pickle.dump(_OsSystemPayload(canary_path), f) + + +# --------------------------------------------------------------------------- +# kv.py — KVCacheMemory +# --------------------------------------------------------------------------- + + +@pytest.fixture +def kv_memory(monkeypatch): + from memos.configs.memory import KVCacheMemoryConfig + from memos.llms import factory as llm_factory + from memos.memories.activation.kv import KVCacheMemory + + monkeypatch.setattr( + llm_factory.LLMFactory, + "from_config", + lambda cfg: MagicMock(build_kv_cache=lambda x: None), + ) + config = MagicMock(spec=KVCacheMemoryConfig) + config.extractor_llm = MagicMock() + config.memory_filename = "kv_cache.pkl" + return KVCacheMemory(config) + + +def test_kv_load_refuses_hostile_pickle(tmp_path: Path, kv_memory) -> None: + """os.system-based __reduce__ must be rejected; canary must not exist.""" + canary = tmp_path / "kv_canary" + _write_hostile_pickle(tmp_path / kv_memory.config.memory_filename, str(canary)) + + kv_memory.load(str(tmp_path)) + + assert not canary.exists(), "SafeUnpickler failed to block os.system payload" + assert kv_memory.kv_cache_memories == {} + + +def test_kv_load_refuses_arbitrary_class(tmp_path: Path, kv_memory) -> None: + """A pickle referencing a random not-allowlisted class must be rejected.""" + import io + + # Pickle referencing `subprocess.Popen` via a bare name; we don't + # actually instantiate — the find_class hook must reject at load time. + raw = pickle.dumps({"kv_cache_memories": []}, protocol=pickle.HIGHEST_PROTOCOL) + # sanity: with allowlist this legit pickle must load + from memos.memories.activation.kv import _SafeUnpickler as SafeUnpickler + + obj = SafeUnpickler(io.BytesIO(raw)).load() + assert isinstance(obj, dict) + assert obj["kv_cache_memories"] == [] + + # now craft one that names os.system — must raise + stream = pickle.dumps( + _OsSystemPayload(str(tmp_path / "should_not_run")), + protocol=pickle.HIGHEST_PROTOCOL, + ) + with pytest.raises(pickle.UnpicklingError): + SafeUnpickler(io.BytesIO(stream)).load() + + +def test_kv_load_round_trip_dump(tmp_path: Path, kv_memory) -> None: + """A legitimate pickle produced by dump() must load fine.""" + from transformers import DynamicCache + + from memos.memories.activation.item import KVCacheItem + + # No dynamic cache tensors — DynamicCache() empty is picklable + item = KVCacheItem(memory=DynamicCache(), metadata={"note": "hi"}) + kv_memory.add([item]) + kv_memory.dump(str(tmp_path)) + + # Fresh instance + kv_memory.kv_cache_memories = {} + kv_memory.load(str(tmp_path)) + assert item.id in kv_memory.kv_cache_memories + + +# --------------------------------------------------------------------------- +# vllmkv.py — VLLMKVCacheMemory +# --------------------------------------------------------------------------- + + +@pytest.fixture +def vllm_memory(monkeypatch): + from memos.configs.memory import KVCacheMemoryConfig + from memos.llms import factory as llm_factory + from memos.memories.activation.vllmkv import VLLMKVCacheMemory + + monkeypatch.setattr( + llm_factory.LLMFactory, + "from_config", + lambda cfg: MagicMock(build_vllm_kv_cache=lambda x: "prompt"), + ) + config = MagicMock(spec=KVCacheMemoryConfig) + config.extractor_llm = MagicMock() + config.memory_filename = "vllm_kv_cache.pkl" + return VLLMKVCacheMemory(config) + + +def test_vllm_load_refuses_hostile_pickle(tmp_path: Path, vllm_memory) -> None: + canary = tmp_path / "vllm_canary" + _write_hostile_pickle(tmp_path / vllm_memory.config.memory_filename, str(canary)) + + vllm_memory.load(str(tmp_path)) + + assert not canary.exists(), "SafeUnpickler failed to block os.system payload" + assert vllm_memory.kv_cache_memories == {} + + +def test_vllm_load_round_trip_dump(tmp_path: Path, vllm_memory) -> None: + """A legitimate pickle produced by dump() must load fine.""" + from memos.memories.activation.item import VLLMKVCacheItem + + item = VLLMKVCacheItem(memory="a prompt", metadata={"note": "hi"}) + vllm_memory.add([item]) + vllm_memory.dump(str(tmp_path)) + + vllm_memory.kv_cache_memories = {} + vllm_memory.load(str(tmp_path)) + assert item.id in vllm_memory.kv_cache_memories