diff --git a/mitogen/core.py b/mitogen/core.py index dacdb444..76baf4d3 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -722,6 +722,10 @@ class Message(object): #: the :class:`mitogen.select.Select` interface. Defaults to :data:`None`. receiver = None + HEADER_FMT = '>hLLLLLL' + HEADER_LEN = struct.calcsize(HEADER_FMT) + HEADER_MAGIC = 0x4d49 # 'MI' + def __init__(self, **kwargs): """ Construct a message from from the supplied `kwargs`. :attr:`src_id` and @@ -732,6 +736,14 @@ class Message(object): vars(self).update(kwargs) assert isinstance(self.data, BytesType) + def pack(self): + return ( + struct.pack(self.HEADER_FMT, self.HEADER_MAGIC, self.dst_id, + self.src_id, self.auth_id, self.handle, + self.reply_to or 0, len(self.data)) + + self.data + ) + def _unpickle_context(self, context_id, name): return _unpickle_context(context_id, name, router=self.router) @@ -930,7 +942,7 @@ def _unpickle_sender(router, context_id, dst_handle): if not (isinstance(router, Router) and isinstance(context_id, (int, long)) and context_id >= 0 and isinstance(dst_handle, (int, long)) and dst_handle > 0): - raise TypeError('cannot unpickle Sender: bad input') + raise TypeError('cannot unpickle Sender: bad input or missing router') return Sender(Context(router, context_id), dst_handle) diff --git a/tests/serialization_test.py b/tests/context_test.py similarity index 50% rename from tests/serialization_test.py rename to tests/context_test.py index 6cf5f8b7..4bc4bd2e 100644 --- a/tests/serialization_test.py +++ b/tests/context_test.py @@ -8,40 +8,7 @@ from mitogen.core import b import testlib -class EvilObject(object): - pass - - -def roundtrip(v): - msg = mitogen.core.Message.pickled(v) - return mitogen.core.Message(data=msg.data).unpickle() - - -class EvilObjectTest(testlib.TestCase): - def test_deserialization_fails(self): - msg = mitogen.core.Message.pickled(EvilObject()) - e = self.assertRaises(mitogen.core.StreamError, - lambda: msg.unpickle() - ) - - -class BlobTest(testlib.TestCase): - klass = mitogen.core.Blob - - # Python 3 pickle protocol 2 does weird stuff depending on whether an empty - # or nonempty bytes is being serialized. For non-empty, it yields a - # _codecs.encode() call. For empty, it yields a bytes() call. - - def test_nonempty_bytes(self): - v = mitogen.core.Blob(b('dave')) - self.assertEquals(b('dave'), roundtrip(v)) - - def test_empty_bytes(self): - v = mitogen.core.Blob(b('')) - self.assertEquals(b(''), roundtrip(v)) - - -class ContextTest(testlib.RouterMixin, testlib.TestCase): +class PickleTest(testlib.RouterMixin, testlib.TestCase): klass = mitogen.core.Context # Ensure Context can be round-tripped by regular pickle in addition to diff --git a/tests/message_test.py b/tests/message_test.py new file mode 100644 index 00000000..79deb2c6 --- /dev/null +++ b/tests/message_test.py @@ -0,0 +1,545 @@ + +import sys +import struct + +import mock +import unittest2 + +import mitogen.core +import mitogen.master +import testlib + +from mitogen.core import b + + +class ConstructorTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_dst_id_default(self): + self.assertEquals(self.klass().dst_id, None) + + def test_dst_id_explicit(self): + self.assertEquals(self.klass(dst_id=1111).dst_id, 1111) + + @mock.patch('mitogen.context_id', 1234) + def test_src_id_default(self): + self.assertEquals(self.klass().src_id, 1234) + + def test_src_id_explicit(self): + self.assertEquals(self.klass(src_id=4321).src_id, 4321) + + @mock.patch('mitogen.context_id', 5555) + def test_auth_id_default(self): + self.assertEquals(self.klass().auth_id, 5555) + + def test_auth_id_explicit(self): + self.assertEquals(self.klass(auth_id=2222).auth_id, 2222) + + def test_handle_default(self): + self.assertEquals(self.klass().handle, None) + + def test_handle_explicit(self): + self.assertEquals(self.klass(handle=1234).handle, 1234) + + def test_reply_to_default(self): + self.assertEquals(self.klass().reply_to, None) + + def test_reply_to_explicit(self): + self.assertEquals(self.klass(reply_to=8888).reply_to, 8888) + + def test_data_default(self): + m = self.klass() + self.assertEquals(m.data, b('')) + self.assertTrue(isinstance(m.data, mitogen.core.BytesType)) + + def test_data_explicit(self): + m = self.klass(data=b('asdf')) + self.assertEquals(m.data, b('asdf')) + self.assertTrue(isinstance(m.data, mitogen.core.BytesType)) + + def test_data_hates_unicode(self): + self.assertRaises(Exception, + lambda: self.klass(data=u'asdf')) + + +class PackTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_header_format_sanity(self): + self.assertEquals(self.klass.HEADER_LEN, + struct.calcsize(self.klass.HEADER_FMT)) + + def test_header_length_correct(self): + s = self.klass(dst_id=123, handle=123).pack() + self.assertEquals(len(s), self.klass.HEADER_LEN) + + def test_magic(self): + s = self.klass(dst_id=123, handle=123).pack() + magic, = struct.unpack('>h', s[:2]) + self.assertEquals(self.klass.HEADER_MAGIC, magic) + + def test_dst_id(self): + s = self.klass(dst_id=123, handle=123).pack() + dst_id, = struct.unpack('>L', s[2:6]) + self.assertEquals(123, dst_id) + + def test_src_id(self): + s = self.klass(src_id=5432, dst_id=123, handle=123).pack() + src_id, = struct.unpack('>L', s[6:10]) + self.assertEquals(5432, src_id) + + def test_auth_id(self): + s = self.klass(auth_id=1919, src_id=5432, dst_id=123, handle=123).pack() + auth_id, = struct.unpack('>L', s[10:14]) + self.assertEquals(1919, auth_id) + + def test_handle(self): + s = self.klass(dst_id=123, handle=9999).pack() + handle, = struct.unpack('>L', s[14:18]) + self.assertEquals(9999, handle) + + def test_reply_to(self): + s = self.klass(dst_id=1231, handle=7777, reply_to=9132).pack() + reply_to, = struct.unpack('>L', s[18:22]) + self.assertEquals(9132, reply_to) + + def test_data_length_empty(self): + s = self.klass(dst_id=1231, handle=7777).pack() + data_length, = struct.unpack('>L', s[22:26]) + self.assertEquals(0, data_length) + + def test_data_length_present(self): + s = self.klass(dst_id=1231, handle=7777, data=b('hello')).pack() + data_length, = struct.unpack('>L', s[22:26]) + self.assertEquals(5, data_length) + + def test_data_empty(self): + s = self.klass(dst_id=1231, handle=7777).pack() + data = s[26:] + self.assertEquals(b(''), data) + + def test_data_present(self): + s = self.klass(dst_id=11, handle=77, data=b('hello')).pack() + data = s[26:] + self.assertEquals(b('hello'), data) + + +class IsDeadTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_is_dead(self): + msg = self.klass(reply_to=mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + + def test_is_not_dead(self): + msg = self.klass(reply_to=5555) + self.assertFalse(msg.is_dead) + + +class DeadTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_no_reason(self): + msg = self.klass.dead() + self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + self.assertEquals(msg.data, b('')) + + def test_with_reason(self): + msg = self.klass.dead(reason=u'oh no') + self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + self.assertEquals(msg.data, b('oh no')) + + +class EvilObject(object): + pass + + +class PickledTest(testlib.TestCase): + # getting_started.html#rpc-serialization-rules + klass = mitogen.core.Message + + def roundtrip(self, v, router=None): + msg = self.klass.pickled(v) + msg2 = self.klass(data=msg.data) + msg2.router = router + return msg2.unpickle() + + def test_bool(self): + for b in True, False: + self.assertEquals(b, self.roundtrip(b)) + + @unittest2.skipIf(condition=sys.version_info < (2, 6), + reason='bytearray missing on <2.6') + def test_bytearray(self): + ba = bytearray(b('123')) + self.assertRaises(mitogen.core.StreamError, + lambda: self.roundtrip(ba) + ) + + def test_bytes(self): + by = b('123') + self.assertEquals(by, self.roundtrip(by)) + + def test_dict(self): + d = {1: 2, u'a': 3, b('b'): 4, 'c': {}} + roundtrip = self.roundtrip(d) + self.assertEquals(d, roundtrip) + self.assertTrue(isinstance(roundtrip, dict)) + for k in d: + self.assertTrue(isinstance(roundtrip[k], type(d[k]))) + + def test_int(self): + self.assertEquals(123, self.klass.pickled(123).unpickle()) + + def test_list(self): + l = [1, u'b', b('c')] + roundtrip = self.roundtrip(l) + self.assertTrue(isinstance(roundtrip, list)) + self.assertEquals(l, roundtrip) + for k in range(len(l)): + self.assertTrue(isinstance(roundtrip[k], type(l[k]))) + + @unittest2.skipIf(condition=sys.version_info > (3, 0), + reason='long missing in >3.x') + def test_long(self): + l = long(0xffffffffffff) + roundtrip = self.roundtrip(l) + self.assertEquals(l, roundtrip) + self.assertTrue(isinstance(roundtrip, long)) + + def test_tuple(self): + l = (1, u'b', b('c')) + roundtrip = self.roundtrip(l) + self.assertEquals(l, roundtrip) + self.assertTrue(isinstance(roundtrip, tuple)) + for k in range(len(l)): + self.assertTrue(isinstance(roundtrip[k], type(l[k]))) + + def test_unicode(self): + u = u'abcd' + roundtrip = self.roundtrip(u) + self.assertEquals(u, roundtrip) + self.assertTrue(isinstance(roundtrip, mitogen.core.UnicodeType)) + + #### custom types. see also: types_test.py, call_error_test.py + + # Python 3 pickle protocol 2 does weird stuff depending on whether an empty + # or nonempty bytes is being serialized. For non-empty, it yields a + # _codecs.encode() call. For empty, it yields a bytes() call. + + def test_blob_nonempty(self): + v = mitogen.core.Blob(b('dave')) + roundtrip = self.roundtrip(v) + self.assertTrue(isinstance(roundtrip, mitogen.core.Blob)) + self.assertEquals(b('dave'), roundtrip) + + def test_blob_empty(self): + v = mitogen.core.Blob(b('')) + roundtrip = self.roundtrip(v) + self.assertTrue(isinstance(roundtrip, mitogen.core.Blob)) + self.assertEquals(b(''), v) + + def test_secret_nonempty(self): + s = mitogen.core.Secret(u'dave') + roundtrip = self.roundtrip(s) + self.assertTrue(isinstance(roundtrip, mitogen.core.Secret)) + self.assertEquals(u'dave', roundtrip) + + def test_secret_empty(self): + s = mitogen.core.Secret(u'') + roundtrip = self.roundtrip(s) + self.assertTrue(isinstance(roundtrip, mitogen.core.Secret)) + self.assertEquals(u'', roundtrip) + + def test_call_error(self): + ce = mitogen.core.CallError('nope') + ce2 = self.assertRaises(mitogen.core.CallError, + lambda: self.roundtrip(ce)) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_context(self): + router = mitogen.master.Router() + try: + c = router.context_by_id(1234) + roundtrip = self.roundtrip(c) + self.assertTrue(isinstance(roundtrip, mitogen.core.Context)) + self.assertEquals(c.context_id, 1234) + finally: + router.broker.shutdown() + router.broker.join() + + def test_sender(self): + router = mitogen.master.Router() + try: + recv = mitogen.core.Receiver(router) + sender = recv.to_sender() + roundtrip = self.roundtrip(sender, router=router) + self.assertTrue(isinstance(roundtrip, mitogen.core.Sender)) + self.assertEquals(roundtrip.context.context_id, mitogen.context_id) + self.assertEquals(roundtrip.dst_handle, sender.dst_handle) + finally: + router.broker.shutdown() + router.broker.join() + + #### + + def test_custom_object_deserialization_fails(self): + self.assertRaises(mitogen.core.StreamError, + lambda: self.roundtrip(EvilObject()) + ) + + +class ReplyTest(testlib.TestCase): + # getting_started.html#rpc-serialization-rules + klass = mitogen.core.Message + + def test_reply_calls_router_route(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + self.assertEquals(1, router.route.call_count) + + def test_reply_pickles_object(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.dst_id, 1234) + self.assertEquals(reply.unpickle(), 123) + + def test_reply_uses_preformatted_message(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + my_reply = mitogen.core.Message.pickled(4444) + msg.reply(my_reply, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertTrue(my_reply is reply) + self.assertEquals(reply.dst_id, 1234) + self.assertEquals(reply.unpickle(), 4444) + + def test_reply_sets_dst_id(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.dst_id, 1234) + + def test_reply_sets_handle(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.handle, 9191) + + +class UnpickleTest(testlib.TestCase): + # mostly done by PickleTest, just check behaviour of parameters + klass = mitogen.core.Message + + def test_throw(self): + ce = mitogen.core.CallError('nope') + m = self.klass.pickled(ce) + ce2 = self.assertRaises(mitogen.core.CallError, + lambda: m.unpickle()) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_no_throw(self): + ce = mitogen.core.CallError('nope') + m = self.klass.pickled(ce) + ce2 = m.unpickle(throw=False) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_throw_dead(self): + m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD) + self.assertRaises(mitogen.core.ChannelError, + lambda: m.unpickle()) + + def test_no_throw_dead(self): + m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD) + self.assertEquals('derp', m.unpickle(throw_dead=False)) + + +class UnpickleCompatTest(testlib.TestCase): + # try weird variations of pickles from different Python versions. + klass = mitogen.core.Message + + def check(self, value, encoded, **kwargs): + if isinstance(encoded, mitogen.core.UnicodeType): + encoded = encoded.encode('latin1') + m = self.klass(data=encoded) + m.router = mitogen.master.Router() + try: + return m.unpickle(**kwargs) + finally: + m.router.broker.shutdown() + m.router.broker.join() + + def test_py24_bytes(self): + self.check('test', + ('\x80\x02U\x04testq\x00.')) + + def test_py24_unicode(self): + self.check(u'test', + ('\x80\x02X\x04\x00\x00\x00testq\x00.')) + + def test_py24_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py24_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py24_dict(self): + self.check({}, + ('\x80\x02}q\x00.')) + + def test_py24_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.')) + + def test_py24_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py24_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.')) + + def test_py24_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x00U\x07bigblobq\x01\x85q\x02Rq\x03.')) + + def test_py24_secret(self): + self.check(mitogen.core.Secret(u'mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.')) + + def test_py24_call_error(self): + self.check(mitogen.core.CallError('big error'), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02R.'), throw=False) + + def test_py24_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.')) + + def test_py24_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.')) + + def test_py27_bytes(self): + self.check(b('test'), + ('\x80\x02U\x04testq\x01.')) + + def test_py27_unicode(self): + self.check(u'test', + ('\x80\x02X\x04\x00\x00\x00testq\x01.')) + + def test_py27_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py27_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py27_dict(self): + self.check({}, + ('\x80\x02}q\x01.')) + + def test_py27_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00b\x87q\x01.')) + + def test_py27_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py27_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x01(K\x01K\x02X\x01\x00\x00\x00be.')) + + def test_py27_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x01U\x07bigblob\x85Rq\x02.')) + + def test_py27_secret(self): + self.check(mitogen.core.Secret(u'mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x01X\n\x00\x00\x00mypassword\x85Rq\x02.')) + + def test_py27_call_error(self): + self.check(mitogen.core.CallError(u'big error',), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x01X\t\x00\x00\x00big errorq\x02\x85Rq\x03.'), throw=False) + + def test_py27_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x01M\xd2\x04N\x86Rq\x02.')) + + def test_py27_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x01M\x03\xd9M\\\x11\x86Rq\x02.')) + + def test_py36_bytes(self): + self.check(b('test'), + ('\x80\x02c_codecs\nencode\nq\x00X\x04\x00\x00\x00testq\x01X\x06\x00\x00\x00latin1q\x02\x86q\x03Rq\x04.')) + + def test_py36_unicode(self): + self.check('test', + ('\x80\x02X\x04\x00\x00\x00testq\x00.')) + + def test_py36_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py36_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py36_dict(self): + self.check({}, + ('\x80\x02}q\x00.')) + + def test_py36_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.')) + + def test_py36_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py36_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.')) + + def test_py36_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x00c_codecs\nencode\nq\x01X\x07\x00\x00\x00bigblobq\x02X\x06\x00\x00\x00latin1q\x03\x86q\x04Rq\x05\x85q\x06Rq\x07.')) + + def test_py36_secret(self): + self.check(mitogen.core.Secret('mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.')) + + def test_py36_call_error(self): + self.check(mitogen.core.CallError('big error'), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02Rq\x03.'), throw=False) + + def test_py36_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.')) + + def test_py36_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.')) + + +class ReprTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_repr(self): + # doesn't crash + repr(self.klass.pickled('test')) + + +if __name__ == '__main__': + unittest2.main()