diff --git a/pyproject.toml b/pyproject.toml index 0ea64c5a..5198bbe2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,7 @@ module = [ "fsspec.*", "funcy", "diskcache", + "diskcache.core", "pygtrie", "dictdiffer", "shortuuid.*", diff --git a/src/dvc_data/hashfile/cache.py b/src/dvc_data/hashfile/cache.py index 7626d56d..6f697a84 100644 --- a/src/dvc_data/hashfile/cache.py +++ b/src/dvc_data/hashfile/cache.py @@ -1,4 +1,6 @@ +import json import os +import os.path as op import pickle import sqlite3 from collections.abc import Iterable, Iterator, Sequence @@ -12,9 +14,33 @@ Index, # noqa: F401 Timeout, # noqa: F401 ) +from diskcache.core import MODE_PICKLE, UNKNOWN from dvc_data.compat import batched +# Protocol 2+ pickles start with the PROTO opcode; no JSON encoding can. +PICKLE_PROTO_OPCODE = b"\x80" + +# diskcache stores SQLite-native values as-is; everything else it pickles. +_SQLITE_INT_MIN = -9223372036854775808 +_SQLITE_INT_MAX = 9223372036854775807 + + +def _stored_raw(value: Any, min_file_size: int) -> bool: + """Whether diskcache would store `value` without pickling it. + + Mirrors the type dispatch in `diskcache.Disk.store`, which uses exact type + checks -- note that `bool` is therefore *not* covered by the `int` case. + """ + type_value = type(value) + if type_value is bytes: + return True + if type_value is str: + return len(value) < min_file_size + if type_value is float: + return True + return type_value is int and _SQLITE_INT_MIN <= value <= _SQLITE_INT_MAX + class DiskError(Exception): def __init__(self, directory: str, type: str) -> None: # noqa: A002 @@ -23,6 +49,10 @@ def __init__(self, directory: str, type: str) -> None: # noqa: A002 super().__init__(f"Could not open disk '{type}' in {directory}") +class LegacyPickleError(Exception): + """A cache entry was written by an older, pickle-serializing dvc-data.""" + + def translate_pickle_error(fn): @wraps(fn) def wrapped(self, *args, **kwargs): @@ -49,6 +79,59 @@ class Disk(_Disk): fetch = translate_pickle_error(_Disk.fetch) +class JSONDisk(Disk): + """Serialize values as JSON rather than pickle. + + diskcache pickles any value that is not a str, int, float, or bytes, and + unpickles it on read. That makes the cache directory a code-execution + surface: anything able to write into it can hand a poisoned payload to the + next reader (CVE-2025-69872 / GHSA-w8v5-vhqr-4h9v, unfixed upstream -- + 5.6.3 is the newest release and both proposed fixes were declined). + + dvc-data does not need pickle's expressiveness. Every value it caches is a + dict, bool, tuple of numbers, or an already-JSON-encoded string, so JSON + covers the whole domain with no code-execution primitive on read. + + Only values the base class would have pickled are re-encoded; str, int, + float, and bytes keep their existing raw storage, so `HashesCache` -- which + writes through this disk but reads back with raw SQL -- is untouched. Keys + are likewise left to the base class, which stores dvc-data's string keys + raw and is therefore already pickle-free. + + JSON payloads reuse the MODE_PICKLE slot, so entries written by an older + dvc-data are still recognised. The two are told apart by content: a + protocol-2+ pickle starts with the PROTO opcode (0x80), which no JSON + encoding can begin with. + """ + + def store(self, value, read, key=UNKNOWN): + if read or _stored_raw(value, self.min_file_size): + # A file-like value is streamed verbatim; str/int/float/bytes are + # stored raw. Neither is pickled, so leave both to the base class. + return super().store(value, read, key=key) + + data = json.dumps(value, separators=(",", ":"), sort_keys=True).encode() + size, _, filename, db_value = super().store(data, False, key=key) + # Reuse the MODE_PICKLE slot so `fetch` knows to decode the payload. + return size, MODE_PICKLE, filename, db_value + + def fetch(self, mode, filename, value, read): + if mode != MODE_PICKLE: + return super().fetch(mode, filename, value, read) + + if value is None: + with open(op.join(self._directory, filename), "rb") as reader: + data = reader.read() + else: + data = bytes(value) + + if data[:1] == PICKLE_PROTO_OPCODE: + # Written by a pre-JSON dvc-data. Refuse to unpickle it; Cache + # turns this into a miss so the entry is recomputed. + raise LegacyPickleError + return json.loads(data) + + class Cache(diskcache.Cache): """Extended to handle pickle errors and use a constant pickle protocol.""" @@ -56,7 +139,7 @@ def __init__( self, directory: Optional[str] = None, timeout: int = 60, - disk: _Disk = Disk, + disk: _Disk = JSONDisk, type: Optional[str] = None, # noqa: A002 **settings: Any, ) -> None: @@ -68,6 +151,58 @@ def __init__( def __getstate__(self): return (*super().__getstate__(), self._type) + def _evict_legacy(self, key) -> None: + try: + super().__delitem__(key, retry=True) + except KeyError: + pass + + def get( + self, + key, + default=None, + read=False, + expire_time=False, + tag=False, + retry=False, + ): + """Return the value for `key`, treating legacy pickled entries as misses. + + These caches all live under `tmp_dir` and are regenerable, so dropping + an entry costs a recomputation rather than data. + """ + try: + return super().get( + key, + default=default, + read=read, + expire_time=expire_time, + tag=tag, + retry=retry, + ) + except LegacyPickleError: + self._evict_legacy(key) + if expire_time and tag: + return default, None, None + if expire_time or tag: + return default, None + return default + + def __getitem__(self, key): + try: + return super().__getitem__(key) + except LegacyPickleError: + self._evict_legacy(key) + raise KeyError(key) from None + + def __contains__(self, key) -> bool: + # `in` must agree with reads: a legacy entry is not readable. + try: + return super().__contains__(key) + except LegacyPickleError: + self._evict_legacy(key) + return False + class HashesCache(Cache): SUPPORTS_UPSERT = sqlite3.sqlite_version_info >= (3, 24, 0) diff --git a/src/dvc_data/hashfile/db/index.py b/src/dvc_data/hashfile/db/index.py index 7272b009..e67c1d2b 100644 --- a/src/dvc_data/hashfile/db/index.py +++ b/src/dvc_data/hashfile/db/index.py @@ -117,7 +117,12 @@ def __contains__(self, hash_: str) -> bool: def dir_hashes(self) -> Iterator[str]: """Iterate over .dir hashes stored in the index.""" - yield from (hash_ for hash_, is_dir in self.index.items() if is_dir) + # Read via get() rather than items(): an entry written by an older, + # pickle-serializing dvc-data is dropped as a miss, and iterating + # items() would race that eviction and raise KeyError. + for hash_ in list(self.index): + if self.index.get(hash_): + yield hash_ def clear(self) -> None: """Clear this index (to force re-indexing later).""" diff --git a/src/dvc_data/hashfile/state.py b/src/dvc_data/hashfile/state.py index 13b0dc0e..a790c735 100644 --- a/src/dvc_data/hashfile/state.py +++ b/src/dvc_data/hashfile/state.py @@ -307,16 +307,25 @@ def get_unused_links(self, used, fs): unused = [] with self.links as ref: - for relative_path in ref: + # Materialize the keys: reading an entry written by an older, + # pickle-serializing dvc-data drops it, which would mutate the + # cache while we iterate it. + for relative_path in list(ref): path = os.path.join(self.root_dir, relative_path) if path in used or not fs.exists(path): continue + entry = ref.get(relative_path) + if entry is None: + continue + inode = get_inode(path) mtime, _ = get_mtime_and_size(path, fs, self.ignore) - if ref[relative_path] == (inode, mtime): + # The link cache round-trips this pair through JSON, which has + # no tuple type, so compare as a sequence rather than by type. + if tuple(entry) == (inode, mtime): logger.debug("Removing '%s' as unused link.", path) unused.append(relative_path) diff --git a/tests/hashfile/test_cache.py b/tests/hashfile/test_cache.py index 0e4c9bb2..34345038 100644 --- a/tests/hashfile/test_cache.py +++ b/tests/hashfile/test_cache.py @@ -1,10 +1,17 @@ +import json +import os +import pathlib import pickle +import sqlite3 +from contextlib import closing from os import fspath from typing import Any +import diskcache import pytest +from diskcache.core import MODE_PICKLE -from dvc_data.hashfile.cache import Cache, DiskError, HashesCache +from dvc_data.hashfile.cache import Cache, Disk, DiskError, HashesCache def set_value(cache: Cache, key: str, value: Any) -> Any: @@ -12,8 +19,27 @@ def set_value(cache: Cache, key: str, value: Any) -> Any: return cache[key] +def pickled_rows(cache: diskcache.Cache) -> list[str]: + """Keys whose stored value is a pickle stream. + + Protocol 2+ pickles start with the PROTO opcode (0x80), which no JSON + encoding can begin with. + """ + rows = cache._sql("SELECT key, value FROM Cache").fetchall() + return sorted( + key + for key, value in rows + if isinstance(value, (bytes, memoryview)) and bytes(value)[:1] == b"\x80" + ) + + @pytest.mark.parametrize("disk_type", [None, "test"]) def test_pickle_protocol_error(tmp_path, disk_type): + """An unusable pickle protocol is still reported as a DiskError. + + Keys are pickled by the base class, so this path stays live even though + values are now serialized as JSON. + """ directory = tmp_path / "test" cache = Cache( fspath(directory), @@ -21,12 +47,23 @@ def test_pickle_protocol_error(tmp_path, disk_type): type=disk_type, ) with pytest.raises(DiskError) as exc, cache as cache: - set_value(cache, "key", ("value1", "value2")) + set_value(cache, ("tuple", "key"), "value") assert exc.value.directory == fspath(directory) assert exc.value.type == "test" assert f"Could not open disk 'test' in {directory}" == str(exc.value) +@pytest.mark.parametrize("disk_type", [None, "test"]) +def test_pickle_protocol_does_not_affect_values(tmp_path, disk_type): + """Values no longer depend on the pickle protocol being usable.""" + with Cache( + fspath(tmp_path / "test"), + disk_pickle_protocol=pickle.HIGHEST_PROTOCOL + 1, + type=disk_type, + ) as cache: + assert set_value(cache, "key", {"loaded": True}) == {"loaded": True} + + @pytest.mark.parametrize( "proto_a, proto_b", [ @@ -34,7 +71,13 @@ def test_pickle_protocol_error(tmp_path, disk_type): (pickle.HIGHEST_PROTOCOL, pickle.HIGHEST_PROTOCOL - 1), ], ) -def test_pickle_backwards_compat(tmp_path, proto_a, proto_b): +def test_readable_across_pickle_protocols(tmp_path, proto_a, proto_b): + """A cache stays readable across `disk_pickle_protocol` settings. + + Values are serialized as JSON, so the pickle protocol no longer applies to + them at all -- which is what makes them readable either way. Tuples come + back as lists because JSON has no tuple type. + """ with Cache( directory=fspath(tmp_path / "test"), disk_pickle_protocol=proto_a, @@ -44,9 +87,8 @@ def test_pickle_backwards_compat(tmp_path, proto_a, proto_b): directory=fspath(tmp_path / "test"), disk_pickle_protocol=proto_b, ) as cache: - assert cache["key"] == ("value1", "value2") - set_value(cache, "key", ("value3", "value4")) - assert cache["key"] == ("value3", "value4") + assert cache["key"] == ["value1", "value2"] + assert set_value(cache, "key", ("value3", "value4")) == ["value3", "value4"] def test_hashes_cache(tmp_path): @@ -88,3 +130,151 @@ def test_hashes_cache_update(tmp_path, upsert): ("key1", "value1"), ("key2", "value2"), ] + + +# The values dvc-data actually caches, per its four Cache/Index call sites: +# hashfile/state.py (links), hashfile/db/index.py (ODB index), index/serialize.py. +DVC_CACHED_VALUES = [ + pytest.param((12345, 1699999999.5), id="links-inode-mtime"), + pytest.param(True, id="index-is-dir"), + pytest.param(False, id="index-is-file"), + pytest.param( + { + "meta": {"size": 3, "isexec": True}, + "hash_info": {"md5": "x"}, + "loaded": True, + }, + id="index-entry-dict", + ), + pytest.param({"loaded": False}, id="index-entry-minimal"), + pytest.param('{"version":1,"checksum":"a","size":12}', id="hashes-json-string"), +] + + +@pytest.mark.parametrize("value", DVC_CACHED_VALUES) +def test_values_are_not_pickled(tmp_path, value): + """No value dvc-data caches may be written as a pickle. + + diskcache unpickles on read, so a pickled value makes the cache directory a + code-execution surface for anything that can write into it + (CVE-2025-69872, unfixed upstream). Guards against a regression back to + pickle serialization. + """ + with Cache(fspath(tmp_path / "test")) as cache: + cache["key"] = value + assert pickled_rows(cache) == [] + + +@pytest.mark.parametrize("value", DVC_CACHED_VALUES) +def test_values_round_trip(tmp_path, value): + """JSON serialization must preserve every value shape dvc-data stores. + + Tuples come back as lists -- JSON has no tuple type -- so compare + structurally. `state.get_unused_links` accounts for this explicitly. + """ + with Cache(fspath(tmp_path / "test")) as cache: + cache["key"] = value + got = cache["key"] + + if isinstance(value, tuple): + assert tuple(got) == value + else: + assert got == value + + +def test_stored_value_is_json(tmp_path): + with Cache(fspath(tmp_path / "test")) as cache: + cache["key"] = {"loaded": True, "meta": {"size": 3}} + ((raw,),) = cache._sql("SELECT value FROM Cache").fetchall() + assert json.loads(bytes(raw)) == {"loaded": True, "meta": {"size": 3}} + + +def _write_legacy_pickled_entries(directory: str) -> None: + """Write entries the way a pre-JSON dvc-data did, via the pickling Disk.""" + with diskcache.Cache(directory, disk=Disk, disk_pickle_protocol=4) as cache: + cache.disk._type = cache._type = "test" + cache["tuple"] = (1, 2.5) + cache["bool"] = True + + +def test_legacy_pickled_entry_is_a_miss(tmp_path): + """A cache written by an older dvc-data must not be unpickled. + + These caches live under `tmp_dir` and are regenerable, so a legacy entry + degrades to a miss (and is evicted) rather than raising or executing. + """ + directory = fspath(tmp_path / "test") + _write_legacy_pickled_entries(directory) + + with Cache(directory) as cache: + assert pickled_rows(cache) == ["bool", "tuple"] + + assert cache.get("tuple", "MISS") == "MISS" + assert cache.get("bool", "MISS") == "MISS" + + # Reads evict the poisoned rows, so nothing pickled survives. + assert pickled_rows(cache) == [] + + +def test_legacy_pickled_entry_raises_key_error(tmp_path): + directory = fspath(tmp_path / "test") + _write_legacy_pickled_entries(directory) + + with Cache(directory) as cache: + with pytest.raises(KeyError): + cache["tuple"] + assert "tuple" not in cache + + +def test_legacy_entry_is_replaced_on_write(tmp_path): + directory = fspath(tmp_path / "test") + _write_legacy_pickled_entries(directory) + + with Cache(directory) as cache: + cache["tuple"] = (3, 4.5) + assert tuple(cache["tuple"]) == (3, 4.5) + assert pickled_rows(cache) == ["bool"] + + +def test_hashes_cache_is_unaffected(tmp_path): + """HashesCache reads and writes via raw SQL, bypassing the disk layer.""" + with HashesCache(fspath(tmp_path / "test")) as cache: + cache.set("key", '{"version":1}') + cache.set_many((("k2", '{"version":2}'),)) + assert cache.get("key") == '{"version":1}' + assert list(cache.get_many(("k2",))) == [("k2", '{"version":2}')] + assert pickled_rows(cache) == [] + + +def test_poisoned_entry_is_not_executed(tmp_path): + """A payload injected into the cache directory must not be unpickled. + + This is the CVE-2025-69872 attack: an attacker who can write to the cache + directory replaces a value with a pickle whose `__reduce__` runs a command, + and the next process to read that key executes it. Against stock diskcache + this succeeds; the JSON disk refuses the entry instead. + """ + marker = tmp_path / "executed" + + class Evil: + def __reduce__(self): + return (pathlib.Path.touch, (marker,)) + + directory = fspath(tmp_path / "test") + with Cache(directory) as cache: + cache["key"] = {"placeholder": True} + + # Overwrite the stored value with a pickle payload, as an attacker with + # write access to the cache directory could. + con = sqlite3.connect(os.path.join(directory, "cache.db")) + with closing(con): + con.execute( + "UPDATE Cache SET mode = ?, value = ?, filename = NULL WHERE key = 'key'", + (MODE_PICKLE, sqlite3.Binary(pickle.dumps(Evil(), protocol=4))), + ) + con.commit() + + with Cache(directory) as cache: + assert cache.get("key", "MISS") == "MISS" + + assert not marker.exists(), "the poisoned payload was executed" diff --git a/tests/hashfile/test_checkout.py b/tests/hashfile/test_checkout.py index 373d0b34..642a0fc4 100644 --- a/tests/hashfile/test_checkout.py +++ b/tests/hashfile/test_checkout.py @@ -217,7 +217,9 @@ def test_recheckout_old_obj(tmp_path, relink): def get_inode_and_mtime(path): - return inode(path), get_mtime_and_size(os.fspath(path), localfs)[0] + # A list, not a tuple: the link cache serializes as JSON, which has no + # tuple type, so `state.links[...]` reads back as a list. + return [inode(path), get_mtime_and_size(os.fspath(path), localfs)[0]] def test_checkout_save_link_dir(request, tmp_path): diff --git a/tests/hashfile/test_db_index.py b/tests/hashfile/test_db_index.py index 5ed6f452..11a7ee82 100644 --- a/tests/hashfile/test_db_index.py +++ b/tests/hashfile/test_db_index.py @@ -42,3 +42,30 @@ def test_intersection(index): expected = {str(i) for i in range(1000)} index.update([], hashes) assert set(index.intersection(expected)) == expected + + +def test_legacy_pickled_index_is_dropped(tmp_path): + """An index written by an older, pickle-serializing dvc-data is discarded. + + The index is a regenerable cache under `tmp_dir`, so refusing to unpickle + it costs a re-index rather than data. `dir_hashes` must not raise while the + stale entries are being evicted. + """ + import diskcache + + from dvc_data.hashfile.cache import Disk + + index_dir = tmp_path / ObjectDBIndex.INDEX_DIR / "foo" + index_dir.mkdir(parents=True) + with diskcache.Cache(str(index_dir), disk=Disk, disk_pickle_protocol=4) as cache: + cache.disk._type = cache._type = "index" + cache["1234.dir"] = True + cache["5678"] = False + + with closing(ObjectDBIndex(tmp_path, "foo")) as index: + assert list(index.dir_hashes()) == [] + assert index.intersection({"1234.dir", "5678"}) is not None + + # Re-indexing repopulates it, now without pickle. + index.update({"1234.dir"}, {"5678"}) + assert set(index.dir_hashes()) == {"1234.dir"} diff --git a/tests/hashfile/test_state.py b/tests/hashfile/test_state.py index 111e672d..ec37660d 100644 --- a/tests/hashfile/test_state.py +++ b/tests/hashfile/test_state.py @@ -208,7 +208,9 @@ def test_state_many(tmp_path, state: State): def test_set_link(tmp_path, state): state.set_link(tmp_path / "foo", 42, "mtime") - assert state.links["foo"] == (42, "mtime") + # The link cache serializes as JSON, which has no tuple type, so the pair + # reads back as a list. `get_unused_links` compares it as a sequence. + assert tuple(state.links["foo"]) == (42, "mtime") def test_state_noop(tmp_path): @@ -249,7 +251,9 @@ def _get_inode_mtime(path): return inode(path), get_mtime_and_size(path, fs)[0] assert len(state.links) == 3 - assert {k: state.links[k] for k in state.links} == { + # Values come back as lists: the link cache serializes as JSON, which has + # no tuple type. + assert {k: tuple(state.links[k]) for k in state.links} == { "foo": _get_inode_mtime(foo), "bar": _get_inode_mtime(bar), "dataset": _get_inode_mtime(dataset), @@ -272,3 +276,36 @@ def _get_inode_mtime(path): assert not foo.exists() assert not bar.exists() assert not dataset.exists() + + +def test_legacy_pickled_links_are_dropped(tmp_path): + """A links cache written by an older dvc-data is discarded, not unpickled. + + `get_unused_links` iterates the cache, and reading a legacy entry evicts + it, so the loop must tolerate the mutation rather than raising KeyError. + """ + import diskcache + + from dvc_data.hashfile.cache import Disk + + tmp = tmp_path / "tmp" + links_dir = tmp / "links" + links_dir.mkdir(parents=True) + + foo = tmp_path / "foo" + foo.write_text("foo content", encoding="utf-8") + + with diskcache.Cache(str(links_dir), disk=Disk, disk_pickle_protocol=4) as cache: + cache.disk._type = cache._type = "links" + cache["foo"] = (inode(os.fspath(foo)), "mtime") + + fs = LocalFileSystem() + state = State(root_dir=os.fspath(tmp_path), tmp_dir=os.fspath(tmp)) + try: + assert state.get_unused_links([], fs) == [] + + # Re-recording the link works, and is stored without pickle. + state.save_link(os.fspath(foo), fs) + assert set(state.get_unused_links([], fs)) == {"foo"} + finally: + state.close()