Merge pull request #1450 from moreati/release-0.3.41

Release 0.3.41

Author: moreati

docs/changelog.rst

@@ -18,6 +18,26 @@ To avail of fixes in an unreleased version, please download a ZIP file
 `directly from GitHub <https://github.com/mitogen-hq/mitogen/>`_.
 
 
+v0.3.41 (2026-02-10)
+--------------------
+
+* :gh:issue:`1441` :mod:`mitogen`: Consolidate :mod:`pickle` imports and
+  backward compatibility handling.
+* :gh:issue:`126` :mod:`mitogen`: Switch :class:`mitogen.core.Unpickler`
+  to default deny policy when handling :data:`pickle.GLOBAL` opcode.
+* :gh:issue:`1430` :mod:`mitogen`: Pickle top-level ``bytes`` objects
+  ourself on Python 3.x, to avoid ``_codecs.encode()`` call injected by
+  :class:`pickle.Pickler`
+* :gh:issue:`1430` :mod:`mitogen`: Remove caching of result of
+  :meth:`mitogen.core.Message.unpickle`
+* :gh:issue:`1430` :mod:`mitogen`: Allow mutiple pickle streams in a single
+  :class:`mitogen.core.Message`, add :meth:`mitogen.core.Message.unpickle_iter`
+* :gh:issue:`1430` :mod:`mitogen`: Speed up :class:`mitogen.core.ResourceReader`
+  using 2 pickle streams in :data:`mitogen.core.LOAD_RESOURCE` messages
+* :gh:issue:`1430` :mod:`mitogen`: Default to :func:`mitogen.core.find_deny`
+  in :meth:`mitogen.core.Message.unpickle_iter`
+
+
 v0.3.40 (2026-02-04)
 --------------------
 

mitogen/__init__.py

@@ -35,7 +35,7 @@
 
 
 #: Library version as a tuple.
-__version__ = (0, 3, 40)
+__version__ = (0, 3, 41)
 
 
 #: This is :data:`False` in slave contexts. Previously it was used to prevent

mitogen/core.py

@@ -70,7 +70,6 @@ def _path_importer_cache(cls, path):
 import itertools
 import logging
 import os
-import pickle as py_pickle
 import pstats
 import pty
 import signal
@@ -113,12 +112,57 @@ def set_blocking(fd, blocking):
     now = time.time
 
 if sys.version_info >= (3, 0):
+    from pickle import PicklingError, Unpickler as _Unpickler, UnpicklingError
+    def find_deny(module, name):
+        raise UnpicklingError('Denied: %s.%s' % (module, name))
+    class Unpickler(_Unpickler):
+        def __init__(self, file, find_class=find_deny):
+            self.find_class = find_class
+            super().__init__(file, encoding='bytes')
+else:
+    from cPickle import PicklingError, Unpickler as _Unpickler, UnpicklingError
+    def find_deny(module, name):
+        raise UnpicklingError('Denied: %s.%s' % (module, name))
+    def Unpickler(file, find_class=find_deny):
+        unpickler = _Unpickler(file)
+        unpickler.find_global = find_class
+        return unpickler
+
+if sys.version_info >= (3, 0):
+    from pickle import Pickler as _Pickler
+    class Pickler(_Pickler):
+        def __init__(self, file, protocol):
+            self._file = file
+            self._protocol = protocol
+            super().__init__(file, protocol)
+        def dump(self, obj):
+            if self._protocol == 2 and type(obj) == bytes:
+                self._file.write(struct.pack('<BBBL', 128, 2, 84, len(obj)))
+                self._file.write(obj)
+                self._file.write(struct.pack('<B', 46))
+            else:
+                super().dump(obj)
     str_partition, str_rpartition = str.partition, str.rpartition
     bytes_partition = bytes.partition
 elif sys.version_info >= (2, 5):
+    from cPickle import Pickler
     str_partition, str_rpartition = unicode.partition, unicode.rpartition
     bytes_partition = str.partition
 else:
+    import pickle
+    class Pickler(pickle.Pickler):
+        def save_exc_inst(self, obj):
+            if isinstance(obj, CallError):
+                func, args = obj.__reduce__()
+                self.save(func)
+                self.save(args)
+                self.write(pickle.REDUCE)
+            else:
+                pickle.Pickler.save_inst(self, obj)
+
+        dispatch = pickle.Pickler.dispatch.copy()
+        dispatch[pickle.InstanceType] = save_exc_inst
+
     def _part(s, sep, find):
         "(str|unicode).(partition|rpartition) polyfill for Python 2.4"
         idx = find(sep)
@@ -217,7 +261,6 @@ def any(it):
 PY24 = sys.version_info < (2, 5)
 PY3 = sys.version_info > (3,)
 if sys.version_info >= (3, 0):
-    import pickle
     import _thread as thread
     from io import BytesIO
     b = str.encode
@@ -229,7 +272,6 @@ def any(it):
     iteritems, iterkeys, itervalues = dict.items, dict.keys, dict.values
     range = range
 else:
-    import cPickle as pickle
     import thread
     from cStringIO import StringIO as BytesIO
     b = str
@@ -746,54 +788,6 @@ def iter_split(buf, delim, func):
     return buf[start:], cont
 
 
-class Py24Pickler(py_pickle.Pickler):
-    """
-    Exceptions were classic classes until Python 2.5. Sadly for 2.4, cPickle
-    offers little control over how a classic instance is pickled. Therefore 2.4
-    uses a pure-Python pickler, so CallError can be made to look as it does on
-    newer Pythons.
-
-    This mess will go away once proper serialization exists.
-    """
-    @classmethod
-    def dumps(cls, obj, protocol):
-        bio = BytesIO()
-        self = cls(bio, protocol=protocol)
-        self.dump(obj)
-        return bio.getvalue()
-
-    def save_exc_inst(self, obj):
-        if isinstance(obj, CallError):
-            func, args = obj.__reduce__()
-            self.save(func)
-            self.save(args)
-            self.write(py_pickle.REDUCE)
-        else:
-            py_pickle.Pickler.save_inst(self, obj)
-
-    if sys.version_info < (2, 5):
-        dispatch = py_pickle.Pickler.dispatch.copy()
-        dispatch[py_pickle.InstanceType] = save_exc_inst
-
-
-if sys.version_info >= (3, 0):
-    # In 3.x Unpickler is a class exposing find_class as an overridable, but it
-    # cannot be overridden without subclassing.
-    class _Unpickler(pickle.Unpickler):
-        def find_class(self, module, func):
-            return self.find_global(module, func)
-    pickle__dumps = pickle.dumps
-elif sys.version_info < (2, 5):
-    # On Python 2.4, we must use a pure-Python pickler.
-    pickle__dumps = Py24Pickler.dumps
-    _Unpickler = pickle.Unpickler
-else:
-    pickle__dumps = pickle.dumps
-    # In 2.x Unpickler is a function exposing a writeable find_global
-    # attribute.
-    _Unpickler = pickle.Unpickler
-
-
 class Message(object):
     """
     Messages are the fundamental unit of communication, comprising fields from
@@ -835,8 +829,6 @@ class Message(object):
     #: :ref:`standard-handles` should explicitly declare an encoding.
     enc = ENC_MGC
 
-    _unpickled = object()
-
     #: The :class:`Router` responsible for routing the message. This is
     #: :data:`None` for locally originated messages.
     router = None
@@ -928,21 +920,23 @@ def encoded(cls, obj, enc, **kwargs):
         raise ValueError('Invalid explicit enc: %r' % (enc,))
 
     @classmethod
-    def pickled(cls, obj, **kwargs):
+    def pickled(cls, *args, **kwargs):
         """
         Construct a pickled message, setting :attr:`data` to the serialization
-        of `obj`, and setting remaining fields using `kwargs`.
+        of each object in `args`, and setting remaining fields using `kwargs`.
 
         :returns:
             The new message.
         """
-        self = cls(enc=cls.ENC_PKL, **kwargs)
-        try:
-            self.data = pickle__dumps(obj, protocol=2)
-        except pickle.PicklingError:
-            e = sys.exc_info()[1]
-            self.data = pickle__dumps(CallError(e), protocol=2)
-        return self
+        f = BytesIO()
+        p = Pickler(f, protocol=2)
+        for obj in args:
+            try:
+                p.dump(obj)
+            except PicklingError:
+                exc = sys.exc_info()[1]
+                p.dump(CallError(exc))
+        return cls(enc=cls.ENC_PKL, data=f.getvalue(), **kwargs)
 
     def reply(self, msg, router=None, **kwargs):
         """
@@ -968,11 +962,6 @@ def reply(self, msg, router=None, **kwargs):
             LOG.debug('dropping reply to message with no return address: %r',
                       msg)
 
-    if sys.version_info >= (3, 0):
-        UNPICKLER_KWARGS = {'encoding': 'bytes'}
-    else:
-        UNPICKLER_KWARGS = {}
-
     def _throw_dead(self):
         if len(self.data):
             raise ChannelError(self.data.decode('utf-8', 'replace'))
@@ -986,49 +975,62 @@ def decode(self, throw=True, throw_dead=True):
         if self.enc == self.ENC_BIN: return self.data
         raise ValueError('Invalid explicit enc: %r' % (self.enc,))
 
-    def unpickle(self, throw=True, throw_dead=True):
+    def unpickle(self, throw=True, throw_dead=True, find_class=None):
+        """
+        Return the first unpickled stream in :attr:`data`, optionally raise
+        :exc:`CallError` if the unpickled object is such.
+
+        `throw` and `throw_dead` behave the same as with :meth:`unpickle_iter`.
+
+        :param find_class:
+            Callable that takes ``(module, func)`` and returns a constructor.
+            Defaults to :meth:`_find_global`.
         """
-        Unpickle :attr:`data`, optionally raising any exceptions present.
+        if find_class is None: find_class = self._find_global
+        return next(self.unpickle_iter(throw, throw_dead, find_class))
+
+    def unpickle_iter(self, throw=True, throw_dead=True, find_class=find_deny):
+        """
+        Return an iterator of objects unpickled from :attr:`data`, optionally
+        raising any :exc:`CallError` exceptions present.
 
         :param bool throw_dead:
             If :data:`True`, raise exceptions, otherwise it is the caller's
             responsibility.
+        :param find_class:
+            Callable that takes ``(module, func)`` and returns a constructor.
+            Default: :func:`find_deny`.
 
         :raises CallError:
             The serialized data contained CallError exception.
         :raises ChannelError:
             The `is_dead` field was set.
         """
-        _vv and IOLOG.debug('%r.unpickle()', self)
         if self.enc not in (self.ENC_MGC, self.ENC_PKL):
             raise ValueError(
                 'Message %r is not pickled, invalid enc=%r', self, self.enc,
             )
         if throw_dead and self.is_dead:
             self._throw_dead()
 
-        obj = self._unpickled
-        if obj is Message._unpickled:
-            fp = BytesIO(self.data)
-            unpickler = _Unpickler(fp, **self.UNPICKLER_KWARGS)
-            unpickler.find_global = self._find_global
+        file = BytesIO(self.data)
+        unpickler = Unpickler(file, find_class)
+        while file.tell() < len(self.data):
             try:
                 # Must occur off the broker thread.
                 try:
                     obj = unpickler.load()
                 except:
                     LOG.error('raw pickle was: %r', self.data)
                     raise
-                self._unpickled = obj
             except (TypeError, ValueError):
                 e = sys.exc_info()[1]
                 raise StreamError('invalid message: %s', e)
 
-        if throw:
-            if isinstance(obj, CallError):
+            if throw and isinstance(obj, CallError):
                 raise obj
 
-        return obj
+            yield obj
 
     def __repr__(self):
         if len(self.data) > 60:
@@ -1809,8 +1811,7 @@ def _request_resource(self, fullname, resource, callback):
     def _on_load_resource(self, msg):
         if msg.is_dead:
             return
-        tup = msg.unpickle()
-        fullname, resource, content = tup
+        (fullname, resource), content = msg.unpickle_iter()
 
         self._lock.acquire()
         try:

mitogen/master.py

@@ -1302,26 +1302,17 @@ def _on_get_resource(self, msg):
             content = importlib.resources.read_binary(fullname, resource)
         except (FileNotFoundError, IsADirectoryError):
             content = None
-        if content is not None:
-            self._send_resource(stream, fullname, resource, content)
-        else:
-            self._send_not_found(stream, fullname, resource)
 
-    def _send_resource(self, stream, fullname, resource, content):
         msg = mitogen.core.Message.pickled(
-            (fullname, resource, content),
+            (fullname, resource), content,
             dst_id=stream.protocol.remote_id,
             handle=mitogen.core.LOAD_RESOURCE,
         )
-        self._router._async_route(msg)
 
-    def _send_not_found(self, stream, fullname, resource):
-        msg = mitogen.core.Message.pickled(
-            (fullname, resource, None),
-            dst_id=stream.protocol.remote_id,
-            handle=mitogen.core.LOAD_RESOURCE,
-        )
-        stream.protocol.send(msg)
+        if content is not None:
+            self._router._async_route(msg)
+        else:
+            stream.protocol.send(msg)
 
 
 class Broker(mitogen.core.Broker):

mitogen/parent.py

@@ -2838,7 +2838,7 @@ def _send_resource(self, stream, fullname, resource):
         content = self.requester._cache[(fullname, resource)]
 
         msg = mitogen.core.Message.pickled(
-            (fullname, resource, content),
+            (fullname, resource), content,
             dst_id=stream.protocol.remote_id,
             handle=mitogen.core.LOAD_RESOURCE,
         )

tests/ansible/hosts/group_vars/all.yml

@@ -38,6 +38,7 @@ pkg_mgr_python_interpreter: python
 
 virtualenv_create_argv:
   - virtualenv
+  - --no-download
   - -p
   - "{{ virtualenv_python }}"
   - "{{ virtualenv_path }}"

tests/message_test.py

@@ -12,6 +12,7 @@
 import testlib
 
 from mitogen.core import b
+from mitogen.core import next
 
 
 class ConstructorTest(testlib.TestCase):
@@ -317,6 +318,30 @@ def test_custom_object_deserialization_fails(self):
         )
 
 
+class UnpickleIterTest(testlib.TestCase):
+    def roundtrip(self, *args, **kwargs):
+        msg1 = mitogen.core.Message.pickled(*args, **kwargs)
+        return msg1.unpickle_iter()
+
+    def test_ints(self):
+        self.assertEqual(list(self.roundtrip(1, 2, 3)), [1, 2, 3])
+
+    def test_mixed(self):
+        msg = mitogen.core.Message.pickled((u'foo.bar', u'baz.txt'), b('abc'))
+        self.assertFalse(b('_codecs') in msg.data)
+        self.assertFalse(b('encode') in msg.data)
+        self.assertFalse(b('latin1') in msg.data)
+
+        parts = msg.unpickle_iter()
+        self.assertEqual(next(parts), (u'foo.bar', u'baz.txt'))
+        self.assertEqual(next(parts), b('abc'))
+        self.assertRaises(StopIteration, next, parts)
+
+    def test_default_find_class_denies(self):
+        msg = mitogen.core.Message.pickled(1j)
+        self.assertRaises(mitogen.core.UnpicklingError, next, msg.unpickle_iter())
+
+
 class ReplyTest(testlib.TestCase):
     # getting_started.html#rpc-serialization-rules
     klass = mitogen.core.Message