From ed90f3fa90ef19e71aa90e6b7df8a4a5f48c9fb2 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 16 Sep 2017 18:08:42 +0530
Subject: [PATCH] pickle: Prevent access to the _Dead and CallError
 constructors

This should be pretty much identical the same behaviour as before, but
the extra assertion makes me feel happier.
---
 mitogen/core.py   | 20 ++++++++++++++++----
 mitogen/master.py |  4 ++--
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index 9e05fb80..dba5648b 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -60,16 +60,22 @@ class CallError(Error):
     exception message.
     """
     def __init__(self, e):
-        s = ''
-        if not isinstance(e, basestring):
-            s += '%s.%s: ' % (type(e).__module__, type(e).__name__)
-        s += str(e)
+        s = '%s.%s: %s' % (type(e).__module__, type(e).__name__, e)
         tb = sys.exc_info()[2]
         if tb:
             s += '\n'
             s += ''.join(traceback.format_tb(tb))
         Error.__init__(self, s)
 
+    def __reduce__(self):
+        return (_unpickle_call_error, (self[0],))
+
+def _unpickle_call_error(s):
+    assert type(s) is str and len(s) < 10000
+    inst = CallError.__new__(CallError)
+    Exception.__init__(inst, s)
+    return inst
+
 
 class ChannelError(Error):
     """Raised when a channel dies or has been closed."""
@@ -87,9 +93,15 @@ class Dead(object):
     def __eq__(self, other):
         return type(other) is Dead
 
+    def __reduce__(self):
+        return (_unpickle_dead, ())
+
     def __repr__(self):
         return '<Dead>'
 
+def _unpickle_dead():
+    return _DEAD
+
 
 #: Sentinel value used to represent :py:class:`Channel` disconnection.
 _DEAD = Dead()
diff --git a/mitogen/master.py b/mitogen/master.py
index c694dfc7..5d8e18d9 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -39,8 +39,8 @@ COMMENT_RE = re.compile(r'^[ ]*#[^\n]*$', re.M)
 IOLOG_RE = re.compile(r'^[ ]*IOLOG.debug\(.+?\)$', re.M)
 
 PERMITTED_CLASSES = set([
-    ('mitogen.core', 'CallError'),
-    ('mitogen.core', 'Dead'),
+    ('mitogen.core', '_unpickle_call_error'),
+    ('mitogen.core', '_unpickle_dead'),
 ])