Allow call args of arbitrary type

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
@@ -151,12 +152,14 @@
     FsPathTypes = (str,)
     BufferType = lambda buf, start: memoryview(buf)[start:]
     long = int
+    from types import SimpleNamespace
 else:
     b = str
     BytesType = str
     FsPathTypes = (str, unicode)
     BufferType = buffer
     UnicodeType = unicode
+    SimpleNamespace = None
 
 AnyTextType = (BytesType, UnicodeType)
 
@@ -767,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
@@ -864,7 +906,16 @@ def _find_global(self, module, func):
             return self._unpickle_bytes
         elif module == '__builtin__' and func == 'bytes':
             return BytesType
-        raise StreamError('cannot unpickle %r/%r', module, func)
+        elif SimpleNamespace and module == 'types' and func == 'SimpleNamespace':
+            return SimpleNamespace
+        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):
@@ -964,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):
@@ -3132,7 +3183,7 @@ def stream_by_id(self, dst_id):
         This can be used from any thread, but its output is only meaningful
         from the context of the :class:`Broker` thread, as disconnection or
         replacement could happen in parallel on the broker thread at any
-        moment. 
+        moment.
         """
         return (
             self._stream_by_id.get(dst_id) or
@@ -3841,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)