Restrict backend key unpickling

datasketch/aio/lsh.py

@@ -5,14 +5,14 @@
 """
 
 import asyncio
-import pickle
 from itertools import chain
 from typing import Optional
 
 from datasketch.aio.storage import (
     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)

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

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):

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

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

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)