diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 508c221..067d20c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,7 +64,7 @@ jobs: with: python-version: ${{ matrix.python }} - name: Install diskcache and test dependencies - run: pip install -e . "django==5.2.*" pytest pytest-cov pytest-django pytest-xdist + run: pip install -e . "django==5.2.*" numpy pydantic pytest pytest-cov pytest-django pytest-xdist - name: Run tests run: pytest diff --git a/CHANGELOG.md b/CHANGELOG.md index ad7183e..386c11f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ diskcache-community forked from DiskCache at upstream commit ebfa37c (3 March 20 ## Unreleased +### Security + +- `SafeDisk.with_types(*types, modules=())` returns a `SafeDisk` subclass whose allowlist adds the given classes, matched by identity, and every class defined in the named modules or below them. A module name may only reach a class defined in that module: a dotted name such as `os.system`, reached through a module that imports `os`, a class the module merely imports, a module and a dunder attribute are refused, and reading never imports a module, so a pickle cannot run a module's top-level code. numpy classes are refused (`UNSAFE_MODULES`), because numpy restores a dtype's flags from the pickle without checking them and a crafted array pickle can corrupt memory. Writes stay exact: a value outside the standard-library types is pickled and the names its pickle looks up are checked against the same allowlist, so a value the reader would refuse raises `TypeError` naming the class before anything is written. Calling `with_types` again adds to the allowlist, and a `Cache`, `FanoutCache`, `Deque`, `Index` or `DjangoCache` using the returned class pickles as before. The plain `SafeDisk` is unchanged (PR #23). + ### Docs - Shared links show a project card: a social preview image in `.github/`, and Open Graph tags on the documentation through `sphinxext-opengraph` (PR #21). diff --git a/diskcache_community/core.py b/diskcache_community/core.py index 47d54fa..d1e73d1 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -4,6 +4,7 @@ import codecs import collections as co import contextlib as cl +import copyreg import datetime as dt import decimal import errno @@ -19,6 +20,7 @@ import pickletools import sqlite3 import struct +import sys import tempfile import threading import time @@ -553,24 +555,180 @@ def fetch(self, mode, filename, value, read): _SAFE_SEQUENCES = (list, tuple, set, frozenset, co.deque) _SAFE_TZ_TYPES = (type(None), dt.timezone) +# Packages whose classes SafeDisk.with_types refuses to allow. numpy restores +# the internal flags of a dtype from its pickle without checking them, so a +# crafted pickle can make an array read raw bytes as object pointers. + +UNSAFE_MODULES = ('numpy',) + +_PICKLE_STRING_OPS = frozenset( + {'BINUNICODE', 'BINUNICODE8', 'SHORT_BINUNICODE', 'UNICODE'} +) +_PICKLE_GET_OPS = frozenset({'BINGET', 'GET', 'LONG_BINGET'}) +_PICKLE_PUT_OPS = frozenset({'BINPUT', 'LONG_BINPUT', 'PUT'}) +_PICKLE_EXT_OPS = frozenset({'EXT1', 'EXT2', 'EXT4'}) + + +def _in_modules(module, modules): + """Return True when `module` is one of `modules` or inside one.""" + return any( + module == prefix or module.startswith(prefix + '.') + for prefix in modules + ) + + +def _pickle_globals(data): + """Return the (module, name) pairs that the pickle `data` looks up. + + Return None when a STACK_GLOBAL takes its operands from something other + than the strings pushed just before it, which the pickler never writes, + or when the pickle uses a copyreg extension code instead of a name. + + """ + found = set() + pushed = [] + memo = {} + + for opcode, arg, _ in pickletools.genops(data): + name = opcode.name + + if name in _PICKLE_EXT_OPS: + return None + + if name in ('GLOBAL', 'INST'): + module, _, attr = arg.partition(' ') + found.add((module, attr)) + elif name == 'STACK_GLOBAL': + if len(pushed) < 2 or None in pushed[-2:]: + return None + found.add((pushed[-2], pushed[-1])) + + if name in _PICKLE_STRING_OPS: + pushed.append(arg) + elif name in _PICKLE_GET_OPS: + pushed.append(memo.get(arg)) + elif name == 'MEMOIZE': + memo[len(memo)] = pushed[-1] if pushed else None + elif name in _PICKLE_PUT_OPS: + memo[arg] = pushed[-1] if pushed else None + elif name not in ('FRAME', 'PROTO'): + pushed.append(None) + + return found + class SafeUnpickler(pickle.Unpickler): - """Unpickler that resolves only the names in SAFE_PICKLE_NAMES. + """Unpickler that resolves only an allowlist of names. Every pickle opcode that can call code (GLOBAL, STACK_GLOBAL, INST, OBJ, NEWOBJ, NEWOBJ_EX and REDUCE) obtains its callable through `find_class`, so restricting `find_class` to a fixed set of constructors is enough to keep a crafted pickle from running code. + By default the allowlist is SAFE_PICKLE_NAMES. `names` replaces that + mapping of module to names, `classes` adds classes by identity, and + `modules` adds every class defined in the given modules or below them. + A class is looked up attribute by attribute and every step must be a + class, so a name such as ``os.system`` that reaches a function through a + module imported by a trusted module is refused. Classes are looked up + only in modules that are already imported; reading a pickle never + imports a module, so it cannot run a module's top-level code. + """ + def __init__(self, file, names=None, classes=(), modules=()): + super().__init__(file) + self._names = SAFE_PICKLE_NAMES if names is None else names + self._classes = frozenset(classes) + self._class_modules = frozenset(cls.__module__ for cls in classes) + self._modules = tuple(modules) + def find_class(self, module, name): - if name in SAFE_PICKLE_NAMES.get(module, ()): + if name in self._names.get(module, ()): return super().find_class(module, name) + + if module in self._class_modules or _in_modules(module, self._modules): + obj = self._resolve(module, name) + if obj is not None and ( + obj in self._classes + or _in_modules(obj.__module__, self._modules) + ): + return obj + raise pickle.UnpicklingError( 'SafeDisk refuses to unpickle %s.%s' % (module, name) ) + @staticmethod + def _resolve(module, name): + """Return the class `name` in `module`, or None if it is not one.""" + parts = name.split('.') + + if any(not part or part.startswith('__') for part in parts): + return None + + obj = sys.modules.get(module) + + for part in parts: + obj = getattr(obj, part, None) + if not isinstance(obj, type): + return None + + return obj + + +class _DerivedDiskType(type): + """Metaclass of the disk classes that SafeDisk.with_types derives. + + A derived class is a subclass of a disk class defined in a module, its + root, with a few class attributes set. It pickles as its root and those + attributes, so a Cache using it can be pickled like any other. + + """ + + def __repr__(cls): + if '_disk_root' not in cls.__dict__: + return super().__repr__() + options = ', '.join( + '%s=%r' % item for item in sorted(cls._disk_options.items()) + ) + return '' % ( + '%s.%s' % (cls.__module__, cls.__qualname__), + options, + ) + + +def _derive_disk(base, options): + """Return a subclass of `base` whose class attributes include `options`. + + Deriving from a derived class starts from its root and its options, so + derived classes never nest. A class defined in a module that subclasses + a derived class is a root of its own. + + """ + # pylint: disable=protected-access + root = base._disk_root if '_disk_root' in base.__dict__ else base + merged = dict(getattr(base, '_disk_options', {})) + merged.update(options) + attrs = dict(merged) + attrs.update( + __module__=root.__module__, + __qualname__=root.__qualname__, + _disk_root=root, + _disk_options=merged, + ) + return _DerivedDiskType(root.__name__, (root,), attrs) + + +def _reduce_derived_disk(cls): + # pylint: disable=protected-access + if '_disk_root' not in cls.__dict__: + return cls.__qualname__ + return _derive_disk, (cls._disk_root, cls._disk_options) + + +copyreg.pickle(_DerivedDiskType, _reduce_derived_disk) + class SafeDisk(Disk): """Disk that reads back only an allowlist of standard-library types. @@ -598,14 +756,134 @@ class SafeDisk(Disk): Objects outside the allowlist, such as instances of your own classes, Django model instances and responses, and the ``type`` objects that - ``memoize(typed=True)`` puts into keys, need the default :class:`Disk`. + ``memoize(typed=True)`` puts into keys, need either the default + :class:`Disk` or a disk from :meth:`with_types` that allows them. """ + safe_types: frozenset = frozenset() + safe_modules: tuple = () + def __init__(self, directory, min_file_size=0, pickle_protocol=0): super().__init__(directory, min_file_size, pickle_protocol) self._real_directory = op.realpath(directory) + @classmethod + def with_types(cls, *types, modules=()): + """Return a SafeDisk subclass that also allows the given types. + + The allowlist of the returned class is the standard-library one plus + `types`, matched by identity, and every class defined in `modules` or + in a module below them. Calling it on a class that it returned adds + to that class's allowlist. Writes are checked against the same + allowlist, so a value that the reader would refuse raises + :exc:`TypeError` before it is written. + + Allow only plain data classes whose construction and ``__setstate__`` + do nothing beyond storing their fields, such as dataclasses, named + tuples, enums and pydantic models. Reading looks classes up only in + modules that are already imported, so import them before reading. + numpy classes are refused, see ``UNSAFE_MODULES``; store arrays as + bytes, for example with ``numpy.save(..., allow_pickle=False)``. + + >>> import collections + >>> Point = collections.namedtuple('Point', 'x y') + >>> PointDisk = SafeDisk.with_types(Point) + >>> issubclass(PointDisk, SafeDisk) + True + + :param types: classes to allow, by identity + :param modules: module names whose classes to allow, such as + ``('litellm.types',)``; a single string is one module + :return: SafeDisk subclass to pass as ``disk=`` + :raises TypeError: if an entry of `types` is not a class + :raises ValueError: if an entry of `modules` is not a module name, or + a type or module is in ``UNSAFE_MODULES`` + + """ + for entry in types: + if not isinstance(entry, type): + raise TypeError('with_types takes classes, not %r' % (entry,)) + if _in_modules(entry.__module__, UNSAFE_MODULES): + raise ValueError('with_types refuses the class %r' % (entry,)) + + if isinstance(modules, str): + modules = (modules,) + prefixes = [] + for entry in modules: + prefix = entry.rstrip('.') if isinstance(entry, str) else '' + if not prefix or not all( + part.isidentifier() for part in prefix.split('.') + ): + raise ValueError('%r is not a module name' % (entry,)) + if any( + _in_modules(prefix, (unsafe,)) + or _in_modules(unsafe, (prefix,)) + for unsafe in UNSAFE_MODULES + ): + raise ValueError('with_types refuses the module %r' % (entry,)) + prefixes.append(prefix) + + return _derive_disk( + cls, + { + 'safe_types': cls.safe_types | frozenset(types), + 'safe_modules': tuple( + dict.fromkeys(cls.safe_modules + tuple(prefixes)) + ), + }, + ) + + def _unpickler(self, data): + return SafeUnpickler( + io.BytesIO(data), + SAFE_PICKLE_NAMES, + self.safe_types, + self.safe_modules, + ) + + def _extended(self): + return bool(self.safe_types or self.safe_modules) + + def _check_object(self, obj, what): + """Raise TypeError unless the reader would accept `obj` back.""" + try: + self._check(obj, what) + except TypeError: + if not self._extended(): + raise + else: + return + + data = pickle.dumps(obj, protocol=self.pickle_protocol) + names = _pickle_globals(data) + + if names is None: + try: + self._unpickler(data).load() + except pickle.UnpicklingError as error: + raise TypeError( + 'SafeDisk cannot store %s: %s' % (what, error) + ) from None + return + + for module, name in sorted(names): + try: + self._unpickler(b'').find_class(module, name) + except pickle.UnpicklingError: + raise TypeError( + 'SafeDisk cannot store %s of type %s.%s: it pickles' + ' %s.%s, which the allowlist refuses; allow it with' + ' SafeDisk.with_types()' + % ( + what, + type(obj).__module__, + type(obj).__qualname__, + module, + name, + ) + ) from None + def _check(self, obj, what): """Raise TypeError unless every object in `obj` is in SAFE_TYPES.""" seen = set() @@ -682,18 +960,18 @@ def _confine(self, filename): def put(self, key): # pylint: disable=unidiomatic-typecheck if type(key) not in (bytes, str, int, float): - self._check(key, 'key') + self._check_object(key, 'key') return super().put(key) def get(self, key, raw): if raw: return super().get(key, raw) - return SafeUnpickler(io.BytesIO(key)).load() + return self._unpickler(key).load() def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck if not read and type(value) not in (bytes, str, int, float): - self._check(value, 'value') + self._check_object(value, 'value') return super().store(value, read, key=key) def fetch(self, mode, filename, value, read): @@ -703,7 +981,7 @@ def fetch(self, mode, filename, value, read): if mode == MODE_PICKLE: if value is None: value = super().fetch(MODE_BINARY, filename, None, False) - return SafeUnpickler(io.BytesIO(value)).load() + return self._unpickler(value).load() return super().fetch(mode, filename, value, read) diff --git a/docs/api.rst b/docs/api.rst index eaf70c7..91e6add 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -151,6 +151,12 @@ Read the :ref:`SafeDisk tutorial ` for details. The Python types that :class:`SafeDisk ` accepts for writing; exactly the types the names above read back. +.. data:: diskcache_community.core.UNSAFE_MODULES + + Packages whose classes :meth:`SafeDisk.with_types + ` refuses to allow, because + unpickling them is not memory-safe. + Timeout ------- diff --git a/docs/tutorial.rst b/docs/tutorial.rst index fbf3262..35b64f1 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -979,6 +979,43 @@ process that writes to a shared directory with the default :class:`Disk :class:`SafeDisk ` raises :exc:`pickle.UnpicklingError` on those rows. +Instead of falling back to the default :class:`Disk +`, :meth:`SafeDisk.with_types +` returns a :class:`SafeDisk +` subclass that also allows your own classes +and every class defined in the modules you name. Everything else stays +refused, both when reading and when writing: + +.. code-block:: python + + from diskcache_community import Cache, SafeDisk + + from myapp.models import Message, Point + + AppDisk = SafeDisk.with_types( + Message, + Point, + modules=('litellm.types', 'openai.types'), + ) + cache = Cache('/tmp/mycache', disk=AppDisk) + +Classes given by name are matched by identity, so a subclass is refused +unless it is listed too. A module prefix allows only classes defined in that +module or below it: a pickle cannot reach a function, a module, or a class +that the trusted module merely imports, such as ``os.system`` through a module +that imports ``os``. Reading looks classes up only in modules that are already +imported and never imports one, so import the modules before reading. List +only plain data classes whose construction and ``__setstate__`` do nothing +more than store their fields, such as dataclasses, named tuples, enums and +pydantic models, because allowing a class lets a crafted pickle create +instances of it with any field values. numpy classes are refused: numpy +restores a dtype's internal flags from the pickle without checking them, so a +crafted array pickle can corrupt memory. Store arrays as bytes instead, for +example with ``numpy.save(buffer, array, allow_pickle=False)``. A write whose +value pickles a class outside the allowlist raises :exc:`TypeError` that names +the class. Calling ``with_types`` on the returned class adds to its allowlist, +and a cache that uses it can be pickled as usual. + .. _tutorial-read-only: Read-Only Caches diff --git a/tests/safedisk_models/__init__.py b/tests/safedisk_models/__init__.py new file mode 100644 index 0000000..c8dd4be --- /dev/null +++ b/tests/safedisk_models/__init__.py @@ -0,0 +1 @@ +"""Classes for the SafeDisk.with_types tests, in a package of their own.""" diff --git a/tests/safedisk_models/models.py b/tests/safedisk_models/models.py new file mode 100644 index 0000000..d34dbff --- /dev/null +++ b/tests/safedisk_models/models.py @@ -0,0 +1,45 @@ +"""Classes that the tests allow through a module prefix. + +The module imports ``os`` and the class Foreign from a sibling module, so +the tests can check that a pickle cannot reach either through this module. +""" + +import collections +import dataclasses +import enum +import os # noqa: F401 + +from tests.safedisk_models.other import Foreign # noqa: F401 + + +@dataclasses.dataclass +class Point: + x: int + y: int + + +@dataclasses.dataclass +class Shape: + name: str + points: list + + +class Color(enum.Enum): + RED = 1 + GREEN = 2 + + +class Outer: + class Inner: + def __init__(self, value): + self.value = value + + def __eq__(self, other): + return type(other) is type(self) and other.value == self.value + + +Pair = collections.namedtuple('Pair', 'left right') + + +class PointSubclass(Point): + pass diff --git a/tests/safedisk_models/other.py b/tests/safedisk_models/other.py new file mode 100644 index 0000000..9e34cdd --- /dev/null +++ b/tests/safedisk_models/other.py @@ -0,0 +1,8 @@ +"""A class defined outside the trusted models module.""" + +import dataclasses + + +@dataclasses.dataclass +class Foreign: + name: str diff --git a/tests/safedisk_models/side_effect.py b/tests/safedisk_models/side_effect.py new file mode 100644 index 0000000..4774a97 --- /dev/null +++ b/tests/safedisk_models/side_effect.py @@ -0,0 +1,3 @@ +"""A module that the tests check a read never imports.""" + +raise RuntimeError('tests.safedisk_models.side_effect was imported') diff --git a/tests/test_safedisk_types.py b/tests/test_safedisk_types.py new file mode 100644 index 0000000..6f51b78 --- /dev/null +++ b/tests/test_safedisk_types.py @@ -0,0 +1,513 @@ +"""Test diskcache_community.SafeDisk.with_types.""" + +import copyreg +import io +import os.path as op +import pickle +import pickletools +import shutil +import sqlite3 +import sys +import tempfile + +import pytest + +import diskcache_community as dc +from diskcache_community.core import ( + MODE_PICKLE, + SafeUnpickler, + _pickle_globals, +) +from tests.safedisk_models import models +from tests.safedisk_models.other import Foreign + +pytestmark = pytest.mark.filterwarnings('ignore', category=dc.EmptyDirWarning) + +MODULE = 'tests.safedisk_models' + + +class AppDisk(dc.SafeDisk.with_types(models.Point)): + """A disk class defined in a module on top of a derived class.""" + + marker = 'app' + + +def _execute(path): + """Marker that a crafted pickle ran code: it creates `path`.""" + with open(path, 'w', encoding='utf-8') as writer: + writer.write('executed') + return 'executed' + + +class Poison: + """Object whose pickle calls `_execute` when loaded.""" + + def __init__(self, path): + self.path = path + + def __reduce__(self): + return (_execute, (self.path,)) + + +def stack_global(module, name, *args): + """Return a protocol 4 pickle that looks up module.name and calls it.""" + + def string(text): + data = text.encode('utf-8') + return b'\x8c' + bytes([len(data)]) + data + + if not args: + return b'\x80\x04' + string(module) + string(name) + b'\x93.' + call = string(module) + string(name) + b'\x93' + string(args[0]) + return b'\x80\x04' + call + b'\x85R.' + + +@pytest.fixture +def directory(): + directory = tempfile.mkdtemp(prefix='diskcache-') + yield directory + shutil.rmtree(directory, ignore_errors=True) + + +def test_with_types_derives_a_safedisk(): + disk = dc.SafeDisk.with_types(models.Point) + + assert issubclass(disk, dc.SafeDisk) + assert disk is not dc.SafeDisk + assert disk.__name__ == 'SafeDisk' + assert disk.safe_types == frozenset({models.Point}) + assert disk.safe_modules == () + assert 'safe_types' in repr(disk) + assert dc.SafeDisk.safe_types == frozenset() + + +def test_with_types_adds_to_a_derived_class(): + first = dc.SafeDisk.with_types(models.Point, modules='a.b') + second = first.with_types(models.Pair, modules=['a.b', 'c']) + + assert second.safe_types == frozenset({models.Point, models.Pair}) + assert second.safe_modules == ('a.b', 'c') + assert second.__bases__ == (dc.SafeDisk,) + assert first.safe_types == frozenset({models.Point}) + + +def test_with_types_on_a_subclass_keeps_the_subclass(): + class LoggingDisk(dc.SafeDisk): + pass + + disk = LoggingDisk.with_types(models.Point) + + assert issubclass(disk, LoggingDisk) + assert disk.__bases__ == (LoggingDisk,) + + +@pytest.mark.parametrize('entry', [1, 'models.Point', models.Point(1, 2)]) +def test_with_types_rejects_non_classes(entry): + with pytest.raises(TypeError, match='with_types takes classes'): + dc.SafeDisk.with_types(entry) + + +@pytest.mark.parametrize('entry', ['', '.', 'a..b', '1abc', 'a.b-c', None]) +def test_with_types_rejects_bad_module_names(entry): + with pytest.raises(ValueError, match='not a module name'): + dc.SafeDisk.with_types(modules=[entry]) + + +def test_module_prefix_drops_a_trailing_dot(): + disk = dc.SafeDisk.with_types(modules='litellm.types.') + assert disk.safe_modules == ('litellm.types',) + + +@pytest.mark.parametrize('protocol', [0, 1, 2, 3, 4, 5]) +@pytest.mark.parametrize('min_file_size', [0, 2**15]) +def test_exact_types_round_trip(directory, protocol, min_file_size): + disk = dc.SafeDisk.with_types( + models.Point, models.Shape, models.Pair, models.Color + ) + values = [ + models.Point(1, 2), + models.Shape('square', [models.Point(0, 0), models.Point(1, 1)]), + models.Pair('left', [1, 2]), + models.Color.GREEN, + {'points': [models.Point(3, 4)], 'color': models.Color.RED}, + ] + + with dc.Cache( + directory, + disk=disk, + disk_pickle_protocol=protocol, + disk_min_file_size=min_file_size, + ) as cache: + for number, value in enumerate(values): + cache[number] = value + assert cache[number] == value + cache[models.Pair(1, 2)] = 'pair key' + assert cache[models.Pair(1, 2)] == 'pair key' + + with dc.Cache(directory) as cache: + assert cache[0] == values[0] + + +def test_nested_class(directory): + disk = dc.SafeDisk.with_types(models.Outer.Inner) + + with dc.Cache(directory, disk=disk) as cache: + cache['inner'] = models.Outer.Inner(5) + assert cache['inner'] == models.Outer.Inner(5) + + +def test_exact_types_refuse_others_at_write(directory): + disk = dc.SafeDisk.with_types(models.Point) + + with dc.Cache(directory, disk=disk) as cache: + with pytest.raises(TypeError, match='with_types'): + cache['shape'] = models.Shape('line', []) + with pytest.raises(TypeError, match='PointSubclass'): + cache['subclass'] = models.PointSubclass(1, 2) + with pytest.raises(TypeError, match='Foreign'): + cache['nested'] = [models.Point(1, 2), Foreign('x')] + with pytest.raises(TypeError, match='key'): + cache[models.Shape('key', [])] = 1 + assert len(cache) == 0 + + +def test_exact_types_refuse_others_at_read(directory): + with dc.Cache(directory) as cache: + cache['shape'] = models.Shape('line', []) + cache['point'] = models.Point(1, 2) + + disk = dc.SafeDisk.with_types(models.Point) + + with dc.Cache(directory, disk=disk) as cache: + assert cache['point'] == models.Point(1, 2) + with pytest.raises(pickle.UnpicklingError, match='Shape'): + cache['shape'] # pylint: disable=pointless-statement + + +def test_poisoned_value_still_refused(directory): + path = op.join(directory, 'marker') + disk = dc.SafeDisk.with_types(models.Point, modules=MODULE) + + with dc.Cache(directory, disk=disk) as cache: + cache['key'] = models.Point(1, 2) + cache._sql( + 'UPDATE Cache SET mode = ?, value = ? WHERE key = ?', + (MODE_PICKLE, sqlite3.Binary(pickle.dumps(Poison(path))), 'key'), + ) + with pytest.raises(pickle.UnpicklingError, match='refuses'): + cache['key'] # pylint: disable=pointless-statement + + assert not op.exists(path) + + +def test_module_prefix_round_trip(directory): + disk = dc.SafeDisk.with_types(modules=[MODULE]) + values = [ + models.Point(1, 2), + models.PointSubclass(3, 4), + models.Outer.Inner('x'), + models.Color.RED, + Foreign('defined in a sibling module'), + ] + + with dc.Cache(directory, disk=disk) as cache: + for number, value in enumerate(values): + cache[number] = value + assert cache[number] == value + + +def test_module_prefix_matches_whole_names(): + unpickler = SafeUnpickler(io.BytesIO(b''), modules=['tests.safedisk']) + + with pytest.raises(pickle.UnpicklingError, match='refuses'): + unpickler.find_class('tests.safedisk_models.models', 'Point') + + +def test_module_prefix_does_not_reach_imported_names(): + disk = dc.SafeDisk.with_types(modules=['tests.safedisk_models.models']) + module = 'tests.safedisk_models.models' + + with pytest.raises(TypeError, match='Foreign'): + disk(tempfile.gettempdir()).store(Foreign('x'), False) + + for name in ( + 'os.system', + 'os', + 'Foreign', + 'collections.OrderedDict', + 'Point.__subclasses__', + 'Point.__init__', + 'Point.', + 'dataclasses.dataclass', + ): + with pytest.raises(pickle.UnpicklingError, match='refuses'): + SafeUnpickler( + io.BytesIO(stack_global(module, name)), modules=[module] + ).load() + + +def test_module_prefix_blocks_dotted_name_to_os_system(directory): + path = op.join(directory, 'marker') + module = 'tests.safedisk_models.models' + data = stack_global(module, 'os.system', 'touch %s' % path) + + with pytest.raises(pickle.UnpicklingError, match='os.system'): + SafeUnpickler(io.BytesIO(data), modules=[MODULE]).load() + + assert not op.exists(path) + + +def test_refused_modules_are_not_imported(): + name = 'tests.safedisk_models.never_imported' + unpickler = SafeUnpickler(io.BytesIO(b''), classes=[models.Point]) + + with pytest.raises(pickle.UnpicklingError, match='refuses'): + unpickler.find_class(name, 'Anything') + + assert name not in sys.modules + + +def test_exact_class_must_be_the_same_object(): + unpickler = SafeUnpickler(io.BytesIO(b''), classes=[models.Point]) + + assert unpickler.find_class(models.__name__, 'Point') is models.Point + with pytest.raises(pickle.UnpicklingError, match='refuses'): + unpickler.find_class(models.__name__, 'PointSubclass') + with pytest.raises(pickle.UnpicklingError, match='refuses'): + unpickler.find_class(models.__name__, 'os') + + +def test_default_allowlist_unchanged(directory): + with dc.Cache(directory, disk=dc.SafeDisk) as cache: + with pytest.raises(TypeError, match='SafeDisk cannot store value'): + cache['point'] = models.Point(1, 2) + + +def test_typed_memoize_with_types(directory): + disk = dc.SafeDisk.with_types(models.Point) + + with dc.Cache(directory, disk=disk) as cache: + + @cache.memoize(typed=True) + def double(value): + return value * 2 + + assert double(2) == 4 + assert double(2) == 4 + assert double(2.0) == 4.0 + assert len(cache) == 2 + + +@pytest.mark.parametrize('protocol', [0, 1, 2, 3, 4, 5]) +def test_pickle_globals_matches_pickletools(protocol): + values = [ + models.Point(1, 2), + [models.Point(1, 2), models.Point(3, 4), models.Pair(1, 2)], + {models.Color.RED: models.Outer.Inner(1)}, + int, + sqlite3.Row, + ] + + for value in values: + data = pickle.dumps(value, protocol=protocol) + expected = set() + for opcode, arg, _ in pickletools.genops(data): + if opcode.name in ('GLOBAL', 'INST'): + expected.add(tuple(arg.split(' ', 1))) + found = _pickle_globals(data) + assert found is not None + assert expected <= found + + +def test_pickle_globals_gives_up_on_unknown_operands(): + data = b'\x80\x04\x8c\x04name\x8c\x04attr\x86\x94\x93.' + assert _pickle_globals(data) is None + + +def test_write_check_falls_back_to_loading(directory, monkeypatch): + disk = dc.SafeDisk.with_types(models.Point) + monkeypatch.setattr( + 'diskcache_community.core._pickle_globals', lambda data: None + ) + + with dc.Cache(directory, disk=disk) as cache: + cache['point'] = models.Point(1, 2) + assert cache['point'] == models.Point(1, 2) + with pytest.raises(TypeError, match='SafeDisk cannot store value'): + cache['shape'] = models.Shape('line', []) + + +def test_cache_with_derived_disk_pickles(directory): + disk = dc.SafeDisk.with_types(models.Point, modules=MODULE) + + with dc.Cache(directory, disk=disk) as cache: + cache['point'] = models.Point(1, 2) + copy = pickle.loads(pickle.dumps(cache)) + try: + assert type(copy.disk).safe_types == disk.safe_types + assert type(copy.disk).safe_modules == disk.safe_modules + assert copy['point'] == models.Point(1, 2) + finally: + copy.close() + + with dc.FanoutCache(directory + '-fanout', disk=disk) as fanout: + copy = pickle.loads(pickle.dumps(fanout)) + try: + assert type(copy._shards[0].disk).safe_types == disk.safe_types + finally: + copy.close() + + +def test_containers_take_derived_disks(directory): + disk = dc.SafeDisk.with_types(models.Point) + + deque = dc.Deque([models.Point(1, 2)], directory=directory, disk=disk) + assert deque[0] == models.Point(1, 2) + deque.cache.close() + + index = dc.Index(op.join(directory, 'index'), disk=disk) + index['point'] = models.Point(3, 4) + assert index['point'] == models.Point(3, 4) + index.cache.close() + + with dc.FanoutCache(op.join(directory, 'fanout'), disk=disk) as fanout: + fanout['point'] = models.Point(5, 6) + assert fanout['point'] == models.Point(5, 6) + assert type(fanout.cache('sub').disk) is disk + + +def test_django_disk_option_with_derived_disk(directory): + pytest.importorskip('django') + disk = dc.SafeDisk.with_types(models.Point) + + django_cache = dc.DjangoCache(directory, {'DISK': disk}) + try: + assert type(django_cache._cache.disk) is disk + assert django_cache.set('point', models.Point(1, 2)) + assert django_cache.get('point') == models.Point(1, 2) + finally: + django_cache.close() + + +def test_numpy_is_refused(directory): + np = pytest.importorskip('numpy') + + for entry in (np.ndarray, np.dtype, np.float64): + with pytest.raises(ValueError, match='refuses the class'): + dc.SafeDisk.with_types(entry) + + for prefix in ('numpy', 'numpy.', 'numpy._core', 'numpy.core.multiarray'): + with pytest.raises(ValueError, match='refuses the module'): + dc.SafeDisk.with_types(modules=prefix) + + disk = dc.SafeDisk.with_types(models.Point, modules=MODULE) + + with dc.Cache(directory, disk=disk) as cache: + for value in (np.arange(3), np.float64(1.5), np.dtype('f8')): + with pytest.raises(TypeError, match='numpy'): + cache['value'] = value + + +def test_modules_are_never_imported_by_a_read(): + name = 'tests.safedisk_models.side_effect' + data = stack_global(name, 'Thing') + + with pytest.raises(pickle.UnpicklingError, match='refuses'): + SafeUnpickler(io.BytesIO(data), modules=[MODULE]).load() + + assert name not in sys.modules + + +def test_trailing_dot_module_is_refused(): + data = stack_global('tests.safedisk_models.models.', 'Point') + + with pytest.raises(pickle.UnpicklingError, match='refuses'): + SafeUnpickler(io.BytesIO(data), modules=[MODULE]).load() + + +def test_subclass_of_a_derived_class_is_its_own_root(directory): + assert repr(AppDisk) == "" + assert AppDisk.safe_types == frozenset({models.Point}) + + more = AppDisk.with_types(models.Pair) + assert issubclass(more, AppDisk) + assert more.marker == 'app' + assert more.safe_types == frozenset({models.Point, models.Pair}) + + with dc.Cache(directory, disk=AppDisk) as cache: + cache['point'] = models.Point(1, 2) + copy = pickle.loads(pickle.dumps(cache)) + try: + assert type(copy.disk) is AppDisk + finally: + copy.close() + + assert pickle.loads(pickle.dumps(more)).safe_types == more.safe_types + + +def test_write_check_handles_extension_codes(directory): + disk = dc.SafeDisk.with_types(models.Point) + code = 0x7FFF0001 + copyreg.add_extension(models.__name__, 'Shape', code) + + try: + data = pickle.dumps(models.Shape('line', []), protocol=2) + assert _pickle_globals(data) is None + with dc.Cache(directory, disk=disk, disk_pickle_protocol=2) as cache: + with pytest.raises(TypeError, match='SafeDisk cannot store value'): + cache['shape'] = models.Shape('line', []) + finally: + copyreg.remove_extension(models.__name__, 'Shape', code) + + +def test_pydantic_models(directory): + pydantic = pytest.importorskip('pydantic') + + class Message(pydantic.BaseModel): + role: str + content: str + + Message.__qualname__ = 'Message' + setattr(sys.modules[__name__], 'Message', Message) + disk = dc.SafeDisk.with_types(Message) + + try: + with dc.Cache(directory, disk=disk) as cache: + cache['message'] = Message(role='user', content='hi') + assert cache['message'] == Message(role='user', content='hi') + finally: + delattr(sys.modules[__name__], 'Message') + + +def test_restricted_disk_subclass_pattern(directory): + # The pattern dspy uses for its own restricted cache: a Disk subclass + # built with type() that overrides fetch for MODE_PICKLE and is passed + # to FanoutCache(disk=...). The fork keeps the hooks it relies on. + class Refused(Exception): + pass + + class Restricted(pickle.Unpickler): + allowed = frozenset({(models.__name__, 'Point')}) + + def find_class(self, module, name): + if (module, name) in self.allowed: + return super().find_class(module, name) + raise Refused('%s.%s' % (module, name)) + + class RestrictedDisk(dc.Disk): + def fetch(self, mode, filename, value, read): + if mode == MODE_PICKLE: + if value is None: + with open(op.join(self._directory, filename), 'rb') as f: + return Restricted(f).load() + return Restricted(io.BytesIO(value)).load() + return super().fetch(mode, filename, value, read) + + disk = type('RestrictedDisk', (RestrictedDisk,), {}) + + with dc.FanoutCache(directory, shards=2, disk=disk) as cache: + cache['point'] = models.Point(1, 2) + assert cache['point'] == models.Point(1, 2) + cache['shape'] = models.Shape('line', []) + with pytest.raises(Refused): + cache['shape'] # pylint: disable=pointless-statement