mitogen: Explicitly mark pickled messages, re-uses magic header field

docs/changelog.rst

@@ -23,6 +23,9 @@ In progress (unreleased)
 
 * :gh:issue:`1430` :mod:`mitogen`: Pickle :data:`mitogen.core.GET_RESOURCE`
   parameters directly as textual strings (rather than ASCII in byte strings)
+* :gh:issue:`1430` :mod:`mitogen`: Explicitly mark messages known to carry
+  pickled data, using :data:`mitogen.core.Message.ENC_PKL`. This repurposes
+  the magic field as a content encoding enumeration.
 
 
 v0.3.38 (2026-01-23)

mitogen/core.py

@@ -801,6 +801,9 @@ class Message(object):
     :class:`mitogen.core.Router` for ingress messages, and helper methods for
     deserialization and generating replies.
     """
+    ENCS = frozenset(range(0x4d49, 0x4d49+2))
+    ENC_MGC, ENC_PKL = sorted(ENCS)
+
     #: Integer target context ID. :class:`Router` delivers messages locally
     #: when their :attr:`dst_id` matches :data:`mitogen.context_id`, otherwise
     #: they are routed up or downstream.
@@ -827,6 +830,11 @@ class Message(object):
     #: Raw message data bytes.
     data = b('')
 
+    #: Encoding of payload in :attr:`data`, one of the ``ENC_*`` constants.
+    #: :attr:`ENC_MGC` is an implicit, legacy value. New features &
+    #: :ref:`standard-handles` should explicitly declare an encoding.
+    enc = ENC_MGC
+
     _unpickled = object()
 
     #: The :class:`Router` responsible for routing the message. This is
@@ -839,7 +847,7 @@ class Message(object):
 
     HEADER_FMT = '>hLLLLLL'
     HEADER_LEN = struct.calcsize(HEADER_FMT)
-    HEADER_MAGIC = 0x4d49  # 'MI'
+    HEADER_MAGIC = ENC_MGC
 
     def __init__(self, **kwargs):
         """
@@ -850,10 +858,12 @@ class Message(object):
         self.auth_id = mitogen.context_id
         vars(self).update(kwargs)
         assert isinstance(self.data, BytesType), 'Message data is not Bytes'
+        if self.enc not in self.ENCS:
+            raise ValueError('Invalid enc: %r' % (self.enc,))
 
     def pack(self):
         return (
-            struct.pack(self.HEADER_FMT, self.HEADER_MAGIC, self.dst_id,
+            struct.pack(self.HEADER_FMT, self.enc, self.dst_id,
                         self.src_id, self.auth_id, self.handle,
                         self.reply_to or 0, len(self.data))
             + self.data
@@ -920,7 +930,7 @@ class Message(object):
         :returns:
             The new message.
         """
-        self = cls(**kwargs)
+        self = cls(enc=cls.ENC_PKL, **kwargs)
         try:
             self.data = pickle__dumps(obj, protocol=2)
         except pickle.PicklingError:
@@ -2386,13 +2396,13 @@ class MitogenProtocol(Protocol):
 
         msg = Message()
         msg.router = self._router
-        (magic, msg.dst_id, msg.src_id, msg.auth_id,
+        (msg.enc, msg.dst_id, msg.src_id, msg.auth_id,
          msg.handle, msg.reply_to, msg_len) = struct.unpack(
             Message.HEADER_FMT,
             self._input_buf[0][:Message.HEADER_LEN],
         )
 
-        if magic != Message.HEADER_MAGIC:
+        if msg.enc not in Message.ENCS:
             LOG.error(self.corrupt_msg, self.stream.name, self._input_buf[0][:2048])
             self.stream.on_disconnect(broker)
             return False

tests/message_test.py

@@ -59,9 +59,14 @@ class ConstructorTest(testlib.TestCase):
         self.assertEqual(m.data, b('asdf'))
         self.assertIsInstance(m.data, mitogen.core.BytesType)
 
-    def test_data_hates_unicode(self):
+    def test_enc(self):
+        self.assertEqual(self.klass().enc, self.klass.ENC_MGC)
+        self.assertEqual(self.klass(enc=self.klass.ENC_PKL).enc, self.klass.ENC_PKL)
+
+    def test_invalid_args(self):
         self.assertRaises(Exception,
             lambda: self.klass(data=u'asdf'))
+        self.assertRaises(ValueError, lambda: self.klass(enc=42))
 
 
 class PackTest(testlib.TestCase):
@@ -75,10 +80,10 @@ class PackTest(testlib.TestCase):
         s = self.klass(dst_id=123, handle=123).pack()
         self.assertEqual(len(s), self.klass.HEADER_LEN)
 
-    def test_magic(self):
+    def test_enc(self):
         s = self.klass(dst_id=123, handle=123).pack()
-        magic, = struct.unpack('>h', s[:2])
-        self.assertEqual(self.klass.HEADER_MAGIC, magic)
+        enc, = struct.unpack('>h', s[:2])
+        self.assertEqual(self.klass.ENC_MGC, enc)
 
     def test_dst_id(self):
         s = self.klass(dst_id=123, handle=123).pack()
@@ -164,7 +169,9 @@ class PickledTest(testlib.TestCase):
 
     def roundtrip(self, v, router=None):
         msg = self.klass.pickled(v)
+        self.assertEqual(self.klass.ENC_PKL, msg.enc)
         msg2 = self.klass(data=msg.data)
+        self.assertEqual(self.klass.ENC_MGC, msg2.enc)
         msg2.router = router
         return msg2.unpickle()