fix(ipv6)!: an extension header parsed with extension=True cannot rebuild from its own info (#1446)

pcapkit/protocols/internet/internet.py

@@ -12,6 +12,7 @@
 
 """
 import collections
+import io
 from typing import TYPE_CHECKING, Generic, cast
 
 from pcapkit.const.reg.transtype import TransType as Enum_TransType
@@ -23,7 +24,7 @@
 from pcapkit.utilities.warnings import RegistryWarning, warn
 
 if TYPE_CHECKING:
-    from typing import Any, Optional, Type
+    from typing import IO, Any, Optional, Type
 
     from typing_extensions import Literal
 
@@ -173,6 +174,36 @@ def register(cls, code: 'Enum_TransType', protocol: 'ModuleDescriptor[ProtocolBa
                  f'{incumbent!r} with {protocol!r}', RegistryWarning)
         cls.__proto__[code] = protocol
 
+    ##########################################################################
+    # Data models.
+    ##########################################################################
+
+    def __post_init__(self, file: 'Optional[IO[bytes] | bytes]' = None,
+                      length: 'Optional[int]' = None, *,
+                      extension: 'bool' = False, **kwargs: 'Any') -> 'None':
+        """Post initialisation hook.
+
+        Args:
+            file: Source packet stream.
+            length: Length of packet data.
+            extension: If the protocol is used as an IPv6 extension header.
+            **kwargs: Arbitrary keyword arguments.
+
+        Notes:
+            An IPv6 extension header is handed the rest of the datagram, but
+            parses only itself: its :attr:`info` holds no payload, and
+            :class:`~pcapkit.protocols.internet.ipv6.IPv6` decodes what follows.
+            So :attr:`data` is cut down to the header's own :attr:`length`, and
+            ``from_data(info).data`` rebuilds it exactly (:issue:`1446`).
+
+        """
+        super().__post_init__(file, length, extension=extension, **kwargs)  # type: ignore[arg-type]
+
+        if extension and file is not None and len(self._data) > self.length:
+            self._data = self._data[:self.length]
+            self._file = io.BytesIO(self._data)
+            self.__cached__.pop('__len__', None)
+
     ##########################################################################
     # Utilities.
     ##########################################################################

tests/protocols/internet/test_ipv6_ext_header_rebuild_unit.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+"""An IPv6 extension header rebuilds byte for byte from its own info.
+
+GitHub issue #1446: :class:`~pcapkit.protocols.internet.ipv6.IPv6` hands each
+extension header the rest of the datagram, with ``extension=True``. The header
+parses only itself -- its info holds no payload -- but its ``data`` kept every
+octet after it, so ``type(ex).from_data(ex.info).data`` came back short: 8 of
+the 28 octets of an 8-octet header followed by a UDP datagram.
+
+``data`` now holds the header alone, which is also what the parent's
+``from_data`` concatenates, so the parent rebuild is pinned alongside.
+
+Every case builds its own octets in memory and reads no capture. Classes are
+imported inside each test, after :func:`~tests._support.reimport_once_per_class`.
+
+"""
+
+import importlib
+import unittest
+
+from tests._support import reimport_once_per_class
+
+#: Module, class and IPv6 Next Header code of each extension header.
+EXTENSIONS = (
+    ('hopopt', 'HOPOPT', 0),
+    ('ipv6_route', 'IPv6_Route', 43),
+    ('ipv6_frag', 'IPv6_Frag', 44),
+    ('ipv6_opts', 'IPv6_Opts', 60),
+    ('mh', 'MH', 135),
+    ('ah', 'AH', 51),
+    ('hip', 'HIP', 139),
+)
+
+
+class TestIPv6ExtHeaderRebuild(unittest.TestCase):
+    """Pin ``from_data(info).data == data`` for extension-mode headers."""
+
+    def setUp(self) -> None:
+        reimport_once_per_class(self)
+
+    def _parse(self, mod: str, cls: str, code: int, upper: bytes) -> 'tuple[bytes, object, object]':
+        from pcapkit.protocols.internet.ipv6 import IPv6
+
+        klass = getattr(importlib.import_module(f'pcapkit.protocols.internet.{mod}'), cls)
+        header = klass(next=17, payload=b'').data
+        ipv6 = IPv6(IPv6(next=code, src='2001:db8::1', dst='2001:db8::2',
+                         payload=header + upper).data)
+        ext = next(iter(ipv6.extension_headers.values()))
+        self.assertIs(type(ext), klass)
+        return header, ext, ipv6
+
+    def test_extension_header_rebuilds_from_info(self) -> None:
+        from pcapkit.protocols.transport.udp import UDP
+
+        upper = UDP(srcport=1, dstport=2, payload=b'hello world!').data
+        for mod, cls, code in EXTENSIONS:
+            with self.subTest(header=cls):
+                header, ext, _ = self._parse(mod, cls, code, upper)
+                self.assertEqual(ext.data, header)
+                self.assertEqual(len(ext), ext.length)
+                self.assertEqual(type(ext).from_data(ext.info).data, ext.data)
+
+    def test_parent_ipv6_still_rebuilds(self) -> None:
+        from pcapkit.protocols.internet.ipv6 import IPv6
+        from pcapkit.protocols.transport.udp import UDP
+
+        upper = UDP(srcport=1, dstport=2, payload=b'hello world!').data
+        for mod, cls, code in EXTENSIONS:
+            with self.subTest(header=cls):
+                _, _, ipv6 = self._parse(mod, cls, code, upper)
+                self.assertEqual(IPv6.from_data(ipv6.info).data, ipv6.data)
+
+    def test_fragment_keeps_header_only(self) -> None:
+        # a non-first fragment, as in ``examples/captures/ipv6.pcap``: the
+        # remainder is fragment data, not a decodable upper layer
+        from pcapkit.protocols.internet.ipv6 import IPv6
+        from pcapkit.protocols.internet.ipv6_frag import IPv6_Frag
+
+        header = IPv6_Frag(next=17, offset=181, id=0x1234).data
+        ipv6 = IPv6(IPv6(next=44, src='2001:db8::1', dst='2001:db8::2',
+                         payload=header + bytes(range(256)) * 2).data)
+        ext = ipv6.extension_headers.get(44)
+        self.assertEqual(ext.data, header)
+        self.assertEqual(IPv6_Frag.from_data(ext.info).data, header)
+        self.assertEqual(IPv6.from_data(ipv6.info).data, ipv6.data)
+
+    def test_standalone_parse_keeps_payload(self) -> None:
+        # without ``extension=True`` the header still owns what follows it
+        from pcapkit.protocols.internet.hopopt import HOPOPT
+        from pcapkit.protocols.transport.udp import UDP
+
+        upper = UDP(srcport=1, dstport=2, payload=b'hello world!').data
+        packet = HOPOPT(next=17, payload=upper).data
+        hop = HOPOPT(packet)
+        self.assertEqual(hop.data, packet)
+        self.assertEqual(HOPOPT.from_data(hop.info).data, packet)
+
+
+if __name__ == '__main__':
+    unittest.main()