diff --git a/pcapkit/foundation/registry/protocols.py b/pcapkit/foundation/registry/protocols.py index bff9d12e28..24774b4c74 100644 --- a/pcapkit/foundation/registry/protocols.py +++ b/pcapkit/foundation/registry/protocols.py @@ -88,6 +88,7 @@ from pcapkit.const.reg.transtype import TransType from pcapkit.const.sctp.payload_protocol_identifier import \ PayloadProtocolIdentifier as SCTP_PayloadProtocolIdentifier + from pcapkit.corekit.enum import EnumRegistry from pcapkit.corekit.sentinels import NullType from pcapkit.protocols.application.httpv2 import FrameConstructor as HTTP_FrameConstructor from pcapkit.protocols.application.httpv2 import FrameParser as HTTP_FrameParser @@ -415,13 +416,32 @@ def _register_name(protocol: 'Type[ProtocolBase]', *targets: 'tuple[type, Any]') #: table so :func:`register_protocol_code` can consult it. :class:`~pcapkit.const.reg.linktype.LinkType` naming #: more than one class is deliberate, not ambiguous -- :func:`register_linktype` #: already fans out to each of them. -_CODE_DESTINATIONS: 'dict[type, tuple[Type[ProtocolBase], ...]]' = { +_CODE_DESTINATIONS: 'dict[Type[EnumRegistry], tuple[Type[ProtocolBase], ...]]' = { Enum_EtherType: (Link,), Enum_TransType: (Internet,), Enum_PayloadProtocolIdentifier: (SCTP,), Enum_LinkType: (Frame, PCAPNG), } +#: Key type of each ``__proto__`` dispatch registry a ``code=`` target can name. +#: +#: Maps the class owning the registry to the enum its typed registrar checks +#: ``code`` against, and whether a bare :class:`int` is a valid key there too -- +#: the same :func:`_check_code` arguments that registrar passes. A destination is +#: looked up by the class it inherits ``__proto__`` from, so a subclass sharing +#: :class:`Link`'s registry is checked as :class:`Link` is, while one keeping a +#: registry of its own, such as :class:`~pcapkit.protocols.link.loopback.Loopback`, +#: names no entry here and is not checked (:issue:`2268`). +_DESTINATION_KINDS: 'dict[Type[ProtocolBase], tuple[Type[EnumRegistry], bool]]' = { + Link: (Enum_EtherType, False), # type: ignore[type-abstract] + Internet: (Enum_TransType, False), # type: ignore[type-abstract] + Frame: (Enum_LinkType, False), + PCAPNG: (Enum_LinkType, False), + TCP: (Enum_AppType, True), + UDP: (Enum_AppType, True), + SCTP: (Enum_PayloadProtocolIdentifier, True), +} + def _iter_code_targets(code: 'Any') -> 'Iterator[tuple[Type[ProtocolBase], Any]]': """Flatten a ``code=`` argument into ``(destination, key)`` pairs. @@ -444,7 +464,7 @@ def _iter_code_targets(code: 'Any') -> 'Iterator[tuple[Type[ProtocolBase], Any]] return if isinstance(code, (enum.Enum, aenum.Enum)): - destinations = _CODE_DESTINATIONS.get(type(code)) + destinations = _CODE_DESTINATIONS.get(cast('Type[EnumRegistry]', type(code))) if destinations is None: raise RegistryError( f'no destination registry is known for enum type {type(code).__name__!r} ' @@ -527,10 +547,22 @@ def register_protocol_code(protocol: 'Type[ProtocolBase]', code: 'Any') -> 'None RegistryError: If a bare :class:`int` (or any other non-enum, non-mapping value) is given without an explicit destination, or if an enum member's type names no known destination registry -- - inference refuses rather than guesses. + inference refuses rather than guesses; or if a key is not of the + type its destination's registry is keyed by, as its typed registrar + would reject it, e.g. ``{Frame: TransType.ICMP}``. Every target is + checked before any is registered, so a rejected ``code`` leaves + every registry as it was (:issue:`2268`). """ - for destination, key in _iter_code_targets(code): + targets = list(_iter_code_targets(code)) + for destination, key in targets: + owner = next((klass for klass in getattr(destination, '__mro__', ()) + if '__proto__' in vars(klass)), None) + if owner in _DESTINATION_KINDS: + kind, number = _DESTINATION_KINDS[owner] + _check_code(key, kind, number=number) + + for destination, key in targets: destination.register(key, protocol) logger.debug('registered %s into %s.__proto__: %s', protocol.__name__, destination.__name__, key) diff --git a/tests/protocols/test_protocol_code_registration_unit.py b/tests/protocols/test_protocol_code_registration_unit.py index 3e2151009f..0341327308 100644 --- a/tests/protocols/test_protocol_code_registration_unit.py +++ b/tests/protocols/test_protocol_code_registration_unit.py @@ -31,10 +31,15 @@ import importlib.util import sys import unittest +import warnings +from typing import TYPE_CHECKING from unittest import mock from tests._support import reimport_once_per_class +if TYPE_CHECKING: + from typing import Any + #: Test tier (#1958): branches on the running interpreter's version or implementation. __tier__ = 'ALL' @@ -407,5 +412,194 @@ def test_backward_compatible_dispatch_tables_are_unaffected(self) -> None: ('pcapkit.protocols.internet.ipv4', 'IPv4')) +@unittest.skipUnless(HAS_RUNTIME, 'runtime dependencies not installed') +class RegisterProtocolCodeKeyCheckTests(unittest.TestCase): + """Each ``code=`` key is checked against its destination's key type, all before any registers (#2268). + + The typed registrars reject a member of another enum (#2224), but + ``register_protocol_code`` handed every ``{destination: key}`` pair straight + to ``destination.register``. The registry enums compare equal to their + values, so ``{Frame: TransType.ICMP}`` replaced the shipped entry at + ``LinkType.ETHERNET`` with only a warning. It also registered each target as + it flattened them, so a bad later target left the earlier ones in place. + + Each case snapshots every dispatch registry and the name registry, expects + :exc:`~pcapkit.utilities.exceptions.RegistryError`, compares each registry + with its snapshot, and restores the snapshot whatever happened. + """ + + def setUp(self) -> None: + reimport_once_per_class(self) + + def _assert_rejected_unchanged(self, call: 'Any') -> None: + from pcapkit.protocols import __proto__ as protocol_registry + from pcapkit.protocols.internet.internet import Internet + from pcapkit.protocols.link.link import Link + from pcapkit.protocols.misc.pcap.frame import Frame + from pcapkit.protocols.misc.pcapng import PCAPNG + from pcapkit.protocols.transport.sctp import SCTP + from pcapkit.protocols.transport.tcp import TCP + from pcapkit.protocols.transport.udp import UDP + from pcapkit.utilities.exceptions import RegistryError + + tables = {klass.__name__: klass.__proto__ + for klass in (Link, Internet, Frame, PCAPNG, TCP, UDP, SCTP)} + tables['pcapkit.protocols'] = protocol_registry + before = {name: dict(table) for name, table in tables.items()} + try: + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + with self.assertRaises(RegistryError): + call() + for name, table in tables.items(): + if name != 'pcapkit.protocols': + self.assertEqual(dict(table), before[name], f'{name}.__proto__ changed') + finally: + for name, table in tables.items(): + table.clear() + table.update(before[name]) + + def test_mapping_rejects_a_member_of_another_enum(self) -> None: + """The issue's example: ``TransType.ICMP`` equals ``LinkType.ETHERNET``.""" + from pcapkit.const.reg.linktype import LinkType + from pcapkit.const.reg.transtype import TransType + from pcapkit.foundation.registry.protocols import register_protocol_code + from pcapkit.protocols.misc.pcap.frame import Frame + from pcapkit.protocols.misc.raw import Raw + + self.assertEqual(TransType.ICMP, LinkType.ETHERNET) + self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {Frame: TransType.ICMP})) + + def test_class_keyword_rejects_a_member_of_another_enum(self) -> None: + from pcapkit.const.reg.transtype import TransType + from pcapkit.protocols.misc.pcap.frame import Frame + from pcapkit.protocols.misc.raw import Raw + + def define() -> None: + class Evil(Raw, code={Frame: TransType.ICMP}): + pass + + self._assert_rejected_unchanged(define) + + def test_subclass_sharing_a_registry_is_checked_as_its_owner(self) -> None: + """``Ethernet`` inherits ``Link.__proto__``, so its keys are ``EtherType``s.""" + from pcapkit.const.reg.transtype import TransType + from pcapkit.foundation.registry.protocols import register_protocol_code + from pcapkit.protocols.link.ethernet import Ethernet + from pcapkit.protocols.misc.raw import Raw + + self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {Ethernet: TransType.ICMP})) + + def test_number_destination_rejects_a_member_of_another_enum(self) -> None: + from pcapkit.const.reg.linktype import LinkType + from pcapkit.foundation.registry.protocols import register_protocol_code + from pcapkit.protocols.misc.raw import Raw + from pcapkit.protocols.transport.udp import UDP + + self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {UDP: LinkType(53)})) + + def test_a_bad_later_target_leaves_the_earlier_ones_unregistered(self) -> None: + """The issue's list example, and one whose bad target fails the key check instead.""" + from pcapkit.const.reg.transtype import TransType + from pcapkit.foundation.registry.protocols import register_protocol_code + from pcapkit.protocols.internet.internet import Internet + from pcapkit.protocols.misc.raw import Raw + + for code in ([TransType.L2TP, 5], [TransType.L2TP, {Internet: 115}]): + with self.subTest(code=code): + self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, code)) + + def test_number_and_own_registry_destinations_still_accept_an_int(self) -> None: + """A port is a valid ``TCP`` key; ``Loopback`` keeps an unchecked registry of its own.""" + from pcapkit.foundation.registry.protocols import register_protocol_code + from pcapkit.protocols.link.loopback import Loopback + from pcapkit.protocols.transport.tcp import TCP + + class FakeProtocol: + pass + + with mock.patch.object(TCP, 'register') as reg_tcp, \ + mock.patch.object(Loopback, 'register') as reg_loopback: + register_protocol_code(FakeProtocol, [{TCP: 54324}, {Loopback: 2}]) + reg_tcp.assert_called_once_with(54324, FakeProtocol) + reg_loopback.assert_called_once_with(2, FakeProtocol) + + def test_every_destination_kind_row_is_enforced(self) -> None: + """Every row of ``_DESTINATION_KINDS`` is exercised, both cells of it. + + The table is what the check is made of, so a row dropped from it, or a + ``kind``/``number`` cell changed in it, silently stops checking a + destination -- the registrars it stands in for are not called here, so + nothing else would notice. Each row is pinned by behaviour rather than + by comparing the table with a copy of itself: + + * ``kind`` -- a member of *another* registry's enum must be rejected, + and must leave every registry as it was. That is the #2224 confusion + the check exists for, so a row that stops being checked, or one + checked against the wrong enum, fails here. + * ``number`` -- a bare :class:`int` is accepted where the registrar + accepts one (a port, a payload protocol identifier) and rejected + where it does not, which is the only observable difference the flag + makes. The :class:`int` given to a rejecting row is the *value* of a + member its registry already ships an entry at, so the check is pinned + as being on the type and the entry at risk is a real one. + + The ``kind`` identity is asserted as well, since a row retyped to an + enum unrelated to every key used here would reject both keys for the + wrong reason. + """ + from pcapkit.const.reg.apptype import AppType + from pcapkit.const.reg.ethertype import EtherType + from pcapkit.const.reg.linktype import LinkType + from pcapkit.const.reg.transtype import TransType + from pcapkit.const.sctp.payload_protocol_identifier import PayloadProtocolIdentifier + from pcapkit.foundation.registry.protocols import (_DESTINATION_KINDS, + register_protocol_code) + from pcapkit.protocols.internet.internet import Internet + from pcapkit.protocols.link.link import Link + from pcapkit.protocols.misc.pcap.frame import Frame + from pcapkit.protocols.misc.pcapng import PCAPNG + from pcapkit.protocols.misc.raw import Raw + from pcapkit.protocols.transport.sctp import SCTP + from pcapkit.protocols.transport.tcp import TCP + from pcapkit.protocols.transport.udp import UDP + + # destination -> (kind, number, a member of another enum, a bare int) + expected = { + # 0x0800 is EtherType.Internet_Protocol_version_4, where Link ships IPv4. + Link: (EtherType, False, TransType.ICMP, 0x0800), + # 6 is TransType.TCP, where Internet ships TCP. + Internet: (TransType, False, EtherType.XEROX_NS_IDP, 6), + # The issue's own example: TransType.ICMP == LinkType.ETHERNET, where + # both Frame and PCAPNG ship Ethernet. + Frame: (LinkType, False, TransType.ICMP, 1), + PCAPNG: (LinkType, False, TransType.ICMP, 1), + TCP: (AppType, True, LinkType(53), 54324), + UDP: (AppType, True, LinkType(53), 54324), + SCTP: (PayloadProtocolIdentifier, True, TransType.ICMP, 54324), + } + + for destination, (kind, number, member, raw) in expected.items(): + with self.subTest(destination=destination.__name__): + self.assertNotIsInstance(member, kind) # the test data, not the table + + self._assert_rejected_unchanged( + lambda dest=destination, key=member: register_protocol_code(Raw, {dest: key})) + + if number: + with mock.patch.object(destination, 'register') as register: + register_protocol_code(Raw, {destination: raw}) + register.assert_called_once_with(raw, Raw) + else: + self._assert_rejected_unchanged( + lambda dest=destination, key=raw: register_protocol_code(Raw, {dest: key})) + + self.assertIs(_DESTINATION_KINDS.get(destination, (None, None))[0], kind, + f'{destination.__name__} is not checked against {kind.__name__}') + + self.assertEqual(set(_DESTINATION_KINDS), set(expected), + '_DESTINATION_KINDS no longer lists the destinations cased here') + + if __name__ == '__main__': unittest.main()