diff --git a/pcapkit/protocols/schema/schema.py b/pcapkit/protocols/schema/schema.py index 18f8e2e030..0c851c2b08 100644 --- a/pcapkit/protocols/schema/schema.py +++ b/pcapkit/protocols/schema/schema.py @@ -99,7 +99,7 @@ class marked with a user's own ``from typing import final`` on 3.10 is temp = ['__map__', '__map_reverse__', '__builtin__', '__fields__', '__buffer__', '__updated__', - '__payload__', '__finalised__'] + '__payload__', '__finalised__', '__remainder__', '__remainder_offset__'] temp.extend(cls.__additional__) for obj in cls.mro(): temp.extend(el for el in dir(obj) if el not in cls.__fields__) @@ -374,6 +374,11 @@ class Schema(Mapping[str, _VT], Generic[_VT], metaclass=SchemaMeta): __buffer__: 'dict[str, bytes]' #: Flag for whether the schema is recently updated. __updated__: 'bool' + #: Octets of an enclosing :class:`~pcapkit.corekit.fields.misc.SchemaField` + #: span that this nested schema left unread; set only when non-empty. + __remainder__: 'bytes' + #: Octets this nested schema read before :attr:`__remainder__`. + __remainder_offset__: 'int' #: Flag for finalised class initialisation. __finalised__: 'FinalisedState' = FinalisedState.NONE @@ -554,6 +559,13 @@ def __update__(self, dict_: 'Optional[Mapping[str, _VT] | Iterable[tuple[str, _V data_iter = itertools.chain(dict_, kwargs.items()) for (key, value) in data_iter: + # NOTE: Not a field, but what :meth:`to_dict` emits for the octets an + # enclosing ``SchemaField`` span held beyond this schema, and where, + # so that a rebuilt schema still re-emits them (:issue:`1380`). + if key in ('__remainder__', '__remainder_offset__'): + self.__dict__[key] = value + continue + if key not in self.__buffer__: warn(f'{key!r} is not a valid field name', UnknownFieldWarning) continue @@ -691,6 +703,15 @@ def to_dict(self) -> 'dict[str, _VT]': types, such as :obj:`list`, :obj:`tuple`, :obj:`set`, etc., we shall not convert them into :obj:`dict` and remain them intact. + Note: + A schema unpacked as the value of a + :class:`~pcapkit.corekit.fields.misc.SchemaField` that left part of + the field's span unread carries those octets under a + ``__remainder__`` key, and the number of octets it did read under + ``__remainder_offset__``. :meth:`from_dict` accepts both back, so + that the rebuilt field still packs them (:issue:`1380`). The keys + are absent otherwise. + """ dict_ = {} # type: dict[str, Any] for (key, value) in self.__dict__.items(): @@ -702,6 +723,10 @@ def to_dict(self) -> 'dict[str, _VT]': dict_[out_key] = value.to_dict() else: dict_[out_key] = value + + for key in ('__remainder__', '__remainder_offset__'): + if key in self.__dict__: + dict_[key] = self.__dict__[key] return dict_ def to_bytes(self) -> 'bytes': @@ -862,19 +887,31 @@ class will consider negative value as a placeholder. 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. + # it left unread travel with the nested value as its + # ``__remainder__`` -- on the schema itself, or as a key of the + # mapping :meth:`to_dict` made of it -- so they are re-emitted + # from there, as long as the nested schema still packs to the + # ``__remainder_offset__`` octets it was parsed at. Re-emitting + # them after a nested schema that has shrunk or grown would put + # stale octets at a shifted offset, so then, or for a schema that + # never had any, the rest of the span is zero-filled. # Either way the field keeps the width its length field declares - # (:issue:`1292`). + # (:issue:`1292`), and a schema rebuilt through :meth:`from_dict` + # packs what the parsed one did (:issue:`1380`). # # 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'' + if isinstance(data, Schema): + record = data.__dict__ # type: Mapping[str, Any] + elif isinstance(data, collections.abc.Mapping): + record = data + else: + record = {} + tail = record.get('__remainder__', b'') + if record.get('__remainder_offset__') != len(temp) or len(temp) + len(tail) > field.length: + tail = b'' temp += tail + bytes(field.length - len(temp) - len(tail)) self.__buffer__[field.name] = temp @@ -998,6 +1035,16 @@ def unpack(cls, data: 'bytes | IO[bytes]', self.__buffer__[field.name] = byte value = field.unpack(byte, packet.copy()) + if isinstance(field, SchemaField) and isinstance(value, Schema): + # NOTE: Kept on the nested schema rather than only in + # ``__buffer__``, so that the octets it left unread go wherever + # the value goes, :meth:`to_dict` included; see :meth:`pack`. + # The width read is summed from the nested buffers rather than + # taken with :func:`len`, which may pack the schema again. + consumed = sum(map(len, value.__buffer__.values())) + if len(byte) > consumed: + value.__dict__['__remainder__'] = byte[consumed:] + value.__dict__['__remainder_offset__'] = consumed setattr(self, field.name, value) packet[field.name] = value diff --git a/tests/protocols/schema/test_schema_field_remainder_unit.py b/tests/protocols/schema/test_schema_field_remainder_unit.py new file mode 100644 index 0000000000..edacb8bf58 --- /dev/null +++ b/tests/protocols/schema/test_schema_field_remainder_unit.py @@ -0,0 +1,208 @@ +# -*- coding: utf-8 -*- +"""A :class:`~pcapkit.corekit.fields.misc.SchemaField`'s unread octets survive a rebuild. + +GitHub issue #1380: a nested schema may read less than its +:class:`~pcapkit.corekit.fields.misc.SchemaField`'s declared span. Parsing then +packing kept the octets it left unread (#1292), but a schema rebuilt through +:meth:`~pcapkit.protocols.schema.schema.Schema.from_dict`, or a protocol rebuilt +through :meth:`~pcapkit.protocols.protocol.ProtocolBase.from_data`, had never +been unpacked and zero-filled them. They now travel with the nested value as its +``__remainder__``. + +No shipped schema declares a :class:`~pcapkit.corekit.fields.misc.SchemaField` +directly, so the schemas and the protocol here are built in the test. Everything +from :mod:`pcapkit` is imported inside each test, after +:func:`~tests._support.reimport_once_per_class`, not at module load. + +""" + +import unittest +import warnings + +from tests._support import reimport_once_per_class + +#: ``n`` = 3, a 3-octet span whose nested schema reads only ``02``, then ``tail``. +RAW = bytes.fromhex('0302eeee04') + + +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 + + +#: A 6-octet span whose nested schema reads ``k`` = 2 and ``aabb``, then ``tail``. +SIZED_RAW = bytes.fromhex('02aabbccddee04') + + +def _sized_schemas(): # type: ignore[no-untyped-def] + """An outer schema whose nested schema's own width depends on its values.""" + from pcapkit.corekit.fields.misc import SchemaField + from pcapkit.corekit.fields.numbers import UInt8Field + from pcapkit.corekit.fields.strings import BytesField + from pcapkit.protocols.schema.schema import Schema, schema_final + + @schema_final + class Inner(Schema): + k: int = UInt8Field() + body: bytes = BytesField(length=lambda pkt: pkt['k']) + + @schema_final + class Outer(Schema): + inner: Inner = SchemaField(length=6, schema=Inner) + tail: int = UInt8Field() + + return Inner, Outer + + +def _protocol(outer): # type: ignore[no-untyped-def] + """A protocol over ``outer`` whose data model holds the nested schema's mapping.""" + from pcapkit.corekit.infoclass import info_final + from pcapkit.corekit.protochain import ProtoChain + from pcapkit.protocols.data.data import Data + from pcapkit.protocols.protocol import ProtocolBase + + @info_final + class OuterData(Data): + n: int + inner: dict + tail: int + + class OuterProtocol(ProtocolBase, schema=outer, data=OuterData): # type: ignore[call-arg,misc] + __layer__ = 'Internet' + + @property + def name(self): # type: ignore[no-untyped-def] + return 'Outer' + + @property + def length(self): # type: ignore[no-untyped-def] + return len(self.__header__) + + def read(self, length=None, **kwargs): # type: ignore[no-untyped-def] + from pcapkit.protocols.misc.null import NoPayload + + schema = self.__header__ + self._next = NoPayload() + self._protos = ProtoChain(type(self), self.alias, basis=self._next.protochain) + # NOTE: A made header still holds the mapping ``make`` was given. + inner = schema.inner if isinstance(schema.inner, dict) else schema.inner.to_dict() + return OuterData(n=schema.n, inner=inner, tail=schema.tail) + + def make(self, n=0, inner=None, tail=0, **kwargs): # type: ignore[no-untyped-def] + return outer(n=n, inner=inner, tail=tail) + + @classmethod + def __index__(cls): # type: ignore[no-untyped-def,override] + return 250 + + return OuterProtocol + + +class TestSchemaFieldRemainder(unittest.TestCase): + """Pin the unread octets of a nested schema through every rebuild.""" + + def setUp(self) -> None: + reimport_once_per_class(self) + + def test_to_dict_carries_the_remainder(self) -> None: + _, outer = _nested_schemas() + + parsed = outer.unpack(RAW, len(RAW), None) + self.assertEqual(parsed.to_dict(), { + 'n': 3, 'inner': {'x': 2, '__remainder__': b'\xee\xee', '__remainder_offset__': 1}, 'tail': 4, + }) + + def test_from_dict_rebuilds_byte_exactly(self) -> None: + _, outer = _nested_schemas() + + parsed = outer.unpack(RAW, len(RAW), None) + with warnings.catch_warnings(): + warnings.simplefilter('error') + rebuilt = outer.from_dict(parsed.to_dict()) + packed = rebuilt.pack() + self.assertEqual(packed, RAW) + + def test_parsed_nested_schema_keeps_its_remainder(self) -> None: + _, outer = _nested_schemas() + + parsed = outer.unpack(RAW, len(RAW), None) + self.assertEqual(outer(n=3, inner=parsed.inner, tail=4).pack(), RAW) + + def test_remainder_is_not_a_field(self) -> None: + _, outer = _nested_schemas() + + inner = outer.unpack(RAW, len(RAW), None).inner + self.assertEqual(dict(inner), {'x': 2}) + self.assertEqual(bytes(inner), b'\x02') + with self.assertRaises(KeyError): + inner['__remainder__'] # pylint: disable=pointless-statement,unsubscriptable-object + + def test_remainder_that_no_longer_fits_is_dropped(self) -> None: + _, outer = _nested_schemas() + + data = outer.unpack(RAW, len(RAW), None).to_dict() + data['n'] = 2 + self.assertEqual(outer.from_dict(data).pack(), bytes.fromhex('02020004')) + + def test_shrunk_in_place_is_zero_filled(self) -> None: + _, outer = _sized_schemas() + + parsed = outer.unpack(SIZED_RAW, len(SIZED_RAW), None) + parsed.inner.body = b'\xaa' + parsed.inner.k = 1 + self.assertEqual(parsed.pack(), bytes.fromhex('01aa0000000004')) + + def test_shrunk_through_from_dict_is_zero_filled(self) -> None: + _, outer = _sized_schemas() + + data = outer.unpack(SIZED_RAW, len(SIZED_RAW), None).to_dict() + data['inner'].update(k=1, body=b'\xaa') + self.assertEqual(outer.from_dict(data).pack(), bytes.fromhex('01aa0000000004')) + + def test_grown_through_from_dict_is_zero_filled(self) -> None: + _, outer = _sized_schemas() + + data = outer.unpack(SIZED_RAW, len(SIZED_RAW), None).to_dict() + data['inner'].update(k=4, body=b'\xaa\xbb\xcc\xdd') + self.assertEqual(outer.from_dict(data).pack(), bytes.fromhex('04aabbccdd0004')) + + def test_sized_round_trips_byte_exactly(self) -> None: + _, outer = _sized_schemas() + + parsed = outer.unpack(SIZED_RAW, len(SIZED_RAW), None) + self.assertEqual(parsed.pack(), SIZED_RAW) + self.assertEqual(outer.from_dict(parsed.to_dict()).pack(), SIZED_RAW) + + def test_no_remainder_without_unread_octets(self) -> None: + _, outer = _nested_schemas() + + raw = bytes.fromhex('010204') + parsed = outer.unpack(raw, len(raw), None) + self.assertNotIn('__remainder__', parsed.to_dict()['inner']) + self.assertEqual(outer.from_dict(parsed.to_dict()).pack(), raw) + + def test_protocol_from_data_rebuilds_byte_exactly(self) -> None: + _, outer = _nested_schemas() + protocol = _protocol(outer) + + parsed = protocol(RAW, len(RAW)) + self.assertEqual(bytes(protocol.from_data(parsed.info)), RAW) + self.assertEqual(bytes(protocol.from_data(parsed.info.to_dict())), RAW) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/protocols/schema/test_schema_roundtrip_unit.py b/tests/protocols/schema/test_schema_roundtrip_unit.py index 5496c3b457..d7b90239f0 100644 --- a/tests/protocols/schema/test_schema_roundtrip_unit.py +++ b/tests/protocols/schema/test_schema_roundtrip_unit.py @@ -11,7 +11,8 @@ 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. +the octets the nested schema left unread were lost; GitHub issue #1380: a +schema rebuilt through ``from_dict`` still zero-filled them. Every case builds its own octets in memory and reads no capture. Everything from :mod:`pcapkit` is imported inside each test, after @@ -98,10 +99,10 @@ def test_schema_field_keeps_declared_width(self) -> None: 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')) + # NOTE: The unread octets travel through ``to_dict`` as the nested + # value's ``__remainder__`` (#1380); a nested schema that never had any + # is zero-filled to the declared width. + self.assertEqual(rebuilt.pack(), raw) self.assertEqual(outer(n=3, inner=inner(x=2), tail=4).pack(), bytes.fromhex('0302000004'))