fix(ipv6): keep the extension-header chain in from_data rebuilds (#1177)

pcapkit/protocols/data/internet/ipv6.py

@@ -8,12 +8,14 @@
 
 if TYPE_CHECKING:
     from ipaddress import IPv6Address
-    from typing import Any
+    from typing import Any, Type
 
     from typing_extensions import Literal
 
     from pcapkit.const.reg.transtype import TransType
+    from pcapkit.protocols.data.data import Data
     from pcapkit.protocols.data.protocol import Packet
+    from pcapkit.protocols.protocol import ProtocolBase
 
 __all__ = [
     'IPv6',
@@ -30,6 +32,8 @@ class IPv6(Protocol):
 
     """
 
+    __excluded__ = ['__exthdr__']
+
     #: Version.
     version: 'Literal[6]'
     #: Traffic class.
@@ -56,6 +60,10 @@ class IPv6(Protocol):
         hdr_len: 'int'
         #: Raw payload length (excluding extension headers).
         raw_len: 'int'
+        #: Extension header chain, in wire order, as pairs of the parser
+        #: class and its info. Kept apart from the per-header keys above,
+        #: which a repeated header overwrites.
+        __exthdr__: 'tuple[tuple[Type[ProtocolBase], Data], ...]'
 
     def __new__(cls, *args: 'Any', **kwargs: 'Any') -> 'IPv6':
         self = super().__new__(cls, *args, **kwargs)

pcapkit/protocols/internet/ipv6.py

@@ -321,6 +321,15 @@ def _make_data(cls, data: 'Data_IPv6') -> 'dict[str, Any]':  # type: ignore[over
             Key-value pairs for protocol construction.
 
         """
+        upper = cls._make_payload(data)
+        payload = upper  # type: bytes | ProtocolBase
+
+        # NOTE: The extension headers sit between this header and the upper
+        # layer payload, so they are rebuilt in wire order ahead of it.
+        exthdr = data.get('__exthdr__')
+        if exthdr:
+            payload = b''.join(proto.from_data(info).data for proto, info in exthdr) + upper.data
+
         return {
             'traffic_class': data['class'],
             'flow_label': data.label,
@@ -333,7 +342,7 @@ def _make_data(cls, data: 'Data_IPv6') -> 'dict[str, Any]':  # type: ignore[over
             # capture rebuilds with the length that was on the wire
             # (:issue:`1155`).
             'payload_length': data.payload,
-            'payload': cls._make_payload(data)
+            'payload': payload,
         }
 
     def _read_ip_hextet(self) -> 'tuple[int, int, int]':
@@ -380,6 +389,7 @@ def _decode_next_layer(self, ipv6: 'Data_IPv6', proto: 'Optional[int]' = None,
         hdr_len = self.length       # header length
         raw_len = ipv6.payload      # payload length
         _protos = []                # ProtoChain buffer
+        _exthdr = []                # (parser class, info) per extension header
 
         # traverse if next header is an extension header
         payload = self.__header__.get_payload()
@@ -406,6 +416,7 @@ def _decode_next_layer(self, ipv6: 'Data_IPv6', proto: 'Optional[int]' = None,
             # record protocol name
             # self._protos = ProtoChain(name, chain, alias)
             _protos.append(next_)
+            _exthdr.append((type(next_), info))
 
             # update header & payload length
             hdr_len += next_.length  # type: ignore[assignment]
@@ -464,6 +475,11 @@ def _decode_next_layer(self, ipv6: 'Data_IPv6', proto: 'Optional[int]' = None,
 
             # update next header
             'protocol': proto,
+
+            # extension header chain in wire order, for ``from_data``; the
+            # per-header keys above cannot carry it, since a repeated header
+            # overwrites its predecessor's key
+            '__exthdr__': tuple(_exthdr),
         })
 
         ipv6_exthdr = ProtoChain.from_list(_protos)  # type: ignore[arg-type]

tests/protocols/internet/test_ipv6_from_data_ext_unit.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+"""IPv6 ``from_data`` keeps the extension-header chain, byte for byte.
+
+GitHub issue #1177: :meth:`IPv6._make_data
+<pcapkit.protocols.internet.ipv6.IPv6._make_data>` rebuilt the payload from
+the upper layer alone, so every extension header between the IPv6 header and
+that layer was dropped from the rebuild. An IPv6 header followed by an 8-octet
+Routing header came back as 40 octets instead of 48.
+
+Every case builds its own octets in memory and reads no capture.
+
+:class:`IPv6` is imported inside each test, after
+:func:`~tests._support.reimport_once_per_class`, not at module load, so that
+it belongs to the same :mod:`pcapkit` import as the extension-header classes
+it dispatches to.
+
+"""
+
+import struct
+import unittest
+
+from tests._support import reimport_once_per_class
+
+#: Next header codes.
+HOPOPT, UDP, ROUTE, NONXT, OPTS, MH = 0, 17, 43, 59, 60, 135
+
+#: PadN option filling an 8-octet options header.
+PADN = bytes.fromhex('0104' '00000000')
+
+
+def ext(next_: int, body: bytes) -> bytes:
+    """Return a generic extension header with ``body`` after ``Hdr Ext Len``."""
+    return struct.pack('!BB', next_, (len(body) + 2) // 8 - 1) + body
+
+
+def ipv6(next_: int, payload: bytes) -> bytes:
+    """Return an IPv6 header carrying ``payload``."""
+    return (bytes.fromhex('60000000') + struct.pack('!HBB', len(payload), next_, 64)
+            + bytes(15) + b'\x01' + bytes(15) + b'\x02' + payload)
+
+
+#: UDP datagram with an 8-octet payload and no checksum.
+UDP_DGRAM = struct.pack('!HHHH', 1234, 5678, 16, 0) + b'payload!'
+
+PACKETS = {
+    'hopopt': ipv6(HOPOPT, ext(NONXT, PADN)),
+    'routing': ipv6(ROUTE, bytes.fromhex('3b00000000000000')),
+    'routing-trailing-raw': ipv6(ROUTE, bytes.fromhex('3b00000000000000') + b'abcd'),
+    'destination-options': ipv6(OPTS, ext(NONXT, PADN)),
+    'mobility-header': ipv6(MH, bytes.fromhex('3b00000000000000')),
+    'chain-with-repeated-header': ipv6(
+        HOPOPT,
+        ext(OPTS, PADN) + ext(ROUTE, PADN) + bytes.fromhex('3c00000000000000')
+        + ext(UDP, PADN) + UDP_DGRAM,
+    ),
+}
+
+
+class TestIPv6FromDataExtensionHeaders(unittest.TestCase):
+    """Pin that a rebuild re-emits every extension header in wire order."""
+
+    def setUp(self) -> None:
+        reimport_once_per_class(self)
+
+    def test_from_data_rebuilds_byte_for_byte(self) -> None:
+        from pcapkit.protocols.internet.ipv6 import IPv6
+
+        for name, packet in PACKETS.items():
+            with self.subTest(packet=name):
+                parsed = IPv6(packet, len(packet))
+                self.assertTrue(parsed.extension_headers, 'no extension header parsed')
+                self.assertEqual(IPv6.from_data(parsed.info).data, packet)
+
+    def test_repeated_header_keeps_both_instances(self) -> None:
+        from pcapkit.protocols.internet.ipv6 import IPv6
+
+        packet = PACKETS['chain-with-repeated-header']
+        parsed = IPv6(packet, len(packet))
+        self.assertEqual(
+            [type(proto).__name__ for _, proto in parsed.extension_headers.items(multi=True)],
+            ['HOPOPT', 'IPv6_Opts', 'IPv6_Route', 'IPv6_Opts'],
+        )
+        self.assertEqual(IPv6.from_data(parsed.info).data, packet)
+
+    def test_chain_is_hidden_from_info_views(self) -> None:
+        from pcapkit.protocols.internet.ipv6 import IPv6
+
+        packet = PACKETS['routing']
+        info = IPv6(packet, len(packet)).info
+        self.assertNotIn('__exthdr__', list(info))
+        self.assertNotIn('__exthdr__', info.to_dict())
+
+
+if __name__ == '__main__':
+    unittest.main()