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 bdfff167..4967ecac 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. @@ -343,7 +345,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] @@ -442,7 +444,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: @@ -494,7 +496,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: @@ -506,7 +508,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: @@ -533,7 +535,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): @@ -570,7 +572,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]]: @@ -596,7 +598,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 b33ec899..4d37c34f 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) @@ -89,6 +102,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)