Fine-tune can_respond of examine_headers and examine_body decorators (16e77cb4) · Commits · Kilter / Kilter Service

kilter/service/options.py

+40 −4

Original line number	Diff line number	Diff line
		@@ -11,6 +11,8 @@ Filter decorators for marking the requested protocol options and actions used
		from __future__ import annotations

		from collections import defaultdict
		from enum import Flag
		from enum import auto
		from typing import Callable
		from typing import Literal
		from typing import NamedTuple
		@@ -22,6 +24,7 @@ from kilter.protocol.messages import Stage
		from .session import Filter

		__all__ = [
		"CanRespond", "NEVER", "BEFORE", "DURING", "AFTER",
		"responds_to_connect", "examine_helo",
		"examine_sender", "examine_recipients",
		"examine_headers", "examine_body",
		@@ -48,6 +51,28 @@ DEFAULT_UNSET = \
		ProtocolFlags.NR_UNKNOWN


		class CanRespond(Flag):
		"""
		Flags for fine indication of which stages during message sending a filter may respond at

		Used with `examine_headers()` and `examine_body()` to further refine which stages during
		message header and body transfer will synchronously block until a response of some kind
		is received.
		"""

		NEVER = 0
		BEFORE = auto()
		DURING = auto()
		AFTER = auto()
		ALL = BEFORE\|DURING\|AFTER


		NEVER = CanRespond.NEVER
		BEFORE = CanRespond.BEFORE
		DURING = CanRespond.DURING
		AFTER = CanRespond.AFTER


		class FlagsTuple(NamedTuple):

		unset_options: ProtocolFlags = ProtocolFlags.NONE
		@@ -205,7 +230,7 @@ def examine_recipients(


		def examine_headers(
		can_respond: bool = False,
		can_respond: bool\|CanRespond = False,
		can_add: bool = False,
		can_modify: bool = False,
		leading_space: bool = False,
		@@ -226,8 +251,14 @@ def examine_headers(
		unset = ProtocolFlags.NO_HEADERS
		opts = ProtocolFlags.NONE
		acts = ActionFlags.NONE
		if can_respond:
		if isinstance(can_respond, bool):
		can_respond = CanRespond.ALL if can_respond else CanRespond.NEVER
		if CanRespond.BEFORE in can_respond:
		unset \|= ProtocolFlags.NO_DATA \| ProtocolFlags.NR_DATA
		if CanRespond.DURING in can_respond:
		unset \|= ProtocolFlags.NR_HEADER
		if CanRespond.AFTER in can_respond:
		unset \|= ProtocolFlags.NO_END_OF_HEADERS \| ProtocolFlags.NR_END_OF_HEADERS
		if can_add:
		acts \|= ActionFlags.ADD_HEADERS
		if can_modify:
		@@ -238,7 +269,7 @@ def examine_headers(


		def examine_body(
		can_respond: bool = False,
		can_respond: bool\|CanRespond = False,
		can_replace: bool = False,
		data_size: SIZES = ProtocolFlags.NONE,
		) -> Decorator:
		@@ -256,8 +287,13 @@ def examine_body(
		during negotiation and the filter will be disabled.
		"""
		unset = ProtocolFlags.NO_BODY
		if can_respond:
		if isinstance(can_respond, bool):
		can_respond = CanRespond.ALL if can_respond else CanRespond.NEVER
		if CanRespond.BEFORE in can_respond:
		unset \|= ProtocolFlags.NO_END_OF_HEADERS \| ProtocolFlags.NR_END_OF_HEADERS
		if CanRespond.DURING in can_respond:
		unset \|= ProtocolFlags.NR_BODY
		# CanRespond.AFTER is implicit
		return modify_flags(
		unset_options=unset, set_options=data_size,
		set_actions=ActionFlags.CHANGE_BODY if can_replace else ActionFlags.NONE,

tests/test_options.py

+278 −0

Original line number	Diff line number	Diff line
		@@ -6,6 +6,9 @@ from kilter.protocol import ProtocolFlags
		from kilter.protocol import Stage
		from kilter.service import Session
		from kilter.service import options
		from kilter.service.options import AFTER
		from kilter.service.options import BEFORE
		from kilter.service.options import DURING

		NO_CONNECT = ProtocolFlags.NO_CONNECT
		NO_HELO = ProtocolFlags.NO_HELO
		@@ -253,6 +256,9 @@ class Tests(TestCase):
		assert HEADER_LEADING_SPACE not in resolve_opts(flags, ProtocolFlags.NONE)
		assert HEADER_LEADING_SPACE in resolve_opts(flags, HEADER_LEADING_SPACE)

		assert NO_DATA in resolve_opts(flags, NO_DATA)
		assert NO_END_OF_HEADERS in resolve_opts(flags, NO_END_OF_HEADERS)

		assert ADD_HEADERS not in flags.set_actions
		assert CHANGE_HEADERS not in flags.set_actions

		@@ -292,6 +298,178 @@ class Tests(TestCase):
		assert ADD_HEADERS in flags.set_actions
		assert CHANGE_HEADERS in flags.set_actions

		def test_examine_headers_respond(self) -> None:
		"""
		Check that flags for skipping stages and responses are set according to 'can_respond'
		"""
		with self.subTest("False"):
		@options.examine_headers(can_respond=False)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# Flags must be left set if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA in resolved
		assert NR_DATA in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER in resolved
		assert NO_END_OF_HEADERS in resolved
		assert NR_END_OF_HEADERS in resolved

		with self.subTest("True"):
		@options.examine_headers(can_respond=True)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# Flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		with self.subTest("BEFORE"):
		@options.examine_headers(can_respond=BEFORE)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# NO_DATA and NR_DATA flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER in resolved
		assert NO_END_OF_HEADERS in resolved
		assert NR_END_OF_HEADERS in resolved

		with self.subTest("DURING"):
		@options.examine_headers(can_respond=DURING)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# NR_HEADER flag must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA in resolved
		assert NR_DATA in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS in resolved
		assert NR_END_OF_HEADERS in resolved

		with self.subTest("AFTER"):
		@options.examine_headers(can_respond=AFTER)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# NO_EOH and NR_EOH flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA in resolved
		assert NR_DATA in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		with self.subTest("BEFORE\|AFTER"):
		@options.examine_headers(can_respond=BEFORE\|AFTER)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER not in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		# NO_DATA, NR_DATA, NO_EOH and NR_EOH flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_DATA\|NR_DATA\|NO_HEADERS\|NR_HEADER\|NO_END_OF_HEADERS\|NR_END_OF_HEADERS,
		)
		assert NO_DATA not in resolved
		assert NR_DATA not in resolved
		assert NO_HEADERS not in resolved
		assert NR_HEADER in resolved
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved

		def test_examine_body(self) -> None:
		"""
		Check that examine_body sets protocol options and actions as appropriate
		@@ -330,6 +508,106 @@ class Tests(TestCase):

		assert CHANGE_BODY in flags.set_actions

		def test_examine_body_respond(self) -> None:
		"""
		Check that flags for skipping stages and responses are set according to 'can_respond'
		"""
		with self.subTest("False"):
		@options.examine_body(can_respond=False)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved

		# Flags must be left set if already set
		resolved = resolve_opts(
		flags,
		NO_END_OF_HEADERS\|NR_END_OF_HEADERS\|NO_BODY\|NR_BODY,
		)
		assert NO_END_OF_HEADERS in resolved
		assert NR_END_OF_HEADERS in resolved
		assert NO_BODY not in resolved
		assert NR_BODY in resolved

		with self.subTest("True"):
		@options.examine_body(can_respond=True)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved

		# Flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_END_OF_HEADERS\|NR_END_OF_HEADERS\|NO_BODY\|NR_BODY,
		)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved

		with self.subTest("BEFORE"):
		@options.examine_body(can_respond=BEFORE)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved

		# NO_EOH and NR_EOH flags must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_END_OF_HEADERS\|NR_END_OF_HEADERS\|NO_BODY\|NR_BODY,
		)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY in resolved

		with self.subTest("DURING"):
		@options.examine_body(can_respond=DURING)
		async def filtr(session: Session) -> Accept:
		return Accept()

		flags = options.get_flags(filtr)

		# Flags must not be set if not already set
		resolved = resolve_opts(flags)
		assert NO_END_OF_HEADERS not in resolved
		assert NR_END_OF_HEADERS not in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved

		# NR_BODY flag must be unset if already set
		resolved = resolve_opts(
		flags,
		NO_END_OF_HEADERS\|NR_END_OF_HEADERS\|NO_BODY\|NR_BODY,
		)
		assert NO_END_OF_HEADERS in resolved
		assert NR_END_OF_HEADERS in resolved
		assert NO_BODY not in resolved
		assert NR_BODY not in resolved


		class MacroTests(TestCase):
		"""