From 7fb78f7fc0e86a33b308490b7cc9093b2040454f Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sat, 18 Jul 2026 23:46:27 -0400 Subject: [PATCH 1/2] Restrict backend key unpickling --- datasketch/aio/lsh.py | 17 +++++---- datasketch/key_serialization.py | 65 +++++++++++++++++++++++++++++++++ datasketch/lsh.py | 20 +++++----- datasketch/lshensemble.py | 3 +- docs/lsh.rst | 12 ++++++ test/test_lsh.py | 52 ++++++++++++++++++++++++++ 6 files changed, 151 insertions(+), 18 deletions(-) create mode 100644 datasketch/key_serialization.py diff --git a/datasketch/aio/lsh.py b/datasketch/aio/lsh.py index 634464dc..c00bbc63 100644 --- a/datasketch/aio/lsh.py +++ b/datasketch/aio/lsh.py @@ -5,7 +5,6 @@ """ import asyncio -import pickle from itertools import chain from typing import Optional @@ -13,6 +12,7 @@ async_ordered_storage, async_unordered_storage, ) +from datasketch.key_serialization import dumps_key, loads_key from datasketch.lsh import _optimal_param from datasketch.minhash import _check_scheme_consistency from datasketch.storage import _random_name, unordered_storage @@ -31,7 +31,8 @@ class AsyncMinHashLSH: If storage_config is None aiomongo storage will be used. :param prepickle (bool, optional): If True, all keys are pickled to bytes before insertion. If None, a default value is chosen based on the - `storage_config`. + `storage_config`. For safe deserialization, keys must contain only + primitive built-in types; custom classes are rejected. For example usage see :ref:`minhash_lsh_async`. Example of supported storage configuration: @@ -260,7 +261,7 @@ async def _insert(self, key, minhash, check_duplication=True, buffer=False): "Either pass bytes keys or use prepickle=True for automatic serialization." ) if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) # `key` is already pickled at this point under prepickle=True; call the # storage primitive directly so we don't re-pickle through has_key(). @@ -286,13 +287,13 @@ async def query(self, minhash): ) candidates = frozenset(chain.from_iterable(await asyncio.gather(*fs))) if self.prepickle: - return [pickle.loads(key) for key in candidates] + return [loads_key(key) for key in candidates] return list(candidates) async def has_key(self, key): """See :class:`datasketch.MinHashLSH`.""" if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) return await self.keys.has_key(key) async def remove(self, key): @@ -301,7 +302,7 @@ async def remove(self, key): async def _remove(self, key, buffer=False): if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) # `key` is already pickled here; call storage primitives directly so # the existence check, lookup, and deletes all use the stored form. @@ -339,7 +340,7 @@ async def _query_b(self, minhash, b): fs.append(hashtable.get(H)) candidates = set(chain.from_iterable(await asyncio.gather(*fs))) if self.prepickle: - return {pickle.loads(key) for key in candidates} + return {loads_key(key) for key in candidates} return candidates async def get_counts(self): @@ -352,7 +353,7 @@ async def get_subset_counts(self, *keys): # Keys in storage are pickled when prepickle is enabled, so we have to # pickle the query keys to match the stored representation. if self.prepickle: - keys = tuple(pickle.dumps(key) for key in keys) + keys = tuple(dumps_key(key) for key in keys) key_set = list(set(keys)) hashtables = [unordered_storage({"type": "dict"}) for _ in range(self.b)] Hss = await self.keys.getmany(*key_set) diff --git a/datasketch/key_serialization.py b/datasketch/key_serialization.py new file mode 100644 index 00000000..a27c1852 --- /dev/null +++ b/datasketch/key_serialization.py @@ -0,0 +1,65 @@ +"""Restricted serialization helpers for LSH keys stored in backends.""" + +from __future__ import annotations + +import io +import pickle +import pickletools +from collections.abc import Hashable +from typing import Any + +_FORBIDDEN_OPCODES = { + "BINPERSID", + "BUILD", + "EXT1", + "EXT2", + "EXT4", + "GLOBAL", + "INST", + "NEWOBJ", + "NEWOBJ_EX", + "OBJ", + "PERSID", + "REDUCE", + "STACK_GLOBAL", +} + + +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler that forbids importing every global callable and class.""" + + def find_class(self, module: str, name: str) -> Any: + global_name = f"{module}.{name}" + raise pickle.UnpicklingError(f"global {global_name!r} is forbidden in an LSH key") + + def persistent_load(self, pid: Any) -> Any: + raise pickle.UnpicklingError("persistent IDs are forbidden in an LSH key") + + +def loads_key(payload: bytes) -> Hashable: + """Load a backend key without allowing global object construction.""" + try: + for opcode, _argument, _position in pickletools.genops(payload): + if opcode.name in _FORBIDDEN_OPCODES: + raise pickle.UnpicklingError(f"opcode {opcode.name!r} is forbidden in an LSH key") + key = _RestrictedUnpickler(io.BytesIO(payload)).load() + except (EOFError, AttributeError, TypeError, ValueError) as error: + raise pickle.UnpicklingError("invalid serialized LSH key") from error + if not isinstance(key, Hashable): + raise pickle.UnpicklingError("serialized LSH key is not hashable") + return key + + +def dumps_key(key: Hashable) -> bytes: + """Serialize a key and ensure the restricted loader can read it back.""" + if not isinstance(key, Hashable): + raise TypeError("LSH keys must be hashable") + payload = pickle.dumps(key) + try: + loads_key(payload) + except pickle.UnpicklingError as error: + raise TypeError( + "prepickle=True supports keys made from primitive built-in types only; " + f"custom type {type(key).__name__!r} cannot be stored safely" + ) from error + return payload diff --git a/datasketch/lsh.py b/datasketch/lsh.py index 64e1360d..85868816 100644 --- a/datasketch/lsh.py +++ b/datasketch/lsh.py @@ -1,12 +1,12 @@ from __future__ import annotations -import pickle import struct from collections.abc import Hashable from typing import Callable, List, Optional, Union from scipy.integrate import quad as integrate +from datasketch.key_serialization import dumps_key, loads_key from datasketch.minhash import MinHash, _check_scheme_consistency from datasketch.storage import ( OrderedStorage, @@ -77,7 +77,9 @@ class MinHashLSH: set this, you will be responsible for ensuring there are no key collisions. prepickle (Optional[bool]): If True, all keys are pickled to bytes before insertion. If not specified, a default value is chosen based on the - `storage_config`. + `storage_config`. For safe deserialization, keys must contain only + primitive built-in types such as strings, integers, floats, bytes, + tuples, and frozensets; custom classes are rejected. hashfunc (Optional[Callable[[bytes], bytes]]): If a hash function is provided it will be used to compress the index keys to reduce the memory footprint. This could cause a higher false positive rate. @@ -345,7 +347,7 @@ def _insert( "Either pass bytes keys or use prepickle=True for automatic serialization." ) if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) if check_duplication and key in self.keys: raise ValueError("The given key already exists") Hs = [self._H(minhash.hashvalues[start:end]) for start, end in self.hashranges] @@ -452,7 +454,7 @@ def query(self, minhash) -> list[Hashable]: for key in hashtable.get(H): candidates.add(key) if self.prepickle: - return [pickle.loads(key) for key in candidates] + return [loads_key(key) for key in candidates] return list(candidates) def add_to_query_buffer(self, minhash: Union[MinHash, WeightedMinHash]) -> None: @@ -504,7 +506,7 @@ def collect_query_buffer(self) -> list[Hashable]: candidates = set.intersection(*per_query_result_sets) if self.prepickle: - return [pickle.loads(key) for key in candidates] + return [loads_key(key) for key in candidates] return list(candidates) def __contains__(self, key: Hashable) -> bool: @@ -516,7 +518,7 @@ def __contains__(self, key: Hashable) -> bool: """ if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) return key in self.keys def remove(self, key: Hashable) -> None: @@ -543,7 +545,7 @@ def _remove(self, key: Hashable, buffer: bool = False) -> None: """ if self.prepickle: - key = pickle.dumps(key) + key = dumps_key(key) if key not in self.keys: raise ValueError("The given key does not exist") for H, hashtable in zip(self.keys[key], self.hashtables): @@ -580,7 +582,7 @@ def _query_b(self, minhash, b): for key in hashtable[H]: candidates.add(key) if self.prepickle: - return {pickle.loads(key) for key in candidates} + return {loads_key(key) for key in candidates} return candidates def get_counts(self) -> list[dict[Hashable, int]]: @@ -606,7 +608,7 @@ def get_subset_counts(self, *keys: Hashable) -> list[dict[Hashable, int]]: list: a list of dictionaries. """ - key_set = [pickle.dumps(key) for key in set(keys)] if self.prepickle else list(set(keys)) + key_set = [dumps_key(key) for key in set(keys)] if self.prepickle else list(set(keys)) hashtables = [unordered_storage({"type": "dict"}) for _ in range(self.b)] Hss = self.keys.getmany(*key_set) for key, Hs in zip(key_set, Hss): diff --git a/datasketch/lshensemble.py b/datasketch/lshensemble.py index b25cca4a..b9f5b21e 100644 --- a/datasketch/lshensemble.py +++ b/datasketch/lshensemble.py @@ -86,7 +86,8 @@ class MinHashLSHEnsemble: set this, you will be responsible for ensuring there are no key collisions. prepickle (Optional[bool]): If True, all keys are pickled to bytes before insertion. If None, a default value is chosen based on the - `storage_config`. + `storage_config`. For safe deserialization, keys must contain only + primitive built-in types; custom classes are rejected. Note: Using more partitions (`num_part`) leads to better accuracy, at the diff --git a/docs/lsh.rst b/docs/lsh.rst index 3f33f000..0adbeeeb 100644 --- a/docs/lsh.rst +++ b/docs/lsh.rst @@ -259,6 +259,12 @@ or keyspace (and thus DROP existing ones), set `drop_tables` and `drop_keyspace` Like the Redis counterpart, you can use insert sessions to reduce the number of network calls during bulk insertion. +When ``prepickle=True``, backend keys are decoded with a restricted unpickler +that cannot import global functions or classes. Keys made from primitive +built-in types (for example strings, integers, floats, bytes, tuples, and frozensets) +are supported; custom class instances are rejected. Existing indexes that +used custom object keys must migrate those keys to supported built-in values. + Connecting to Existing MinHash LSH ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -269,6 +275,12 @@ share it across multiple processes. There are two ways to do it: The recommended way is to use "pickling". The MinHash LSH object is serializable so you can call `pickle`: +.. warning:: + + Only unpickle a serialized MinHash LSH object from a trusted source. + This is separate from the restricted decoding applied to keys returned by + an external storage backend. + .. code:: python import pickle diff --git a/test/test_lsh.py b/test/test_lsh.py index a9f238ec..1bc1a01a 100644 --- a/test/test_lsh.py +++ b/test/test_lsh.py @@ -9,6 +9,19 @@ from datasketch.minhash import MinHash from datasketch.weighted_minhash import WeightedMinHashGenerator +backend_payload_executed = False + + +def execute_backend_payload(): + global backend_payload_executed + backend_payload_executed = True + return "compromised" + + +class UnsafeKey: + def __reduce__(self): + return execute_backend_payload, () + def fake_redis(**kwargs): redis = mockredis.mock_redis_client(**kwargs) @@ -97,6 +110,45 @@ def test_query(self): m3 = MinHash(18) self.assertRaises(ValueError, lsh.query, m3) + def test_prepickle_round_trips_primitive_builtin_keys(self): + lsh = MinHashLSH(threshold=0.5, num_perm=16, prepickle=True) + minhash = MinHash(16) + minhash.update(b"value") + keys = ["string", 42, b"bytes", ("tuple", 3), frozenset({"set", 4}), None] + for key in keys: + lsh.insert(key, minhash) + + self.assertEqual(set(lsh.query(minhash)), set(keys)) + + def test_prepickle_rejects_custom_key_types_before_insert(self): + global backend_payload_executed + backend_payload_executed = False + lsh = MinHashLSH(threshold=0.5, num_perm=16, prepickle=True) + minhash = MinHash(16) + minhash.update(b"value") + + with self.assertRaisesRegex(TypeError, "primitive built-in types"): + lsh.insert(UnsafeKey(), minhash) + + self.assertFalse(backend_payload_executed) + self.assertEqual(len(lsh.keys), 0) + + def test_query_rejects_executable_pickle_from_backend(self): + global backend_payload_executed + backend_payload_executed = False + lsh = MinHashLSH(threshold=0.5, num_perm=16, prepickle=True) + minhash = MinHash(16) + minhash.update(b"value") + lsh.insert("safe", minhash) + + start, end = lsh.hashranges[0] + band_hash = lsh._H(minhash.hashvalues[start:end]) + lsh.hashtables[0].insert(band_hash, pickle.dumps(UnsafeKey())) + + with self.assertRaisesRegex(pickle.UnpicklingError, "forbidden"): + lsh.query(minhash) + self.assertFalse(backend_payload_executed) + def test_query_buffer(self): lsh = MinHashLSH(threshold=0.5, num_perm=16) m1 = MinHash(16) From a43a2209680ffd05d5b033bed529459919f190bb Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sun, 19 Jul 2026 01:05:29 -0400 Subject: [PATCH 2/2] Harden restricted key validation --- datasketch/key_serialization.py | 20 ++++++++- test/test_lsh.py | 76 ++++++++++++++++++++++++++++++++- 2 files changed, 93 insertions(+), 3 deletions(-) diff --git a/datasketch/key_serialization.py b/datasketch/key_serialization.py index a27c1852..c2b5b415 100644 --- a/datasketch/key_serialization.py +++ b/datasketch/key_serialization.py @@ -25,6 +25,22 @@ } +def _is_hashable(key: Any) -> bool: + """Return whether ``key`` can actually be hashed. + + ``isinstance(key, Hashable)`` is insufficient for containers such as a + tuple that contains a list: the tuple exposes ``__hash__`` but calling it + still raises ``TypeError``. + """ + if not isinstance(key, Hashable): + return False + try: + hash(key) + except TypeError: + return False + return True + + class _RestrictedUnpickler(pickle.Unpickler): """Unpickler that forbids importing every global callable and class.""" @@ -45,14 +61,14 @@ def loads_key(payload: bytes) -> Hashable: key = _RestrictedUnpickler(io.BytesIO(payload)).load() except (EOFError, AttributeError, TypeError, ValueError) as error: raise pickle.UnpicklingError("invalid serialized LSH key") from error - if not isinstance(key, Hashable): + if not _is_hashable(key): raise pickle.UnpicklingError("serialized LSH key is not hashable") return key def dumps_key(key: Hashable) -> bytes: """Serialize a key and ensure the restricted loader can read it back.""" - if not isinstance(key, Hashable): + if not _is_hashable(key): raise TypeError("LSH keys must be hashable") payload = pickle.dumps(key) try: diff --git a/test/test_lsh.py b/test/test_lsh.py index 1bc1a01a..9b711b44 100644 --- a/test/test_lsh.py +++ b/test/test_lsh.py @@ -1,10 +1,13 @@ +import io import pickle import unittest -from unittest.mock import patch +from unittest.mock import AsyncMock, Mock, patch import mockredis import numpy as np +from datasketch.aio import AsyncMinHashLSH +from datasketch.key_serialization import _RestrictedUnpickler, dumps_key, loads_key from datasketch.lsh import MinHashLSH from datasketch.minhash import MinHash from datasketch.weighted_minhash import WeightedMinHashGenerator @@ -149,6 +152,35 @@ def test_query_rejects_executable_pickle_from_backend(self): lsh.query(minhash) self.assertFalse(backend_payload_executed) + def test_restricted_key_loader_defensive_errors(self): + unpickler = _RestrictedUnpickler(io.BytesIO(b"")) + with self.assertRaisesRegex(pickle.UnpicklingError, "builtins.str"): + unpickler.find_class("builtins", "str") + with self.assertRaisesRegex(pickle.UnpicklingError, "persistent IDs"): + unpickler.persistent_load("backend-reference") + + with self.assertRaisesRegex(pickle.UnpicklingError, "invalid serialized"): + loads_key(b"not a pickle") + with self.assertRaisesRegex(pickle.UnpicklingError, "not hashable"): + loads_key(pickle.dumps(["list"])) + with self.assertRaisesRegex(TypeError, "must be hashable"): + dumps_key(["list"]) + + def test_prepickle_rejects_structurally_unhashable_tuple(self): + key = ([],) + with self.assertRaisesRegex(TypeError, "must be hashable"): + dumps_key(key) + with self.assertRaisesRegex(pickle.UnpicklingError, "not hashable"): + loads_key(pickle.dumps(key)) + + def test_get_subset_counts_with_prepickled_key(self): + lsh = MinHashLSH(threshold=0.5, num_perm=16, prepickle=True) + minhash = MinHash(16) + minhash.update(b"value") + lsh.insert("safe", minhash) + + self.assertEqual(lsh.get_subset_counts("safe"), lsh.get_counts()) + def test_query_buffer(self): lsh = MinHashLSH(threshold=0.5, num_perm=16) m1 = MinHash(16) @@ -633,5 +665,47 @@ def test_pickle(self): self.assertTrue("b" in result) +class TestAsyncMinHashLSHKeySerialization(unittest.IsolatedAsyncioTestCase): + async def test_prepickle_paths_use_restricted_serialization(self): + minhash = MinHash(16) + minhash.update(b"value") + lsh = AsyncMinHashLSH( + num_perm=16, + params=(4, 4), + storage_config={"type": "aiomongo"}, + prepickle=True, + ) + serialized_key = dumps_key("safe") + band_hashes = [lsh._H(minhash.hashvalues[start:end]) for start, end in lsh.hashranges] + lsh.keys = Mock( + has_key=AsyncMock(return_value=True), + insert=AsyncMock(), + get=AsyncMock(return_value=band_hashes), + getmany=AsyncMock(return_value=[band_hashes]), + remove=AsyncMock(), + ) + lsh.hashtables = [ + Mock( + insert=AsyncMock(), + get=AsyncMock(return_value=[serialized_key]), + has_key=AsyncMock(return_value=True), + remove_val=AsyncMock(), + remove=AsyncMock(), + ) + for _ in range(lsh.b) + ] + + await lsh._insert("safe", minhash, check_duplication=False) + self.assertEqual(await lsh.query(minhash), ["safe"]) + self.assertTrue(await lsh.has_key("safe")) + self.assertEqual(await lsh._query_b(minhash, lsh.b), {"safe"}) + self.assertEqual(len(await lsh.get_subset_counts("safe")), lsh.b) + + for hashtable in lsh.hashtables: + hashtable.get.return_value = [] + await lsh.remove("safe") + lsh.keys.remove.assert_awaited_once_with(serialized_key, buffer=False) + + if __name__ == "__main__": unittest.main()