diff --git a/hippolyzer/lib/base/message/message.py b/hippolyzer/lib/base/message/message.py index 857204a..7e63dd3 100644 --- a/hippolyzer/lib/base/message/message.py +++ b/hippolyzer/lib/base/message/message.py @@ -193,12 +193,21 @@ class Message: # should be set once a packet is sent / dropped to prevent accidental # re-sending or re-dropping self.finalized = False - # Whether message is owned by the queue or should be sent immediately + # Whether message is owned by a queue or should be sent immediately self.queued: bool = False self._blocks: BLOCK_DICT = {} self.add_blocks(args) + def __reduce_ex__(self, protocol): + reduced: Tuple[Any] = super().__reduce_ex__(protocol) + # https://docs.python.org/3/library/pickle.html#object.__reduce__ + # We need to make some changes to the object state to make it serializable + state_dict: Dict = reduced[2][1] + # Have to remove the deserializer weak ref so we can pickle + state_dict['deserializer'] = None + return reduced + @property def packet_id(self) -> Optional[int]: return self._packet_id diff --git a/hippolyzer/lib/proxy/message_logger.py b/hippolyzer/lib/proxy/message_logger.py index 45266f1..fc11377 100644 --- a/hippolyzer/lib/proxy/message_logger.py +++ b/hippolyzer/lib/proxy/message_logger.py @@ -546,7 +546,6 @@ class LLUDPMessageLogEntry(AbstractMessageLogEntry): # These are expensive to keep around. pickle them and un-pickle on # an as-needed basis. self._deserializer = self.message.deserializer - self.message.deserializer = None self._frozen_message = pickle.dumps(self._message, protocol=pickle.HIGHEST_PROTOCOL) self._message = None diff --git a/tests/base/test_message_wrapper.py b/tests/base/test_message_wrapper.py index c421077..7297cda 100644 --- a/tests/base/test_message_wrapper.py +++ b/tests/base/test_message_wrapper.py @@ -126,8 +126,6 @@ class TestMessage(unittest.TestCase): def test_partial_decode_pickle(self): msg = self.deserial.deserialize(self.serial.serialize(self.chat_msg)) self.assertEqual(msg.deserializer(), self.deserial) - # Have to remove the weak ref so we can pickle - msg.deserializer = None msg = pickle.loads(pickle.dumps(msg, protocol=pickle.HIGHEST_PROTOCOL)) # We should still have the raw body at this point