fix(udp)!: keep the octets past the UDP Length field as captured (#1461)

pcapkit/protocols/data/transport/udp.py

@@ -26,4 +26,8 @@ class UDP(Protocol):
     checksum: 'bytes'
 
     if TYPE_CHECKING:
+        #: Octets captured past the Length field, present only when there are
+        #: any.
+        trailer: 'bytes'
+
         def __init__(self, srcport: 'AppType', dstport: 'AppType', len: 'int', checksum: 'bytes') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements,redefined-builtin

pcapkit/protocols/schema/transport/udp.py

@@ -104,9 +104,14 @@ class UDP(Schema):
     len: 'int' = UInt16Field()
     #: Checksum of UDP packet.
     checksum: 'bytes' = BytesField(length=2)
-    #: Payload.
-    payload: 'bytes' = PayloadField()
+    #: Payload. A Length below the header's own 8 octets leaves the length
+    #: negative, which takes the rest of the datagram, as before.
+    payload: 'bytes' = PayloadField(length=lambda pkt: pkt['len'] - 8)
+    #: Octets captured past the Length field, c.f.
+    #: :attr:`IPv4.trailer <pcapkit.protocols.schema.internet.ipv4.IPv4.trailer>`.
+    trailer: 'bytes' = PayloadField(default=b'')
 
     if TYPE_CHECKING:
         def __init__(self, srcport: 'Enum_AppType | int', dstport: 'Enum_AppType | int', len: 'int',
-                     checksum: 'bytes', payload: 'bytes | Schema | ProtocolBase') -> 'None': ...
+                     checksum: 'bytes', payload: 'bytes | Schema | ProtocolBase',
+                     trailer: 'bytes' = b'') -> 'None': ...

pcapkit/protocols/transport/udp.py

@@ -173,6 +173,14 @@ def read(self, length: 'Optional[int]' = None, **kwargs: 'Any') -> 'Data_UDP':
             checksum=schema.checksum,
         )
 
+        # NOTE: Octets captured past the Length field belong to this layer,
+        # since its own length field is what leaves them out of the payload,
+        # and the rebuild keeps them as captured (:issue:`1461`).
+        if schema.trailer:
+            udp.__update__([
+                ('trailer', schema.trailer),
+            ])
+
         return self._decode_next_layer(udp, (udp.srcport.port, udp.dstport.port), udp.len - 8)
 
     def make(self,
@@ -181,6 +189,7 @@ def make(self,
              checksum: 'bytes' = b'\x00\x00',
              total_length: 'Optional[int]' = None,
              payload: 'bytes | Schema | ProtocolBase' = b'',
+             trailer: 'bytes' = b'',
              **kwargs: 'Any') -> 'Schema_UDP':
         """Make (construct) packet data.
 
@@ -191,6 +200,8 @@ def make(self,
             total_length: Length of the datagram, header included; computed from
                 the header and the payload when omitted.
             payload: Payload data.
+            trailer: Octets after the datagram, outside the Length field;
+                written as is and not counted in a computed ``total_length``.
             **kwargs: Arbitrary keyword arguments.
 
         Returns:
@@ -211,6 +222,7 @@ def make(self,
             len=total_length,
             checksum=checksum,
             payload=payload,
+            trailer=trailer,
         )
 
     ##########################################################################
@@ -257,4 +269,6 @@ def _make_data(cls, data: 'Data_UDP') -> 'dict[str, Any]':  # type: ignore[overr
             # the length that was on the wire (:issue:`1155`).
             'total_length': data.len,
             'payload': cls._make_payload(data),
+            # NOTE: And the octets captured past the Length field (:issue:`1461`).
+            'trailer': data.get('trailer', b''),
         }

tests/protocols/transport/test_udp_trailer_kept_unit.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+"""Octets captured past a UDP datagram's own Length field survive a rebuild.
+
+GitHub issue #1461: :class:`~pcapkit.protocols.transport.udp.UDP` reads its
+Length field, but the payload ran to the end of the capture and nothing kept
+what lay past an empty one, so
+:meth:`~pcapkit.protocols.protocol.ProtocolBase.from_data` dropped it: a header
+of Length 8 followed by two octets rebuilt as 8 octets of its 10.
+
+The payload now stops at Length and the rest is ``info.trailer``, as IPv4, IPv6
+(#1209) and IPX (#1433) do, and :meth:`make` writes a ``trailer`` outside the
+length it declares.
+
+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`,
+c.f. ``tests/protocols/link/test_ethernet_mac_roundtrip_unit.py``.
+
+"""
+
+import unittest
+import warnings
+
+from tests._support import reimport_once_per_class
+
+
+def _udp(length: int, port: int = 53) -> bytes:
+    """Return an 8-octet UDP header between ``port`` and itself, of Length ``length``."""
+    return port.to_bytes(2, 'big') * 2 + length.to_bytes(2, 'big') + b'\x12\x34'
+
+
+#: Datagrams with octets past the Length field: (octets, payload, trailer).
+TRAILER = {
+    'header only': (_udp(8) + b'pp', b'', b'pp'),
+    'payload': (_udp(10) + b'ab' + b'pp', b'ab', b'pp'),
+    'HTTP port': (_udp(12, 80) + b'GET ' + b'pp', b'GET ', b'pp'),
+}
+
+#: Datagrams that end at or before their own Length, or whose Length is below
+#: the header: (octets, payload).
+NO_TRAILER = {
+    'exact': (_udp(10) + b'ab', b'ab'),
+    'truncated': (_udp(255) + b'abcd', b'abcd'),
+    'Length below header': (_udp(3) + b'abcd', b'abcd'),
+    'Length zero': (_udp(0) + b'abcd', b'abcd'),
+}
+
+
+class TestUDPTrailerKept(unittest.TestCase):
+    """Pin the octets past a UDP Length field through ``from_data``."""
+
+    def setUp(self) -> None:
+        reimport_once_per_class(self)
+
+    def test_datagram_rebuilds_byte_for_byte(self) -> None:
+        from pcapkit.protocols.transport.udp import UDP
+
+        for name, (raw, payload, trailer) in TRAILER.items():
+            with self.subTest(datagram=name):
+                parsed = UDP(raw, len(raw))
+                self.assertEqual(parsed.info.trailer, trailer)
+                self.assertEqual(bytes(parsed.__header__.payload), payload)
+                self.assertEqual(UDP.from_data(parsed.info).data, raw)
+                self.assertEqual(UDP.from_data(parsed.info.to_dict()).data, raw)
+
+    def test_no_trailer_key_without_excess(self) -> None:
+        from pcapkit.protocols.transport.udp import UDP
+
+        for name, (raw, payload) in NO_TRAILER.items():
+            with self.subTest(datagram=name), warnings.catch_warnings():
+                warnings.simplefilter('error')
+                parsed = UDP(raw, len(raw))
+                self.assertNotIn('trailer', parsed.info)
+                self.assertEqual(bytes(parsed.__header__.payload), payload)
+                self.assertEqual(UDP.from_data(parsed.info).data, raw)
+                self.assertEqual(UDP.from_data(parsed.info.to_dict()).data, raw)
+
+    def test_make_writes_trailer_outside_length(self) -> None:
+        from pcapkit.protocols.transport.udp import UDP
+
+        made = UDP(srcport=53, dstport=53, payload=b'ab', trailer=b'pp')
+        self.assertEqual(made.data, _udp(10)[:6] + b'\x00\x00' + b'ab' + b'pp')
+        self.assertEqual(made.info.len, 10)
+        self.assertEqual(made.info.trailer, b'pp')
+
+        parsed = UDP(made.data, len(made.data))
+        self.assertEqual(UDP.from_data(parsed.info.to_dict()).data, made.data)
+
+    def test_ipv4_keeps_both_trailers(self) -> None:
+        from pcapkit.protocols.internet.ipv4 import IPv4
+
+        udp = _udp(9) + b'a' + b'XYZ'
+        total = (20 + len(udp)).to_bytes(2, 'big')
+        raw = (b'\x45\x00' + total + bytes(4) + b'\x40\x11' + bytes(2)
+               + bytes([10, 0, 0, 1, 10, 0, 0, 2]) + udp + b'\x00\x00')
+
+        parsed = IPv4(raw, len(raw))
+        self.assertEqual(parsed.info.trailer, b'\x00\x00')
+        self.assertEqual(parsed.info.udp.trailer, b'XYZ')
+        self.assertEqual(IPv4.from_data(parsed.info.to_dict()).data, raw)
+
+
+if __name__ == '__main__':
+    unittest.main()