Serialization: add user-specified pickle whitelist

End users can now control which types can be unpickled using
`Connection.pickle_whitelist_patterns`, which allows remote
function call arguments of any type to be used.

Author: jamesob

docs/howitworks.rst

@@ -505,16 +505,24 @@ The primary reason for using :py:mod:`cPickle` is that it is computationally
 efficient, and avoids including a potentially large body of serialization code
 in the bootstrap.
 
-The pickler will instantiate only built-in types and one of 3 constructor
-functions, to support unpickling :py:class:`CallError
+The pickler will, by default, instantiate only built-in types and one of 3
+constructor functions, to support unpickling :py:class:`CallError
 <mitogen.core.CallError>`, :py:class:`mitogen.core.Sender`,and
-:py:class:`Context <mitogen.core.Context>`.
+:py:class:`Context <mitogen.core.Context>`. If you want to allow the deserialization
+of arbitrary types to, for example, allow passing remote function call arguments of an
+arbitrary type, you can use :py:func:`mitogen.core.set_pickle_whitelist` to set a
+list of allowable patterns that match against a global's 
+:code:`[module].[func]` string.
 
 The choice of Pickle is one area to be revisited later. All accounts suggest it
 cannot be used securely, however few of those accounts appear to be expert, and
 none mention any additional attacks that would not be prevented by using a
 restrictive class whitelist.
 
+In the future, pickled data could include an HMAC that is based upon a
+preshared key (specified by the parent during child boot) to reduce the risk
+of malicioius tampering.
+
 
 The IO Multiplexer
 ------------------

mitogen/core.py

@@ -46,6 +46,7 @@
 import os
 import pickle as py_pickle
 import pstats
+import re
 import signal
 import socket
 import struct
@@ -769,6 +770,45 @@ def find_class(self, module, func):
     _Unpickler = pickle.Unpickler
 
 
+#: A list of compiled regex patterns which allow end-users to selectively opt
+#: into deserializing certain globals.
+_PICKLE_GLOBAL_WHITELIST_PATTERNS = None
+_PICKLE_GLOBAL_WHITELIST = None
+
+
+def set_pickle_whitelist(pattern_strings):
+    """
+    Specify regex patterns that control allowable global unpickling functions.
+
+    `pattern_strings` is sequence of pattern strings that will be fed into
+    `re.compile` and then used to authenticate pickle calls. In order for a
+    non-trivially typed message to unpickle, one of these patterns must
+    match against a complete [module].[function] string.
+    """
+    if not isinstance(pattern_strings, (tuple, list, set)):
+        pattern_strings = (pattern_strings,)
+
+    global _PICKLE_GLOBAL_WHITELIST
+    global _PICKLE_GLOBAL_WHITELIST_PATTERNS
+
+    _PICKLE_GLOBAL_WHITELIST = pattern_strings
+    _PICKLE_GLOBAL_WHITELIST_PATTERNS = []
+
+    for patt_str in pattern_strings:
+        if not patt_str.endswith('$'):
+            patt_str += '$'
+        _PICKLE_GLOBAL_WHITELIST_PATTERNS.append(re.compile(patt_str))
+
+
+def _test_pickle_whitelist_accept(module, func):
+    if not _PICKLE_GLOBAL_WHITELIST_PATTERNS:
+        return False
+
+    test_str = "{}.{}".format(module, func)
+    return bool(any(
+        patt.match(test_str) for patt in _PICKLE_GLOBAL_WHITELIST_PATTERNS))
+
+
 class Message(object):
     """
     Messages are the fundamental unit of communication, comprising fields from
@@ -868,7 +908,14 @@ def _find_global(self, module, func):
             return BytesType
         elif SimpleNamespace and module == 'types' and func == 'SimpleNamespace':
             return SimpleNamespace
-        raise StreamError('cannot unpickle %r/%r', module, func)
+        elif _test_pickle_whitelist_accept(module, func):
+            try:
+                return getattr(import_module(module), func)
+            except AttributeError as e:
+                LOG.info(str(e))
+        raise StreamError(
+            'cannot unpickle %r/%r - try using `set_pickle_whitelist`',
+            module, func)
 
     @property
     def is_dead(self):
@@ -968,8 +1015,8 @@ def unpickle(self, throw=True, throw_dead=True):
                 # Must occur off the broker thread.
                 try:
                     obj = unpickler.load()
-                except:
-                    LOG.error('raw pickle was: %r', self.data)
+                except Exception as e:
+                    LOG.error('raw pickle was: %r (exc: %r)', self.data, e)
                     raise
                 self._unpickled = obj
             except (TypeError, ValueError):
@@ -3845,6 +3892,9 @@ def _setup_master(self):
         Router.max_message_size = self.config['max_message_size']
         if self.config['profiling']:
             enable_profiling()
+        if self.config['pickle_whitelist_patterns']:
+            set_pickle_whitelist(self.config['pickle_whitelist_patterns'])
+
         self.broker = Broker(activate_compat=False)
         self.router = Router(self.broker)
         self.router.debug = self.config.get('debug', False)

mitogen/parent.py

@@ -639,7 +639,7 @@ def __init__(self):
     def get_timeout(self):
         """
         Return the floating point seconds until the next event is due.
-        
+
         :returns:
             Floating point delay, or 0.0, or :data:`None` if no events are
             scheduled.
@@ -1504,6 +1504,7 @@ def get_econtext_config(self):
             'blacklist': self._router.get_module_blacklist(),
             'max_message_size': self.options.max_message_size,
             'version': mitogen.__version__,
+            'pickle_whitelist_patterns': mitogen.core._PICKLE_GLOBAL_WHITELIST,
         }
 
     def get_preamble(self):

tests/call_function_test.py

@@ -77,9 +77,13 @@ def test_bad_return_value(self):
             lambda: self.local.call(func_with_bad_return_value))
         self.assertEqual(
                 exc.args[0],
-                "cannot unpickle '%s'/'CrazyType'" % (__name__,),
+                "cannot unpickle '%s'/'CrazyType' - try using `set_pickle_whitelist`" % (__name__,),
         )
 
+        mitogen.core.set_pickle_whitelist(r'.*CrazyType')
+        self.assertIsInstance(self.local.call(func_with_bad_return_value), CrazyType)
+        mitogen.core.set_pickle_whitelist([])
+
     def test_aborted_on_local_context_disconnect(self):
         stream = self.router._stream_by_id[self.local.context_id]
         self.broker.stop_receive(stream)