From cecef992b09573c90af5c04a8292846a47ad92ad Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Fri, 4 May 2018 21:33:19 +0100
Subject: [PATCH] issue #218: core: add Secret and Blob types.

---
 docs/getting_started.rst |  2 ++
 mitogen/core.py          | 36 +++++++++++++++++++-
 mitogen/utils.py         | 14 +++++---
 tests/types_test.py      | 72 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 119 insertions(+), 5 deletions(-)
 create mode 100644 tests/types_test.py

diff --git a/docs/getting_started.rst b/docs/getting_started.rst
index aec86bd7..ddaccf87 100644
--- a/docs/getting_started.rst
+++ b/docs/getting_started.rst
@@ -320,6 +320,8 @@ remote procedure calls:
 
 User-defined types may not be used, except for:
 
+* :py:class:`mitogen.core.Blob`
+* :py:class:`mitogen.core.Secret`
 * :py:class:`mitogen.core.CallError`
 * :py:class:`mitogen.core.Context`
 * :py:class:`mitogen.core.Sender`
diff --git a/mitogen/core.py b/mitogen/core.py
index f47b8380..a3d13353 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -81,8 +81,12 @@ IS_DEAD = 999
 PY3 = sys.version_info > (3,)
 if PY3:
     b = lambda s: s.encode('latin-1')
+    BytesType = bytes
+    UnicodeType = unicode
 else:
     b = str
+    BytesType = str
+    UnicodeType = unicode
 
 CHUNK_SIZE = 131072
 _tls = threading.local()
@@ -109,6 +113,25 @@ class LatchError(Error):
     pass
 
 
+class Blob(BytesType):
+    def __repr__(self):
+        return '[blob: %d bytes]' % len(self)
+
+    def __reduce__(self):
+        return (_unpickle_blob, (BytesType(self),))
+
+
+class Secret(UnicodeType):
+    def __repr__(self):
+        return '[secret]'
+
+    def __str__(self):
+        return UnicodeType(self)
+
+    def __reduce__(self):
+        return (_unpickle_secret, (UnicodeType(self),))
+
+
 class CallError(Error):
     def __init__(self, fmt=None, *args):
         if not isinstance(fmt, Exception):
@@ -127,6 +150,14 @@ class CallError(Error):
         return (_unpickle_call_error, (self[0],))
 
 
+def _unpickle_blob(s):
+    return Blob(s)
+
+
+def _unpickle_secret(s):
+    return Secret(s)
+
+
 def _unpickle_call_error(s):
     if not (type(s) is str and len(s) < 10000):
         raise TypeError('cannot unpickle CallError: bad input')
@@ -310,7 +341,10 @@ class Message(object):
                 return self._unpickle_sender
             elif func == '_unpickle_context':
                 return self._unpickle_context
-
+            elif func == '_unpickle_blob':
+                return _unpickle_blob
+            elif func == '_unpickle_secret':
+                return _unpickle_secret
         raise StreamError('cannot unpickle %r/%r', module, func)
 
     @property
diff --git a/mitogen/utils.py b/mitogen/utils.py
index 480940ab..df876ea0 100644
--- a/mitogen/utils.py
+++ b/mitogen/utils.py
@@ -101,19 +101,25 @@ def with_router(func):
     return wrapper
 
 
+PASSTHROUGH = (
+    int, float, bool,
+    type(None),
+    mitogen.core.Context,
+    mitogen.core.CallError,
+    mitogen.core.Blob,
+    mitogen.core.Secret,
+)
+
 def cast(obj):
     if isinstance(obj, dict):
         return {cast(k): cast(v) for k, v in obj.iteritems()}
     if isinstance(obj, (list, tuple)):
         return [cast(v) for v in obj]
-    if obj is None or isinstance(obj, (int, float)):
+    if isinstance(obj, PASSTHROUGH):
         return obj
     if isinstance(obj, unicode):
         return unicode(obj)
     if isinstance(obj, str):
         return str(obj)
-    if isinstance(obj, (mitogen.core.Context,
-                        mitogen.core.CallError)):
-        return obj
 
     raise TypeError("Cannot serialize: %r: %r" % (type(obj), obj))
diff --git a/tests/types_test.py b/tests/types_test.py
new file mode 100644
index 00000000..44296656
--- /dev/null
+++ b/tests/types_test.py
@@ -0,0 +1,72 @@
+
+import cStringIO
+
+import unittest2
+
+import mitogen.core
+
+
+class BlobTest(unittest2.TestCase):
+    klass = mitogen.core.Blob
+
+    def make(self):
+        return self.klass('x' * 128)
+
+    def test_repr(self):
+        blob = self.make()
+        self.assertEquals('[blob: 128 bytes]', repr(blob))
+
+    def test_decays_on_constructor(self):
+        blob = self.make()
+        self.assertEquals('x'*128, mitogen.core.BytesType(blob))
+
+    def test_decays_on_write(self):
+        blob = self.make()
+        io = cStringIO.StringIO()
+        io.write(blob)
+        self.assertEquals(128, io.tell())
+        self.assertEquals('x'*128, io.getvalue())
+
+    def test_message_roundtrip(self):
+        blob = self.make()
+        msg = mitogen.core.Message.pickled(blob)
+        blob2 = msg.unpickle()
+        self.assertEquals(type(blob), type(blob2))
+        self.assertEquals(repr(blob), repr(blob2))
+        self.assertEquals(mitogen.core.BytesType(blob),
+                          mitogen.core.BytesType(blob2))
+
+
+class SecretTest(unittest2.TestCase):
+    klass = mitogen.core.Secret
+
+    def make(self):
+        return self.klass('password')
+
+    def test_repr(self):
+        secret = self.make()
+        self.assertEquals('[secret]', repr(secret))
+
+    def test_decays_on_constructor(self):
+        secret = self.make()
+        self.assertEquals('password', mitogen.core.UnicodeType(secret))
+
+    def test_decays_on_write(self):
+        secret = self.make()
+        io = cStringIO.StringIO()
+        io.write(secret)
+        self.assertEquals(8, io.tell())
+        self.assertEquals('password', io.getvalue())
+
+    def test_message_roundtrip(self):
+        secret = self.make()
+        msg = mitogen.core.Message.pickled(secret)
+        secret2 = msg.unpickle()
+        self.assertEquals(type(secret), type(secret2))
+        self.assertEquals(repr(secret), repr(secret2))
+        self.assertEquals(mitogen.core.BytesType(secret),
+                          mitogen.core.BytesType(secret2))
+
+
+if __name__ == '__main__':
+    unittest2.main()