fix(schema): Schema.pack keeps option padding, SchemaField spans and iterated keys (#1289, #1291, #1292)

pcapkit/protocols/schema/schema.py

@@ -10,7 +10,8 @@
 
 from pcapkit.corekit.fields.collections import ListField, OptionField
 from pcapkit.corekit.fields.field import NO_VALUE, FieldBase
-from pcapkit.corekit.fields.misc import ConditionalField, ForwardMatchField, PayloadField
+from pcapkit.corekit.fields.misc import (ConditionalField, ForwardMatchField, PayloadField,
+                                         SchemaField)
 from pcapkit.corekit.fields.strings import PaddingField
 from pcapkit.corekit.infoclass import FinalisedState
 from pcapkit.utilities.compat import Mapping, final
@@ -626,10 +627,16 @@ def __iter__(self) -> 'Iterator[str]':
             yield self.__map_reverse__.get(key, key)
 
     def __getitem__(self, name: 'str') -> '_VT':
-        if name in self.__fields__:
+        # NOTE: Serves every key :meth:`__iter__` yields, not only the declared
+        # fields: a ``post_process`` may store derived values on the instance,
+        # as :class:`~pcapkit.protocols.schema.internet.ipv4.TSOption` does, and
+        # iteration reports them, so ``dict(schema)`` must be able to read them
+        # back (:issue:`1291`). A mangled key is reachable only under the name
+        # it was mangled from, as in iteration.
+        if name in self.__fields__ or (name not in self.__builtin__ and name not in self.__map_reverse__):
             key = self.__map__.get(name, name)
             return self.__dict__[key]
-        return super().__getitem__(name)
+        raise KeyError(name)
 
     def __setattr__(self, name: 'str', value: '_VT') -> 'None':
         if name in self.__fields__:
@@ -817,6 +824,16 @@ class will consider negative value as a placeholder.
                     self.__buffer__[field.name] = field.pack(list(data), packet)
                 else:
                     raise ProtocolUnbound(f'unsupported type {type(data)}')
+
+                if isinstance(field, OptionField):
+                    # NOTE: The packing mirror of the ``__option_padding__`` that
+                    # :meth:`unpack` records: whatever of the declared option area
+                    # the packed options do not fill. Without it every padding
+                    # field sized from that key packs nothing, and the header
+                    # comes out shorter than its own length field says
+                    # (:issue:`1289`). An undeclared area (a negative length)
+                    # leaves no remainder.
+                    packet['__option_padding__'] = max(field.length - len(self.__buffer__[field.name]), 0)
                 continue
 
             if isinstance(field, PaddingField):
@@ -841,6 +858,24 @@ class will consider negative value as a placeholder.
                 temp = field.pack(data, packet)
             except NoDefaultValue:
                 temp = bytes(field.length)
+
+            if isinstance(field, SchemaField) and len(temp) < field.length:
+                # NOTE: :meth:`unpack` hands a nested schema its whole declared
+                # span, and the schema may read less of it than that. The octets
+                # it left unread are still in the buffer :meth:`unpack` stored, so
+                # they are re-emitted from there, as long as the buffer still
+                # starts with what the nested schema packs to now; otherwise, or
+                # for a schema that was never unpacked, the span is zero-filled.
+                # Either way the field keeps the width its length field declares
+                # (:issue:`1292`).
+                #
+                # A field reached through a
+                # :class:`~pcapkit.corekit.fields.misc.SwitchField` is not handled
+                # here: the PCAP-NG block selector declares the rest of the
+                # *stream* as its span, which is no part of the field.
+                buffer = self.__buffer__[field.name]
+                tail = buffer[len(temp):field.length] if buffer.startswith(temp) else b''
+                temp += tail + bytes(field.length - len(temp) - len(tail))
             self.__buffer__[field.name] = temp
 
         self.post_process(packet)

tests/protocols/schema/test_schema_roundtrip_unit.py

@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+"""Base :class:`~pcapkit.protocols.schema.schema.Schema` round trips.
+
+GitHub issue #1289: nothing on the packing path set ``__option_padding__``, so
+a padding field sized from it packed nothing, and a TCP or IPv4 header rebuilt
+from its schema came out shorter than its own length field said.
+
+GitHub issue #1291: iterating a schema yielded the values its
+``post_process`` stores on the instance, which ``__getitem__`` refused, so
+``dict(schema)`` raised :exc:`KeyError`.
+
+GitHub issue #1292: a :class:`~pcapkit.corekit.fields.misc.SchemaField` read
+its whole declared span but packed only what its nested schema consumed, so
+the octets the nested schema left unread were lost.
+
+Every case builds its own octets in memory and reads no capture. Everything
+from :mod:`pcapkit` is imported inside each test, after
+:func:`~tests._support.reimport_once_per_class`, not at module load.
+
+"""
+
+import unittest
+
+from tests._support import reimport_once_per_class
+
+#: TCP header, Data Offset 6: NOP, EOOL, then two octets of padding.
+TCP_HEADER = bytes.fromhex('0050c3500000000100000000' '6002ffff00000000' '01000000')
+#: IPv4 header, IHL 6: NOP, EOOL, then two octets of padding.
+IPV4_HEADER = bytes.fromhex('460000180000000040060000c0a80001c0a80002' '01000000')
+
+
+def _nested_schemas():  # type: ignore[no-untyped-def]
+    """An outer schema whose nested schema reads less than its span."""
+    from pcapkit.corekit.fields.misc import SchemaField
+    from pcapkit.corekit.fields.numbers import UInt8Field
+    from pcapkit.protocols.schema.schema import Schema, schema_final
+
+    @schema_final
+    class Inner(Schema):
+        x: int = UInt8Field()
+
+    @schema_final
+    class Outer(Schema):
+        n: int = UInt8Field()
+        inner: Inner = SchemaField(length=lambda pkt: pkt['n'], schema=Inner)
+        tail: int = UInt8Field()
+
+    return Inner, Outer
+
+
+class TestSchemaRoundTrip(unittest.TestCase):
+    """Pin the base schema's packing and mapping behaviour."""
+
+    def setUp(self) -> None:
+        reimport_once_per_class(self)
+
+    def test_tcp_rebuild_keeps_option_padding(self) -> None:
+        from pcapkit.protocols.transport.tcp import TCP
+
+        parsed = TCP(TCP_HEADER, len(TCP_HEADER))
+        rebuilt = TCP.from_schema(parsed.schema.to_dict())
+        self.assertEqual(bytes(rebuilt), TCP_HEADER)
+
+    def test_ipv4_rebuild_keeps_option_padding(self) -> None:
+        from pcapkit.protocols.internet.ipv4 import IPv4
+
+        parsed = IPv4(IPV4_HEADER, len(IPV4_HEADER))
+        rebuilt = IPv4.from_schema(parsed.schema.to_dict())
+        self.assertEqual(bytes(rebuilt), IPV4_HEADER)
+
+    def test_iterated_keys_are_readable(self) -> None:
+        from pcapkit.protocols.schema.internet.ipv4 import TSOption
+
+        schema = TSOption.unpack(bytes.fromhex('4408090000000001'), 8, None)
+        keys = list(schema)
+        self.assertIn('timestamp', keys)
+        for key in keys:
+            with self.subTest(key=key):
+                schema[key]  # pylint: disable=pointless-statement
+        self.assertEqual(list(dict(schema)), keys)
+
+    def test_builtin_names_are_not_items(self) -> None:
+        from pcapkit.protocols.schema.internet.ipv4 import TSOption
+
+        schema = TSOption.unpack(bytes.fromhex('4408090000000001'), 8, None)
+        for name in ('__buffer__', 'pack', 'nonexistent'):
+            with self.subTest(name=name), self.assertRaises(KeyError):
+                schema[name]  # pylint: disable=pointless-statement
+
+    def test_schema_field_keeps_unread_octets(self) -> None:
+        _, outer = _nested_schemas()
+
+        raw = bytes.fromhex('0302eeee04')
+        self.assertEqual(outer.unpack(raw, len(raw), None).pack(), raw)
+
+    def test_schema_field_keeps_declared_width(self) -> None:
+        inner, outer = _nested_schemas()
+
+        raw = bytes.fromhex('0302eeee04')
+        rebuilt = outer.from_dict(outer.unpack(raw, len(raw), None).to_dict())
+        # NOTE: The width is kept but the unread octets are zero-filled, since a
+        # rebuilt schema was never unpacked; that open half is #1380, which will
+        # change this expectation to ``0302eeee04``.
+        self.assertEqual(rebuilt.pack(), bytes.fromhex('0302000004'))
+        self.assertEqual(outer(n=3, inner=inner(x=2), tail=4).pack(), bytes.fromhex('0302000004'))
+
+
+if __name__ == '__main__':
+    unittest.main()