Skip to content

Commit 79b88ee

Browse files
committed
fix(foundation): check every register_protocol_code key before registering any (#2268)
register_protocol_code (and the code= class keyword) handed each {destination: key} pair straight to destination.register, so a member of another enum, e.g. {Frame: TransType.ICMP}, replaced the shipped LinkType.ETHERNET entry; and it registered targets as it flattened them, so a bad later target left the earlier ones registered. Flatten every target first, check each key with _check_code against the enum of the class owning its destination's __proto__ (_DESTINATION_KINDS), then register. A destination keeping a registry of its own (Loopback, IPX) is not checked. Tests: RegisterProtocolCodeKeyCheckTests snapshot and restore the real registries; the five rejection cases fail on main. A table-driven test pins every _DESTINATION_KINDS row against the real registrars, killing all 14 single-cell mutants (dropping a row, or flipping its number flag).
1 parent 7b57bb0 commit 79b88ee

2 files changed

Lines changed: 230 additions & 4 deletions

File tree

‎pcapkit/foundation/registry/protocols.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
from pcapkit.const.reg.transtype import TransType
8989
from pcapkit.const.sctp.payload_protocol_identifier import \
9090
PayloadProtocolIdentifier as SCTP_PayloadProtocolIdentifier
91+
from pcapkit.corekit.enum import EnumRegistry
9192
from pcapkit.corekit.sentinels import NullType
9293
from pcapkit.protocols.application.httpv2 import FrameConstructor as HTTP_FrameConstructor
9394
from pcapkit.protocols.application.httpv2 import FrameParser as HTTP_FrameParser
@@ -415,13 +416,32 @@ def _register_name(protocol: 'Type[ProtocolBase]', *targets: 'tuple[type, Any]')
415416
#: table so :func:`register_protocol_code` can consult it. :class:`~pcapkit.const.reg.linktype.LinkType` naming
416417
#: more than one class is deliberate, not ambiguous -- :func:`register_linktype`
417418
#: already fans out to each of them.
418-
_CODE_DESTINATIONS: 'dict[type, tuple[Type[ProtocolBase], ...]]' = {
419+
_CODE_DESTINATIONS: 'dict[Type[EnumRegistry], tuple[Type[ProtocolBase], ...]]' = {
419420
Enum_EtherType: (Link,),
420421
Enum_TransType: (Internet,),
421422
Enum_PayloadProtocolIdentifier: (SCTP,),
422423
Enum_LinkType: (Frame, PCAPNG),
423424
}
424425

426+
#: Key type of each ``__proto__`` dispatch registry a ``code=`` target can name.
427+
#:
428+
#: Maps the class owning the registry to the enum its typed registrar checks
429+
#: ``code`` against, and whether a bare :class:`int` is a valid key there too --
430+
#: the same :func:`_check_code` arguments that registrar passes. A destination is
431+
#: looked up by the class it inherits ``__proto__`` from, so a subclass sharing
432+
#: :class:`Link`'s registry is checked as :class:`Link` is, while one keeping a
433+
#: registry of its own, such as :class:`~pcapkit.protocols.link.loopback.Loopback`,
434+
#: names no entry here and is not checked (:issue:`2268`).
435+
_DESTINATION_KINDS: 'dict[type, tuple[Type[EnumRegistry], bool]]' = {
436+
Link: (Enum_EtherType, False),
437+
Internet: (Enum_TransType, False),
438+
Frame: (Enum_LinkType, False),
439+
PCAPNG: (Enum_LinkType, False),
440+
TCP: (Enum_AppType, True),
441+
UDP: (Enum_AppType, True),
442+
SCTP: (Enum_PayloadProtocolIdentifier, True),
443+
}
444+
425445

426446
def _iter_code_targets(code: 'Any') -> 'Iterator[tuple[Type[ProtocolBase], Any]]':
427447
"""Flatten a ``code=`` argument into ``(destination, key)`` pairs.
@@ -444,7 +464,7 @@ def _iter_code_targets(code: 'Any') -> 'Iterator[tuple[Type[ProtocolBase], Any]]
444464
return
445465

446466
if isinstance(code, (enum.Enum, aenum.Enum)):
447-
destinations = _CODE_DESTINATIONS.get(type(code))
467+
destinations = _CODE_DESTINATIONS.get(cast('Type[EnumRegistry]', type(code)))
448468
if destinations is None:
449469
raise RegistryError(
450470
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
527547
RegistryError: If a bare :class:`int` (or any other non-enum,
528548
non-mapping value) is given without an explicit destination, or if
529549
an enum member's type names no known destination registry --
530-
inference refuses rather than guesses.
550+
inference refuses rather than guesses; or if a key is not of the
551+
type its destination's registry is keyed by, as its typed registrar
552+
would reject it, e.g. ``{Frame: TransType.ICMP}``. Every target is
553+
checked before any is registered, so a rejected ``code`` leaves
554+
every registry as it was (:issue:`2268`).
531555
532556
"""
533-
for destination, key in _iter_code_targets(code):
557+
targets = list(_iter_code_targets(code))
558+
for destination, key in targets:
559+
owner = next((klass for klass in getattr(destination, '__mro__', ())
560+
if '__proto__' in vars(klass)), None)
561+
if owner in _DESTINATION_KINDS:
562+
kind, number = _DESTINATION_KINDS[owner]
563+
_check_code(key, kind, number=number)
564+
565+
for destination, key in targets:
534566
destination.register(key, protocol)
535567
logger.debug('registered %s into %s.__proto__: %s', protocol.__name__,
536568
destination.__name__, key)

‎tests/protocols/test_protocol_code_registration_unit.py‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,15 @@
3131
import importlib.util
3232
import sys
3333
import unittest
34+
import warnings
35+
from typing import TYPE_CHECKING
3436
from unittest import mock
3537

3638
from tests._support import reimport_once_per_class
3739

40+
if TYPE_CHECKING:
41+
from typing import Any
42+
3843
#: Test tier (#1958): branches on the running interpreter's version or implementation.
3944
__tier__ = 'ALL'
4045

@@ -407,5 +412,194 @@ def test_backward_compatible_dispatch_tables_are_unaffected(self) -> None:
407412
('pcapkit.protocols.internet.ipv4', 'IPv4'))
408413

409414

415+
@unittest.skipUnless(HAS_RUNTIME, 'runtime dependencies not installed')
416+
class RegisterProtocolCodeKeyCheckTests(unittest.TestCase):
417+
"""Each ``code=`` key is checked against its destination's key type, all before any registers (#2268).
418+
419+
The typed registrars reject a member of another enum (#2224), but
420+
``register_protocol_code`` handed every ``{destination: key}`` pair straight
421+
to ``destination.register``. The registry enums compare equal to their
422+
values, so ``{Frame: TransType.ICMP}`` replaced the shipped entry at
423+
``LinkType.ETHERNET`` with only a warning. It also registered each target as
424+
it flattened them, so a bad later target left the earlier ones in place.
425+
426+
Each case snapshots every dispatch registry and the name registry, expects
427+
:exc:`~pcapkit.utilities.exceptions.RegistryError`, compares each registry
428+
with its snapshot, and restores the snapshot whatever happened.
429+
"""
430+
431+
def setUp(self) -> None:
432+
reimport_once_per_class(self)
433+
434+
def _assert_rejected_unchanged(self, call: 'Any') -> None:
435+
from pcapkit.protocols import __proto__ as protocol_registry
436+
from pcapkit.protocols.internet.internet import Internet
437+
from pcapkit.protocols.link.link import Link
438+
from pcapkit.protocols.misc.pcap.frame import Frame
439+
from pcapkit.protocols.misc.pcapng import PCAPNG
440+
from pcapkit.protocols.transport.sctp import SCTP
441+
from pcapkit.protocols.transport.tcp import TCP
442+
from pcapkit.protocols.transport.udp import UDP
443+
from pcapkit.utilities.exceptions import RegistryError
444+
445+
tables = {klass.__name__: klass.__proto__
446+
for klass in (Link, Internet, Frame, PCAPNG, TCP, UDP, SCTP)}
447+
tables['pcapkit.protocols'] = protocol_registry
448+
before = {name: dict(table) for name, table in tables.items()}
449+
try:
450+
with warnings.catch_warnings():
451+
warnings.simplefilter('ignore')
452+
with self.assertRaises(RegistryError):
453+
call()
454+
for name, table in tables.items():
455+
if name != 'pcapkit.protocols':
456+
self.assertEqual(dict(table), before[name], f'{name}.__proto__ changed')
457+
finally:
458+
for name, table in tables.items():
459+
table.clear()
460+
table.update(before[name])
461+
462+
def test_mapping_rejects_a_member_of_another_enum(self) -> None:
463+
"""The issue's example: ``TransType.ICMP`` equals ``LinkType.ETHERNET``."""
464+
from pcapkit.const.reg.linktype import LinkType
465+
from pcapkit.const.reg.transtype import TransType
466+
from pcapkit.foundation.registry.protocols import register_protocol_code
467+
from pcapkit.protocols.misc.pcap.frame import Frame
468+
from pcapkit.protocols.misc.raw import Raw
469+
470+
self.assertEqual(TransType.ICMP, LinkType.ETHERNET)
471+
self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {Frame: TransType.ICMP}))
472+
473+
def test_class_keyword_rejects_a_member_of_another_enum(self) -> None:
474+
from pcapkit.const.reg.transtype import TransType
475+
from pcapkit.protocols.misc.pcap.frame import Frame
476+
from pcapkit.protocols.misc.raw import Raw
477+
478+
def define() -> None:
479+
class Evil(Raw, code={Frame: TransType.ICMP}):
480+
pass
481+
482+
self._assert_rejected_unchanged(define)
483+
484+
def test_subclass_sharing_a_registry_is_checked_as_its_owner(self) -> None:
485+
"""``Ethernet`` inherits ``Link.__proto__``, so its keys are ``EtherType``s."""
486+
from pcapkit.const.reg.transtype import TransType
487+
from pcapkit.foundation.registry.protocols import register_protocol_code
488+
from pcapkit.protocols.link.ethernet import Ethernet
489+
from pcapkit.protocols.misc.raw import Raw
490+
491+
self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {Ethernet: TransType.ICMP}))
492+
493+
def test_number_destination_rejects_a_member_of_another_enum(self) -> None:
494+
from pcapkit.const.reg.linktype import LinkType
495+
from pcapkit.foundation.registry.protocols import register_protocol_code
496+
from pcapkit.protocols.misc.raw import Raw
497+
from pcapkit.protocols.transport.udp import UDP
498+
499+
self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, {UDP: LinkType(53)}))
500+
501+
def test_a_bad_later_target_leaves_the_earlier_ones_unregistered(self) -> None:
502+
"""The issue's list example, and one whose bad target fails the key check instead."""
503+
from pcapkit.const.reg.transtype import TransType
504+
from pcapkit.foundation.registry.protocols import register_protocol_code
505+
from pcapkit.protocols.internet.internet import Internet
506+
from pcapkit.protocols.misc.raw import Raw
507+
508+
for code in ([TransType.L2TP, 5], [TransType.L2TP, {Internet: 115}]):
509+
with self.subTest(code=code):
510+
self._assert_rejected_unchanged(lambda: register_protocol_code(Raw, code))
511+
512+
def test_number_and_own_registry_destinations_still_accept_an_int(self) -> None:
513+
"""A port is a valid ``TCP`` key; ``Loopback`` keeps an unchecked registry of its own."""
514+
from pcapkit.foundation.registry.protocols import register_protocol_code
515+
from pcapkit.protocols.link.loopback import Loopback
516+
from pcapkit.protocols.transport.tcp import TCP
517+
518+
class FakeProtocol:
519+
pass
520+
521+
with mock.patch.object(TCP, 'register') as reg_tcp, \
522+
mock.patch.object(Loopback, 'register') as reg_loopback:
523+
register_protocol_code(FakeProtocol, [{TCP: 54324}, {Loopback: 2}])
524+
reg_tcp.assert_called_once_with(54324, FakeProtocol)
525+
reg_loopback.assert_called_once_with(2, FakeProtocol)
526+
527+
def test_every_destination_kind_row_is_enforced(self) -> None:
528+
"""Every row of ``_DESTINATION_KINDS`` is exercised, both cells of it.
529+
530+
The table is what the check is made of, so a row dropped from it, or a
531+
``kind``/``number`` cell changed in it, silently stops checking a
532+
destination -- the registrars it stands in for are not called here, so
533+
nothing else would notice. Each row is pinned by behaviour rather than
534+
by comparing the table with a copy of itself:
535+
536+
* ``kind`` -- a member of *another* registry's enum must be rejected,
537+
and must leave every registry as it was. That is the #2224 confusion
538+
the check exists for, so a row that stops being checked, or one
539+
checked against the wrong enum, fails here.
540+
* ``number`` -- a bare :class:`int` is accepted where the registrar
541+
accepts one (a port, a payload protocol identifier) and rejected
542+
where it does not, which is the only observable difference the flag
543+
makes. The :class:`int` given to a rejecting row is the *value* of a
544+
member its registry already ships an entry at, so the check is pinned
545+
as being on the type and the entry at risk is a real one.
546+
547+
The ``kind`` identity is asserted as well, since a row retyped to an
548+
enum unrelated to every key used here would reject both keys for the
549+
wrong reason.
550+
"""
551+
from pcapkit.const.reg.apptype import AppType
552+
from pcapkit.const.reg.ethertype import EtherType
553+
from pcapkit.const.reg.linktype import LinkType
554+
from pcapkit.const.reg.transtype import TransType
555+
from pcapkit.const.sctp.payload_protocol_identifier import PayloadProtocolIdentifier
556+
from pcapkit.foundation.registry.protocols import (_DESTINATION_KINDS,
557+
register_protocol_code)
558+
from pcapkit.protocols.internet.internet import Internet
559+
from pcapkit.protocols.link.link import Link
560+
from pcapkit.protocols.misc.pcap.frame import Frame
561+
from pcapkit.protocols.misc.pcapng import PCAPNG
562+
from pcapkit.protocols.misc.raw import Raw
563+
from pcapkit.protocols.transport.sctp import SCTP
564+
from pcapkit.protocols.transport.tcp import TCP
565+
from pcapkit.protocols.transport.udp import UDP
566+
567+
# destination -> (kind, number, a member of another enum, a bare int)
568+
expected = {
569+
# 0x0800 is EtherType.Internet_Protocol_version_4, where Link ships IPv4.
570+
Link: (EtherType, False, TransType.ICMP, 0x0800),
571+
# 6 is TransType.TCP, where Internet ships TCP.
572+
Internet: (TransType, False, EtherType.XEROX_NS_IDP, 6),
573+
# The issue's own example: TransType.ICMP == LinkType.ETHERNET, where
574+
# both Frame and PCAPNG ship Ethernet.
575+
Frame: (LinkType, False, TransType.ICMP, 1),
576+
PCAPNG: (LinkType, False, TransType.ICMP, 1),
577+
TCP: (AppType, True, LinkType(53), 54324),
578+
UDP: (AppType, True, LinkType(53), 54324),
579+
SCTP: (PayloadProtocolIdentifier, True, TransType.ICMP, 54324),
580+
}
581+
582+
for destination, (kind, number, member, raw) in expected.items():
583+
with self.subTest(destination=destination.__name__):
584+
self.assertNotIsInstance(member, kind) # the test data, not the table
585+
586+
self._assert_rejected_unchanged(
587+
lambda dest=destination, key=member: register_protocol_code(Raw, {dest: key}))
588+
589+
if number:
590+
with mock.patch.object(destination, 'register') as register:
591+
register_protocol_code(Raw, {destination: raw})
592+
register.assert_called_once_with(raw, Raw)
593+
else:
594+
self._assert_rejected_unchanged(
595+
lambda dest=destination, key=raw: register_protocol_code(Raw, {dest: key}))
596+
597+
self.assertIs(_DESTINATION_KINDS.get(destination, (None, None))[0], kind,
598+
f'{destination.__name__} is not checked against {kind.__name__}')
599+
600+
self.assertEqual(set(_DESTINATION_KINDS), set(expected),
601+
'_DESTINATION_KINDS no longer lists the destinations cased here')
602+
603+
410604
if __name__ == '__main__':
411605
unittest.main()

0 commit comments

Comments
 (0)