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 specify :py:attr:`Connection.pickle_whitelist_patterns` as 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
@@ -755,6 +756,30 @@ def find_class(self, module, func):
     _Unpickler = pickle.Unpickler
 
 
+#: A compiled regexp which allows end-users to selectively opt into deserializing
+#: certain globals.
+_PICKLE_GLOBAL_WHITELIST_PATTERNS = None
+
+
+def _set_pickle_whitelist(pattern_strings):
+    global _PICKLE_GLOBAL_WHITELIST_PATTERNS
+    _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
@@ -854,7 +879,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 `pickle_whitelist_patterns`',
+            module, func)
 
     @property
     def is_dead(self):
@@ -3816,6 +3848,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.
@@ -1306,9 +1306,15 @@ class Options(object):
     #: UNIX timestamp after which the connection attempt should be abandoned.
     connect_deadline = None
 
+    #: A sequence of pattern strings that will be fed into `re.compile` and then used
+    #: to authenticate pickle calls. One of these patterns must match against a
+    #: complete [module].[function] string in order to allow unpickling .
+    pickle_whitelist_patterns = None
+
     def __init__(self, max_message_size, name=None, remote_name=None,
                  python_path=None, debug=False, connect_timeout=None,
-                 profiling=False, unidirectional=False, old_router=None):
+                 profiling=False, unidirectional=False, old_router=None,
+                 pickle_whitelist_patterns=None):
         self.name = name
         self.max_message_size = max_message_size
         if python_path:
@@ -1326,6 +1332,7 @@ def __init__(self, max_message_size, name=None, remote_name=None,
         self.unidirectional = unidirectional
         self.max_message_size = max_message_size
         self.connect_deadline = mitogen.core.now() + self.connect_timeout
+        self.pickle_whitelist_patterns = pickle_whitelist_patterns
 
 
 class Connection(object):
@@ -1504,6 +1511,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': self.options.pickle_whitelist_patterns,
         }
 
     def get_preamble(self):

tests/call_function_test.py

@@ -81,7 +81,7 @@ def test_bad_return_value(self):
             lambda: self.local.call(func_with_bad_return_value))
         self.assertEquals(
                 exc.args[0],
-                "cannot unpickle '%s'/'CrazyType'" % (__name__,),
+                "cannot unpickle '%s'/'CrazyType' - try using `pickle_whitelist_patterns`" % (__name__,),
         )
 
     def test_aborted_on_local_context_disconnect(self):