switch RestrictedUnpickler from blocklist to allowlist for pickle hardening

The previous blocklist approach could be bypassed via modules not on the list
(e.g. types, functools, _thread). The allowlist only permits modules
TensorWatch actually needs (builtins, collections, numpy, torch, pandas,
tensorwatch, and pickle internals), blocking everything else by default.

Also adds security logging for HMAC verification failures.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: sytelus

README.md

@@ -58,8 +58,11 @@ TensorWatch supports Python 3.x and is tested with PyTorch 0.4-1.x. Most feature
 > - All incoming ZMQ messages are HMAC-SHA256 verified **before**
 >   deserialization (`ZmqWrapper.verify_and_loads`). Messages with invalid
 >   signatures are rejected without being deserialized.
-> - A `RestrictedUnpickler` blocks known-dangerous modules (`os`,
->   `subprocess`, `socket`, `ctypes`, etc.) as defense-in-depth.
+> - A `RestrictedUnpickler` uses an **allowlist** of permitted modules
+>   (builtins, collections, numpy, torch, pandas, tensorwatch, and pickle
+>   internals) as defense-in-depth. Any module not explicitly approved is
+>   blocked, which is significantly stronger than the previous blocklist
+>   approach.
 > - For multi-process setups, set the `TENSORWATCH_HMAC_KEY` environment
 >   variable to a shared hex-encoded secret (e.g.
 >   `export TENSORWATCH_HMAC_KEY=$(python -c "import os; print(os.urandom(32).hex())")`).
@@ -76,9 +79,10 @@ TensorWatch supports Python 3.x and is tested with PyTorch 0.4-1.x. Most feature
 > from files. A crafted pickle file can execute arbitrary code when loaded.
 >
 > **Mitigations in place:**
-> - A `RestrictedUnpickler` blocks known-dangerous modules (`os`,
->   `subprocess`, etc.) as defense-in-depth. This is **not** a complete
->   sandbox — determined attackers may find bypasses.
+> - A `RestrictedUnpickler` uses an **allowlist** of permitted modules
+>   as defense-in-depth. Only modules TensorWatch needs (builtins,
+>   collections, numpy, torch, pandas, tensorwatch) are allowed; all
+>   others are blocked by default.
 >
 > **User responsibilities:**
 > - **Only open TensorWatch data files (`.log`, `.pkl`) that you created
@@ -97,8 +101,8 @@ TensorWatch supports Python 3.x and is tested with PyTorch 0.4-1.x. Most feature
 > | Risk | Mitigation | User Action |
 > |------|-----------|-------------|
 > | `eval()` on expressions from clients | HMAC auth + localhost binding | Never expose ports to untrusted networks |
-> | `pickle.loads()` from ZMQ | HMAC + RestrictedUnpickler | Keep HMAC key secret |
-> | `pickle.load()` from files | RestrictedUnpickler (defense-in-depth) | Only load trusted files |
+> | `pickle.loads()` from ZMQ | HMAC + allowlist RestrictedUnpickler | Keep HMAC key secret |
+> | `pickle.load()` from files | Allowlist RestrictedUnpickler (defense-in-depth) | Only load trusted files |
 > | YAML deserialization | `yaml.SafeLoader` by default | Do not override with unsafe loaders |
 > | ZMQ port exposure | Binds to `127.0.0.1` by default | Do not change to `0.0.0.0` in untrusted environments |
 

tensorwatch/safe_pickle.py

@@ -3,41 +3,52 @@
 
 """Defense-in-depth RestrictedUnpickler for TensorWatch.
 
-Blocks modules commonly exploited in pickle deserialization attacks while
-allowing the data types that TensorWatch legitimately serializes (numpy
-arrays, torch tensors, TensorWatch data classes, built-in collections, etc.).
-
-WARNING: This is NOT a complete sandbox. A determined attacker may still find
-bypass techniques. Do not load pickle data from untrusted sources.
+Uses an **allowlist** approach: only modules that TensorWatch legitimately
+needs for serialization are permitted.  Everything else is blocked by default.
+
+Allowed module families:
+  - Python builtins and standard-library data types (collections, datetime,
+    decimal, fractions, numbers, uuid)
+  - Pickle internals (_codecs, copyreg, _collections — used by the pickle
+    protocol itself)
+  - numpy, torch, pandas — data-science libraries whose objects appear in
+    StreamItem values
+  - tensorwatch — the library's own data classes
+
+WARNING: This is NOT a complete sandbox.  Pickle allowlists reduce the attack
+surface dramatically compared to blocklists, but do not eliminate all risk.
+Do not load pickle data from untrusted sources.
 """
 
 import io
 import pickle
 import logging
 
-_BLOCKED_MODULES = frozenset({
-    # OS / filesystem access
-    'os', 'posix', 'nt', 'os.path',
-    'shutil', 'pathlib',
-    'tempfile', 'glob', 'fnmatch',
-    # Process / subprocess execution
-    'subprocess', 'multiprocessing',
-    'pty', 'commands',
-    # Code compilation / execution
-    'code', 'codeop', 'compileall',
-    'importlib', 'runpy', 'pkgutil',
-    # Network
-    'socket', 'http', 'urllib', 'ftplib', 'smtplib', 'xmlrpc',
-    'socketserver', 'asyncio',
-    # Low-level / FFI
-    'ctypes', 'mmap',
-    # Interactive / debug
-    'pdb', 'profile', 'webbrowser',
-    # Signal handling
-    'signal',
+# ---------------------------------------------------------------------------
+# Allowlist — a module is permitted if its top-level package appears here.
+# Any module NOT in this set is rejected outright.
+# ---------------------------------------------------------------------------
+_ALLOWED_PREFIXES = frozenset({
+    # Python built-in / standard-library data types
+    'builtins',
+    'collections', '_collections',   # _collections holds C-accelerated types
+    'datetime',
+    'decimal',
+    'fractions',
+    'numbers',
+    'uuid',
+    # Pickle protocol internals (required for __reduce__ reconstruction)
+    'copyreg',
+    '_codecs',
+    # Data-science libraries
+    'numpy',
+    'torch',
+    'pandas',
+    # TensorWatch's own types
+    'tensorwatch',
 })
 
-# Specific names blocked from builtins module
+# Even within allowed modules, these builtins are too dangerous to permit.
 _BLOCKED_BUILTINS = frozenset({
     'eval', 'exec', 'compile', '__import__',
     'open', 'input', 'breakpoint',
@@ -46,25 +57,32 @@
     'getattr', 'setattr', 'delattr',
 })
 
+_log = logging.getLogger(__name__)
+
 
 class RestrictedUnpickler(pickle.Unpickler):
-    """Unpickler that blocks known-dangerous modules and callables.
+    """Unpickler that only allows classes from explicitly approved modules.
 
-    Allowed: numpy, torch, tensorwatch, collections, standard data types.
-    Blocked: os, subprocess, socket, ctypes, importlib, etc.
+    Uses an allowlist (not a blocklist), so unknown or newly-introduced
+    dangerous modules are blocked by default.
     """
 
     def find_class(self, module, name):
         top_module = module.split('.')[0]
 
-        if top_module in _BLOCKED_MODULES:
+        # Reject any module not in the allowlist
+        if top_module not in _ALLOWED_PREFIXES:
+            _log.warning("Pickle restricted: blocked %s.%s (module '%s' not in allowlist)",
+                         module, name, top_module)
             raise pickle.UnpicklingError(
-                "Blocked: unpickling {}.{} is not allowed "
-                "(module '{}' is restricted)".format(module, name, top_module))
+                "Blocked: module '{}' is not in the allowlist "
+                "(attempted to load {}.{})".format(top_module, module, name))
 
+        # Block dangerous builtins even though 'builtins' is allowed
         if top_module == 'builtins' and name in _BLOCKED_BUILTINS:
+            _log.warning("Pickle restricted: blocked builtins.%s", name)
             raise pickle.UnpicklingError(
-                "Blocked: unpickling builtins.{} is not allowed".format(name))
+                "Blocked: builtins.{} is not allowed".format(name))
 
         return super().find_class(module, name)
 

tensorwatch/zmq_wrapper.py

@@ -15,6 +15,8 @@
 from .safe_pickle import restricted_loads
 import weakref, logging
 
+_log = logging.getLogger(__name__)
+
 class ZmqWrapper:
 
     _thread:Thread = None
@@ -56,11 +58,18 @@ def verify_and_loads(signed_data: bytes):
         Raises ValueError if the signature does not match.
         """
         if len(signed_data) < 32:
+            _log.warning("ZMQ security: rejected message too short for HMAC "
+                         "(%d bytes) - possible malformed or unsigned message",
+                         len(signed_data))
             raise ValueError("Message too short to contain HMAC signature")
         sig = signed_data[:32]
         payload = signed_data[32:]
         expected = hmac.new(ZmqWrapper.get_hmac_key(), payload, hashlib.sha256).digest()
         if not hmac.compare_digest(sig, expected):
+            _log.warning("ZMQ security: HMAC verification failed - "
+                         "rejecting untrusted message (%d bytes payload). "
+                         "This may indicate a connection from an unauthorized "
+                         "process or a tampering attempt.", len(payload))
             raise ValueError("HMAC verification failed - rejecting untrusted message")
         return restricted_loads(payload)