diff --git a/lib/python/base_cli/command_protocol.py b/lib/python/base_cli/command_protocol.py index 8521769..cc24467 100644 --- a/lib/python/base_cli/command_protocol.py +++ b/lib/python/base_cli/command_protocol.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +import unicodedata from collections.abc import Mapping from dataclasses import dataclass @@ -25,6 +26,14 @@ MAX_RECORD_COUNT = 1_000_000 +def _validate_protocol_header(protocol_header: str) -> None: + """Reject headers that can alter the line-delimited wire framing.""" + if not isinstance(protocol_header, str) or not protocol_header: + raise CommandProtocolError("protocol_header must be a non-empty framing-safe string") + if any(unicodedata.category(character) in {"Cc", "Cf"} for character in protocol_header): + raise CommandProtocolError("protocol_header must be a framing-safe single line") + + class CommandProtocolError(ValueError): """Raised when a command-protocol schema or payload violates its contract.""" @@ -159,6 +168,7 @@ def dumps_records( registry: CommandSchemaRegistry | None = None, ) -> str: """Serialize a sequence of typed command records using the protocol framing.""" + _validate_protocol_header(protocol_header) active_registry = registry or DEFAULT_SCHEMA_REGISTRY schema = active_registry.schema(record_type) if len(records) > MAX_RECORD_COUNT: @@ -187,6 +197,7 @@ def loads_records( registry: CommandSchemaRegistry | None = None, ) -> tuple[str, tuple[dict[str, RecordValue], ...]]: """Validate and decode protocol-framed command records.""" + _validate_protocol_header(protocol_header) active_registry = registry or DEFAULT_SCHEMA_REGISTRY # The wire framing is LF-delimited. `str.splitlines()` also accepts CR, # vertical tab, form feed, and Unicode separators, which would make the diff --git a/tests/test_command_protocol.py b/tests/test_command_protocol.py index 5023c2d..fecb791 100644 --- a/tests/test_command_protocol.py +++ b/tests/test_command_protocol.py @@ -138,6 +138,29 @@ def test_consumer_can_preserve_a_legacy_wire_header(self) -> None: ) self.assertEqual(decoded, (generic_record(),)) + def test_custom_framing_safe_header_round_trips(self) -> None: + header = "CUSTOM_COMMAND_PROTOCOL_V1" + payload = dumps_record(RECORD_TYPE, generic_record(), protocol_header=header) + + self.assertEqual( + loads_records(payload, expected_record_type=RECORD_TYPE, protocol_header=header), + (RECORD_TYPE, (generic_record(),)), + ) + + def test_rejects_empty_and_control_bearing_protocol_headers(self) -> None: + payload = dumps_record(RECORD_TYPE, generic_record()) + for header in ("", "BAD\nHEADER", "BAD\rHEADER", "BAD\0HEADER", "BAD\u202eHEADER"): + with self.subTest(header=repr(header)): + with self.assertRaisesRegex(CommandProtocolError, "framing-safe") as dump_error: + dumps_record(RECORD_TYPE, generic_record(), protocol_header=header) + with self.assertRaisesRegex(CommandProtocolError, "framing-safe") as load_error: + loads_records(payload, protocol_header=header) + for error in (dump_error.exception, load_error.exception): + self.assertNotIn("\n", str(error)) + self.assertNotIn("\r", str(error)) + self.assertNotIn("\0", str(error)) + self.assertNotIn("\u202e", str(error)) + def test_protocol_has_stable_generic_version_and_explicit_field_names(self) -> None: payload = dumps_record(RECORD_TYPE, generic_record())