fix(foundation): check every register_protocol_code key before registering any (#2268)

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)

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()