diff --git a/pcapkit/protocols/schema/schema.py b/pcapkit/protocols/schema/schema.py index f380d5750d..18f8e2e030 100644 --- a/pcapkit/protocols/schema/schema.py +++ b/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) diff --git a/tests/protocols/schema/test_schema_roundtrip_unit.py b/tests/protocols/schema/test_schema_roundtrip_unit.py new file mode 100644 index 0000000000..5496c3b457 --- /dev/null +++ b/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()