Pickle - support extra slots

Author: guedou

scapy/packet.py

@@ -84,6 +84,7 @@ def _rebuild_pkt(
     fields,  # type: Dict[str, Any]
     payload,  # type: Optional[Packet]
     metadata,  # type: Dict[str, Any]
+    extra_slots=None,  # type: Dict[str, Any]
 ):
     # type: (...) -> Packet
     """Helper for unpickling Packet instances via field values."""
@@ -98,6 +99,9 @@ def _rebuild_pkt(
     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
 
 
@@ -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.__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