Pickle - support extra slots

Author: guedou

scapy/packet.py

@@ -79,28 +79,6 @@
 _T = TypeVar("_T", Dict[str, Any], Optional[Dict[str, Any]])
 
 
-def _rebuild_pkt(
-    cls,  # type: Type[Packet]
-    fields,  # type: Dict[str, Any]
-    payload,  # type: Optional[Packet]
-    metadata,  # type: Dict[str, Any]
-):
-    # type: (...) -> Packet
-    """Helper for unpickling Packet instances via field values."""
-    # Create the instance using the field values
-    pkt = cls(**fields)
-    if payload is not None:
-        pkt.add_payload(payload)
-    # Restore metadata
-    pkt.time = metadata['time']
-    pkt.sent_time = metadata['sent_time']
-    pkt.direction = metadata['direction']
-    pkt.sniffed_on = metadata['sniffed_on']
-    pkt.wirelen = metadata['wirelen']
-    pkt.comments = metadata['comments']
-    return pkt
-
-
 class Packet(
     BasePacket,
     _CanvasDumpExtended,
@@ -256,6 +234,32 @@ def comment(self, value):
         else:
             self.comments = None
 
+    @classmethod
+    def _rebuild_pkt(
+        cls,  # type: Type[Packet]
+        fields,  # type: Dict[str, Any]
+        payload,  # type: Optional[Packet]
+        metadata,  # type: Dict[str, Any]
+        extra_slots={},  # type: Dict[str, Any]
+    ):
+        # type: (...) -> Packet
+        """Helper for unpickling Packet instances via field values."""
+        # Create the instance using the field values
+        pkt = cls(**fields)
+        if payload is not None:
+            pkt.add_payload(payload)
+        # Restore metadata
+        pkt.time = metadata['time']
+        pkt.sent_time = metadata['sent_time']
+        pkt.direction = metadata['direction']
+        pkt.sniffed_on = metadata['sniffed_on']
+        pkt.wirelen = metadata['wirelen']
+        pkt.comments = metadata['comments']
+        # Restore any extra __slots__ defined by subclasses
+        for attr, value in extra_slots.items():
+            setattr(pkt, attr, value)
+        return pkt
+
     def __reduce__(self):
         # type: () -> Tuple[Any, ...]
         """Used by pickling methods.
@@ -279,9 +283,14 @@ def __reduce__(self):
             'wirelen': self.wirelen,
             'comments': self.comments,
         }
+        # Collect any extra __slots__ defined by subclasses
+        extra_slots = {}
+        for attr in type(self).__all_slots__ - set(Packet.__slots__):
+            if hasattr(self, attr):
+                extra_slots[attr] = getattr(self, attr)
         return (
-            _rebuild_pkt,
-            (self.__class__, fields, payload, metadata),
+            self._rebuild_pkt,
+            (self.__class__, fields, payload, metadata, extra_slots),
         )
 
     def __deepcopy__(self,

test/regression.uts

@@ -653,6 +653,37 @@ p2 = pickle.loads(pickle.dumps(p))
 assert p2.src == '10.0.0.1'
 assert raw(p2) == raw(p)
 
+= Pickle preserves custom __slots__ from subclasses
+
+import pickle
+import scapy.packet as _pkt_mod
+
+class _PickleTestPacket(Packet):
+    __slots__ = ["custom_id", "custom_tag"]
+    name = "PickleTestPacket"
+    fields_desc = [ByteField("val", 0)]
+
+# Make the class discoverable by pickle
+_pkt_mod._PickleTestPacket = _PickleTestPacket
+_PickleTestPacket.__module__ = 'scapy.packet'
+
+p = _PickleTestPacket(val=42)
+p.custom_id = 0x123
+p.custom_tag = "hello"
+p2 = pickle.loads(pickle.dumps(p))
+assert p2.val == 42
+assert p2.custom_id == 0x123
+assert p2.custom_tag == "hello"
+
+# Slots not explicitly set are not serialized
+p3 = _PickleTestPacket(val=7)
+assert not hasattr(p3, 'custom_id')
+p4 = pickle.loads(pickle.dumps(p3))
+assert p4.val == 7
+assert not hasattr(p4, 'custom_id')
+
+del _pkt_mod._PickleTestPacket
+
 = Usage test
 
 from scapy.main import _usage