Fix cache poisoning vulnerability in IntermediateCache

osipy/common/caching.py

@@ -7,8 +7,13 @@
 
 from __future__ import annotations
 
+import contextlib
+import getpass
 import hashlib
 import json
+import os
+import re
+import sys
 import tempfile
 import time
 import warnings
@@ -130,12 +135,59 @@ def __init__(self, config: CacheConfig | None = None) -> None:
 
         # Set up cache directory for persistent storage
         if self.config.cache_dir is None:
-            self._cache_dir = Path(tempfile.gettempdir()) / "osipy_cache"
+            self._cache_dir = self._default_cache_dir()
         else:
             self._cache_dir = Path(self.config.cache_dir)
 
+        self._prepare_cache_dir()
+
+    @staticmethod
+    def _default_cache_dir() -> Path:
+        """Build a cache directory scoped to the current user.
+
+        Using a fixed, shared directory name (e.g. ``<tmp>/osipy_cache``)
+        would let any other local user read or pre-plant files there.
+        Scoping by username keeps the default cache private on
+        multi-user systems (HPC clusters, CI runners, shared containers).
+        """
+        try:
+            user = getpass.getuser()
+        except OSError:
+            user = "default"
+        safe_user = re.sub(r"[^A-Za-z0-9_.-]", "_", user)
+        return Path(tempfile.gettempdir()) / f"osipy_cache_{safe_user}"
+
+    def _prepare_cache_dir(self) -> None:
+        """Create the cache directory and ensure it is private to this user.
+
+        On POSIX systems this refuses to use a cache directory that is a
+        symlink or that another user owns, and restricts permissions to
+        the owner only. This closes the local cache-poisoning path where
+        another user on a shared machine pre-creates or hijacks the cache
+        directory to plant files.
+        """
         self._cache_dir.mkdir(parents=True, exist_ok=True)
 
+        if sys.platform == "win32":
+            return
+
+        if self._cache_dir.is_symlink():
+            raise RuntimeError(
+                f"Cache directory {self._cache_dir} is a symlink; refusing "
+                "to use it. Configure `cache_dir` to a private, non-symlink "
+                "location."
+            )
+
+        owner_uid = self._cache_dir.stat().st_uid
+        if owner_uid != os.getuid():
+            raise RuntimeError(
+                f"Cache directory {self._cache_dir} is owned by another "
+                "user; refusing to use it. Configure `cache_dir` to a "
+                "private location."
+            )
+
+        self._cache_dir.chmod(0o700)
+
     def get_policy(self, result_type: str) -> RetentionPolicy:
         """Get retention policy for a result type.
 
@@ -348,37 +400,45 @@ def _evict_oldest(self) -> None:
         self._total_memory_bytes -= entry.size_bytes
 
     def _put_disk(self, entry: CacheEntry) -> None:
-        """Store entry on disk."""
+        """Store entry on disk.
+
+        Writes to a temporary file in the cache directory first, then
+        atomically renames it into place, so a reader never observes a
+        partially written file and a concurrent writer can't corrupt an
+        in-progress read.
+        """
         cache_file = self._get_cache_path(entry.key)
+        fd, tmp_name = tempfile.mkstemp(
+            dir=self._cache_dir, prefix=cache_file.stem, suffix=".tmp"
+        )
 
         try:
-            if isinstance(entry.data, np.ndarray):
-                if self.config.compression:
-                    np.savez_compressed(
-                        cache_file,
+            with os.fdopen(fd, "wb") as tmp_file:
+                if isinstance(entry.data, np.ndarray):
+                    save = np.savez_compressed if self.config.compression else np.savez
+                    save(
+                        tmp_file,
                         data=entry.data,
                         metadata=json.dumps(entry.metadata),
                         result_type=entry.result_type,
                         created_at=entry.created_at,
                     )
-                else:
-                    np.savez(
-                        cache_file,
-                        data=entry.data,
+                elif isinstance(entry.data, dict):
+                    np.savez_compressed(
+                        tmp_file,
+                        **{f"data_{k}": v for k, v in entry.data.items()},
                         metadata=json.dumps(entry.metadata),
                         result_type=entry.result_type,
                         created_at=entry.created_at,
                     )
-            elif isinstance(entry.data, dict):
-                np.savez_compressed(
-                    cache_file,
-                    **{f"data_{k}": v for k, v in entry.data.items()},
-                    metadata=json.dumps(entry.metadata),
-                    result_type=entry.result_type,
-                    created_at=entry.created_at,
-                )
+                else:
+                    return
+            Path(tmp_name).replace(cache_file)
         except Exception as e:
             warnings.warn(f"Failed to save cache to disk: {e}", stacklevel=2)
+        finally:
+            with contextlib.suppress(FileNotFoundError):
+                Path(tmp_name).unlink()
 
     def _get_disk(self, key: str) -> NDArray[Any] | dict[str, Any] | None:
         """Retrieve entry from disk."""
@@ -388,7 +448,11 @@ def _get_disk(self, key: str) -> NDArray[Any] | dict[str, Any] | None:
             return None
 
         try:
-            loaded = np.load(cache_file, allow_pickle=True)
+            # allow_pickle=False is intentional: cache entries only ever
+            # contain numeric arrays and JSON-encoded strings, so pickle
+            # support is never needed and would let a file planted by
+            # another user execute arbitrary code on load.
+            loaded = np.load(cache_file, allow_pickle=False)
 
             # Check age
             created_at = float(loaded.get("created_at", 0))

tests/unit/common/test_caching.py

@@ -5,7 +5,11 @@
 
 from __future__ import annotations
 
+import hashlib
+import os
+import sys
 import tempfile
+import time
 from pathlib import Path
 
 import numpy as np
@@ -19,6 +23,12 @@
     get_cache,
 )
 
+# Cache-directory hardening only applies on POSIX; on Windows there is no
+# os.getuid()/chmod() and the check is skipped in IntermediateCache itself.
+posix_only = pytest.mark.skipif(
+    sys.platform == "win32", reason="POSIX-only cache directory hardening"
+)
+
 
 class TestRetentionPolicy:
     """Tests for RetentionPolicy enum."""
@@ -218,6 +228,113 @@ def test_configure_cache(self) -> None:
         assert stats["max_memory_mb"] == 512
 
 
+_poison_executed = False
+
+
+def _poison_payload() -> None:
+    """Stand-in for arbitrary code a malicious cache file could run."""
+    global _poison_executed
+    _poison_executed = True
+
+
+class _EvilReduce:
+    """Object whose pickle reconstruction runs `_poison_payload`."""
+
+    def __reduce__(self) -> tuple:
+        return (_poison_payload, ())
+
+
+class TestCachePoisoningRegression:
+    """Regression tests for GH-171: cache poisoning via pickle deserialization.
+
+    A malicious ``.npz`` file placed at a cache path (e.g. by another user
+    on a shared machine, since the filename is a predictable hash of the
+    cache key) must not be able to execute code when loaded, and the
+    default cache directory must not be a single shared location that
+    every local user can write into.
+    """
+
+    def test_planted_pickle_payload_is_not_executed(self) -> None:
+        """A pre-planted malicious npz at the predictable cache path must
+        not execute code when read back via get()."""
+        global _poison_executed
+        _poison_executed = False
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config = CacheConfig(
+                cache_dir=Path(tmpdir),
+                policies={"t1_map": RetentionPolicy.PERSISTENT},
+            )
+            cache = IntermediateCache(config)
+
+            # Plant a file at the exact path osipy would use for this key,
+            # containing an object whose __reduce__ runs arbitrary code.
+            full_key = "t1_map:subject01"
+            key_hash = hashlib.md5(full_key.encode()).hexdigest()
+            poisoned_path = Path(tmpdir) / f"{key_hash}.npz"
+
+            np.savez(
+                poisoned_path,
+                data=np.array(_EvilReduce(), dtype=object),
+                created_at=np.array(time.time()),
+            )
+
+            result = cache.get("t1_map", "subject01")
+
+            assert _poison_executed is False
+            assert result is None
+
+    def test_default_cache_dir_is_scoped_per_user(self) -> None:
+        """The default cache directory must not be a single name shared by
+        every local user (the predictable, world-writable path in GH-171)."""
+        cache_dir = IntermediateCache._default_cache_dir()
+
+        assert cache_dir.name != "osipy_cache"
+        assert cache_dir.name.startswith("osipy_cache_")
+
+    @posix_only
+    def test_default_cache_dir_is_restricted_to_owner(self) -> None:
+        """The cache directory should not be readable/writable by other
+        local users."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config = CacheConfig(cache_dir=Path(tmpdir) / "osipy_cache")
+            cache = IntermediateCache(config)
+
+            mode = cache._cache_dir.stat().st_mode & 0o777
+            assert mode == 0o700
+
+    @posix_only
+    def test_symlinked_cache_dir_is_rejected(self) -> None:
+        """A cache_dir that is a symlink (e.g. planted by another user to
+        redirect osipy's writes elsewhere) must be refused."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            real_target = Path(tmpdir) / "real_target"
+            real_target.mkdir()
+            symlink_path = Path(tmpdir) / "osipy_cache_link"
+            symlink_path.symlink_to(real_target, target_is_directory=True)
+
+            config = CacheConfig(cache_dir=symlink_path)
+            with pytest.raises(RuntimeError, match="symlink"):
+                IntermediateCache(config)
+
+    @posix_only
+    def test_cache_dir_owned_by_another_user_is_rejected(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A pre-existing cache directory owned by a different user must be
+        refused, even if permissions happen to look fine."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_dir = Path(tmpdir) / "osipy_cache"
+            cache_dir.mkdir()
+            real_uid = cache_dir.stat().st_uid
+
+            monkeypatch.setattr(os, "getuid", lambda: real_uid + 1)
+
+            config = CacheConfig(cache_dir=cache_dir)
+            with pytest.raises(RuntimeError, match="owned by another user"):
+                IntermediateCache(config)
+
+
 class TestCacheWithPolicies:
     """Tests for cache with different policies per type."""