Skip to content

Commit 0ff7478

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.
1 parent 7b57bb0 commit 0ff7478

2 files changed

Lines changed: 151 additions & 2 deletions

File tree

‎pcapkit/foundation/registry/protocols.py‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,25 @@ def _register_name(protocol: 'Type[ProtocolBase]', *targets: 'tuple[type, Any]')
422422
Enum_LinkType: (Frame, PCAPNG),
423423
}
424424

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

426445
def _iter_code_targets(code: 'Any') -> 'Iterator[tuple[Type[ProtocolBase], Any]]':
427446
"""Flatten a ``code=`` argument into ``(destination, key)`` pairs.
@@ -527,10 +546,22 @@ def register_protocol_code(protocol: 'Type[ProtocolBase]', code: 'Any') -> 'None
527546
RegistryError: If a bare :class:`int` (or any other non-enum,
528547
non-mapping value) is given without an explicit destination, or if
529548
an enum member's type names no known destination registry --
530-
inference refuses rather than guesses.
549+
inference refuses rather than guesses; or if a key is not of the
550+
type its destination's registry is keyed by, as its typed registrar
551+
would reject it, e.g. ``{Frame: TransType.ICMP}``. Every target is
552+
checked before any is registered, so a rejected ``code`` leaves
553+
every registry as it was (:issue:`2268`).
531554
532555
"""
533-
for destination, key in _iter_code_targets(code):
556+
targets = list(_iter_code_targets(code))
557+
for destination, key in targets:
558+
owner = next((klass for klass in getattr(destination, '__mro__', ())
559+
if '__proto__' in vars(klass)), None)
560+
if owner in _DESTINATION_KINDS:
561+
kind, number = _DESTINATION_KINDS[owner]
562+
_check_code(key, kind, number=number)
563+
564+
for destination, key in targets:
534565
destination.register(key, protocol)
535566
logger.debug('registered %s into %s.__proto__: %s', protocol.__name__,
536567
destination.__name__, key)

‎tests/protocols/test_protocol_code_registration_unit.py‎

Lines changed: 118 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,118 @@ 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+
410528
if __name__ == '__main__':
411529
unittest.main()

0 commit comments

Comments
 (0)