diff --git a/amqtt/mqtt5/_ack.py b/amqtt/mqtt5/_ack.py new file mode 100644 index 00000000..e0c868ff --- /dev/null +++ b/amqtt/mqtt5/_ack.py @@ -0,0 +1,208 @@ +"""Shared MQTT 5.0 PUBACK/PUBREC/PUBREL/PUBCOMP packet helpers.""" +from __future__ import annotations + +import struct +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar +from typing_extensions import Self + +from amqtt.codecs_amqtt import bytes_to_int, int_to_bytes, read_exact +from amqtt.errors import AMQTTError, CodecError, MQTTError, NoDataError +from amqtt.mqtt3.packet import MQTTFixedHeader, MQTTPacket, MQTTVariableHeader +from amqtt.mqtt5.properties import Properties +from amqtt.mqtt5.reason_codes import ReasonCode + +if TYPE_CHECKING: + from amqtt.adapters import ReaderAdapter + from amqtt.mqtt5.property_ids import PacketName + + +class AcknowledgementVariableHeader(MQTTVariableHeader): + """MQTT 5.0 acknowledgement variable header shared by PUBACK/PUBREC/PUBREL/PUBCOMP.""" + + PACKET_NAME: ClassVar[PacketName] + + __slots__ = ("packet_id", "properties", "reason_code") + + def __init__( + self, + packet_id: int, + reason_code: ReasonCode = ReasonCode.SUCCESS, + properties: Properties | None = None, + ) -> None: + super().__init__() + self.packet_id = packet_id + try: + self.reason_code = ReasonCode(reason_code) + except ValueError as exc: + msg = f"Unknown MQTT 5.0 reason code: {reason_code}" + raise MQTTError(msg) from exc + self.properties = Properties.for_packet(self.PACKET_NAME, properties) + self.validate() + + def validate(self) -> None: + """Validate packet identifier constraints.""" + if not 1 <= self.packet_id <= 0xFFFF: + msg = "[MQTT-2.2.1-3] QoS acknowledgement packets require a non-zero Packet Identifier." + raise MQTTError(msg) + + def to_bytes(self) -> bytes | bytearray: + """Encode packet id, optional reason code, and optional properties.""" + self.validate() + out = bytearray(int_to_bytes(self.packet_id, 2)) + if self.reason_code is ReasonCode.SUCCESS and self.properties.is_empty(): + return out + out.append(int(self.reason_code)) + if not self.properties.is_empty(): + out.extend(self.properties.encode()) + return out + + @classmethod + async def from_stream(cls: type[Self], reader: ReaderAdapter, fixed_header: MQTTFixedHeader) -> Self: + """Decode the acknowledgement variable header from a stream.""" + if fixed_header.remaining_length < 2: + msg = f"MQTT 5.0 {cls.PACKET_NAME} remaining length must include a Packet Identifier" + raise MQTTError(msg) + body = await read_exact(reader, fixed_header.remaining_length, f"{cls.PACKET_NAME} packet body") + return cls.from_bytes(body) + + @classmethod + def from_bytes(cls: type[Self], data: bytes | bytearray) -> Self: + """Decode the acknowledgement variable header from packet-body bytes.""" + if len(data) < 2: + msg = f"MQTT 5.0 {cls.PACKET_NAME} remaining length must include a Packet Identifier" + raise MQTTError(msg) + + packet_id = bytes_to_int(bytes(data[0:2])) + reason_code = ReasonCode.SUCCESS + properties = Properties(packet_name=cls.PACKET_NAME) + + if len(data) >= 3: + try: + reason_code = ReasonCode(data[2]) + except ValueError as exc: + msg = f"Unknown MQTT 5.0 reason code: {data[2]}" + raise MQTTError(msg) from exc + if len(data) >= 4: + properties = Properties.decode(data[3:], packet_name=cls.PACKET_NAME) + + return cls(packet_id, reason_code, properties) + + def __repr__(self) -> str: + """Return a developer-friendly representation.""" + return ( + f"{type(self).__name__}(packet_id={self.packet_id!r}, " + f"reason_code={self.reason_code!r}, properties={self.properties!r})" + ) + + +_AckVariableHeader = TypeVar("_AckVariableHeader", bound=AcknowledgementVariableHeader) + + +class AcknowledgementPacket(MQTTPacket[_AckVariableHeader, None, MQTTFixedHeader], Generic[_AckVariableHeader]): + """Base MQTT 5.0 acknowledgement packet.""" + + VARIABLE_HEADER: type[_AckVariableHeader] + PAYLOAD = None + PACKET_TYPE: ClassVar[int] + PACKET_NAME: ClassVar[PacketName] + EXPECTED_FLAGS: ClassVar[int] + + def __init__( + self, + fixed: MQTTFixedHeader | None = None, + variable_header: _AckVariableHeader | None = None, + payload: None = None, + ) -> None: + if fixed is None: + header = MQTTFixedHeader(self.PACKET_TYPE, self.EXPECTED_FLAGS) + else: + self._validate_fixed_header(fixed) + header = fixed + super().__init__(header, variable_header, payload) + + @classmethod + def _validate_fixed_header(cls, fixed_header: MQTTFixedHeader) -> None: + if fixed_header.packet_type != cls.PACKET_TYPE: + msg = f"Invalid fixed packet type {fixed_header.packet_type} for {cls.__name__} init" + raise AMQTTError(msg) from None + if fixed_header.flags != cls.EXPECTED_FLAGS: + msg = f"Invalid fixed header flags for MQTT 5.0 {cls.PACKET_NAME}" + raise MQTTError(msg) + + @classmethod + def build( + cls, + packet_id: int, + reason_code: ReasonCode = ReasonCode.SUCCESS, + properties: Properties | None = None, + ) -> Self: + """Build an outgoing MQTT 5.0 acknowledgement packet.""" + variable_header = cls.VARIABLE_HEADER(packet_id, reason_code, properties) # pylint: disable=not-callable + return cls(variable_header=variable_header) + + @classmethod + async def from_stream( + cls, + reader: ReaderAdapter, + fixed_header: MQTTFixedHeader | None = None, + variable_header: _AckVariableHeader | None = None, + ) -> Self: + """Decode an MQTT 5.0 acknowledgement packet from a stream.""" + if fixed_header is None: + try: + fixed_header = await cls.FIXED_HEADER.from_stream(reader) + except (CodecError, MQTTError, NoDataError, struct.error) as exc: + msg = f"Malformed MQTT 5.0 {cls.PACKET_NAME} fixed header" + raise MQTTError(msg) from exc + if fixed_header is None: + msg = f"No MQTT 5.0 {cls.PACKET_NAME} fixed header available" + raise MQTTError(msg) + cls._validate_fixed_header(fixed_header) + + if variable_header is None: + try: + variable_header = await cls.VARIABLE_HEADER.from_stream(reader, fixed_header) + except (CodecError, MQTTError, NoDataError, struct.error) as exc: + msg = f"Malformed MQTT 5.0 {cls.PACKET_NAME} packet" + raise MQTTError(msg) from exc + + return cls(fixed_header, variable_header) + + @property + def packet_id(self) -> int: + """Return the Packet Identifier.""" + if self.variable_header is None: + msg = "Variable header is not set" + raise ValueError(msg) + return self.variable_header.packet_id + + @packet_id.setter + def packet_id(self, val: int) -> None: + if self.variable_header is None: + msg = "Variable header is not set" + raise ValueError(msg) + self.variable_header.packet_id = val + self.variable_header.validate() + + @property + def reason_code(self) -> ReasonCode: + """Return the acknowledgement reason code.""" + if self.variable_header is None: + msg = "Variable header is not set" + raise ValueError(msg) + return self.variable_header.reason_code + + @reason_code.setter + def reason_code(self, val: ReasonCode) -> None: + if self.variable_header is None: + msg = "Variable header is not set" + raise ValueError(msg) + self.variable_header.reason_code = ReasonCode(val) + + @property + def properties(self) -> Properties: + """Return the acknowledgement properties.""" + if self.variable_header is None: + msg = "Variable header is not set" + raise ValueError(msg) + return self.variable_header.properties diff --git a/amqtt/mqtt5/properties.py b/amqtt/mqtt5/properties.py index 6e8e6c57..b215eaea 100644 --- a/amqtt/mqtt5/properties.py +++ b/amqtt/mqtt5/properties.py @@ -149,6 +149,10 @@ def has(self, identifier: int) -> bool: """Return whether a property identifier is present.""" return identifier in self._values + def is_empty(self) -> bool: + """Return whether this property set has no entries.""" + return not self._values + def encode(self) -> bytes: r"""Encode properties with the MQTT 5.0 length prefix. diff --git a/amqtt/mqtt5/puback.py b/amqtt/mqtt5/puback.py new file mode 100644 index 00000000..add518f3 --- /dev/null +++ b/amqtt/mqtt5/puback.py @@ -0,0 +1,21 @@ +"""MQTT 5.0 PUBACK packet (§3.4).""" +from __future__ import annotations + +from amqtt.mqtt3.packet import PUBACK +from amqtt.mqtt5._ack import AcknowledgementPacket, AcknowledgementVariableHeader +from amqtt.mqtt5.property_ids import PACKET_PUBACK + + +class PubackVariableHeader(AcknowledgementVariableHeader): + """MQTT 5.0 PUBACK variable header.""" + + PACKET_NAME = PACKET_PUBACK + + +class PubackPacket(AcknowledgementPacket[PubackVariableHeader]): + """MQTT 5.0 PUBACK packet.""" + + VARIABLE_HEADER = PubackVariableHeader # could be inferred, but explicitly set for readability + PACKET_TYPE = PUBACK + PACKET_NAME = PACKET_PUBACK + EXPECTED_FLAGS = 0x00 diff --git a/amqtt/mqtt5/pubcomp.py b/amqtt/mqtt5/pubcomp.py new file mode 100644 index 00000000..894cfe0b --- /dev/null +++ b/amqtt/mqtt5/pubcomp.py @@ -0,0 +1,21 @@ +"""MQTT 5.0 PUBCOMP packet (§3.7).""" +from __future__ import annotations + +from amqtt.mqtt3.packet import PUBCOMP +from amqtt.mqtt5._ack import AcknowledgementPacket, AcknowledgementVariableHeader +from amqtt.mqtt5.property_ids import PACKET_PUBCOMP + + +class PubcompVariableHeader(AcknowledgementVariableHeader): + """MQTT 5.0 PUBCOMP variable header.""" + + PACKET_NAME = PACKET_PUBCOMP + + +class PubcompPacket(AcknowledgementPacket[PubcompVariableHeader]): + """MQTT 5.0 PUBCOMP packet.""" + + VARIABLE_HEADER = PubcompVariableHeader # could be inferred, but explicitly set for readability + PACKET_TYPE = PUBCOMP + PACKET_NAME = PACKET_PUBCOMP + EXPECTED_FLAGS = 0x00 diff --git a/amqtt/mqtt5/pubrec.py b/amqtt/mqtt5/pubrec.py new file mode 100644 index 00000000..9abe18d3 --- /dev/null +++ b/amqtt/mqtt5/pubrec.py @@ -0,0 +1,21 @@ +"""MQTT 5.0 PUBREC packet (§3.5).""" +from __future__ import annotations + +from amqtt.mqtt3.packet import PUBREC +from amqtt.mqtt5._ack import AcknowledgementPacket, AcknowledgementVariableHeader +from amqtt.mqtt5.property_ids import PACKET_PUBREC + + +class PubrecVariableHeader(AcknowledgementVariableHeader): + """MQTT 5.0 PUBREC variable header.""" + + PACKET_NAME = PACKET_PUBREC + + +class PubrecPacket(AcknowledgementPacket[PubrecVariableHeader]): + """MQTT 5.0 PUBREC packet.""" + + VARIABLE_HEADER = PubrecVariableHeader # could be inferred, but explicitly set for readability + PACKET_TYPE = PUBREC + PACKET_NAME = PACKET_PUBREC + EXPECTED_FLAGS = 0x00 diff --git a/amqtt/mqtt5/pubrel.py b/amqtt/mqtt5/pubrel.py new file mode 100644 index 00000000..90a886ed --- /dev/null +++ b/amqtt/mqtt5/pubrel.py @@ -0,0 +1,21 @@ +"""MQTT 5.0 PUBREL packet (§3.6).""" +from __future__ import annotations + +from amqtt.mqtt3.packet import PUBREL +from amqtt.mqtt5._ack import AcknowledgementPacket, AcknowledgementVariableHeader +from amqtt.mqtt5.property_ids import PACKET_PUBREL + + +class PubrelVariableHeader(AcknowledgementVariableHeader): + """MQTT 5.0 PUBREL variable header.""" + + PACKET_NAME = PACKET_PUBREL + + +class PubrelPacket(AcknowledgementPacket[PubrelVariableHeader]): + """MQTT 5.0 PUBREL packet.""" + + VARIABLE_HEADER = PubrelVariableHeader # could be inferred, but explicitly set for readability + PACKET_TYPE = PUBREL + PACKET_NAME = PACKET_PUBREL + EXPECTED_FLAGS = 0x02 diff --git a/tests/mqtt5/test_properties.py b/tests/mqtt5/test_properties.py index ba5fd2b7..bbf800f5 100644 --- a/tests/mqtt5/test_properties.py +++ b/tests/mqtt5/test_properties.py @@ -120,16 +120,22 @@ def test_property_wire_bytes(identifier: int, value, expected_wire: bytes): def test_decode_spec_example_empty_properties(): - assert Properties(packet_name=PACKET_CONNECT).encode() == b"\x00" + properties = Properties(packet_name=PACKET_CONNECT) + + assert properties.is_empty() is True + assert properties.encode() == b"\x00" assert Properties.decode(b"\x00", packet_name=PACKET_CONNECT) == Properties(packet_name=PACKET_CONNECT) + properties.set(SESSION_EXPIRY_INTERVAL, 300) + assert properties.is_empty() is False + def test_properties_must_be_packet_scoped(): with pytest.raises(TypeError): - Properties() # type: ignore[call-arg] + Properties() # type: ignore[call-arg] # pylint: disable=missing-kwoa with pytest.raises(TypeError): - Properties.decode(b"\x00") # type: ignore[call-arg] + Properties.decode(b"\x00") # type: ignore[call-arg] # pylint: disable=missing-kwoa def test_duplicate_non_repeatable_property_raises(): diff --git a/tests/mqtt5/test_puback.py b/tests/mqtt5/test_puback.py new file mode 100644 index 00000000..8fd416e6 --- /dev/null +++ b/tests/mqtt5/test_puback.py @@ -0,0 +1,170 @@ +import asyncio + +from hypothesis import given, strategies as st +import pytest + +from amqtt.errors import AMQTTError, MQTTError +from amqtt.mqtt3.puback import PubackPacket as PubackV3Packet +from amqtt.mqtt3.pubcomp import PubcompPacket as PubcompV3Packet +from amqtt.mqtt3.pubrec import PubrecPacket as PubrecV3Packet +from amqtt.mqtt3.pubrel import PubrelPacket as PubrelV3Packet +from amqtt.mqtt5.properties import Properties +from amqtt.mqtt5.property_ids import ( + PACKET_CONNECT, + PACKET_PUBACK, + PACKET_PUBCOMP, + PACKET_PUBREC, + PACKET_PUBREL, + REASON_STRING, + SESSION_EXPIRY_INTERVAL, +) +from amqtt.mqtt5.puback import PubackPacket +from amqtt.mqtt5.pubcomp import PubcompPacket +from amqtt.mqtt5.pubrec import PubrecPacket +from amqtt.mqtt5.pubrel import PubrelPacket +from amqtt.mqtt5.reason_codes import ReasonCode + +ACK_PACKET_CASES = [ + pytest.param(PubackPacket, PACKET_PUBACK, b"\x40", id="PUBACK"), + pytest.param(PubrecPacket, PACKET_PUBREC, b"\x50", id="PUBREC"), + pytest.param(PubrelPacket, PACKET_PUBREL, b"\x62", id="PUBREL"), + pytest.param(PubcompPacket, PACKET_PUBCOMP, b"\x70", id="PUBCOMP"), +] + +V3_ACK_PACKET_CASES = [ + pytest.param(PubackV3Packet, b"\x40\x02\x00\x0a", id="PUBACK"), + pytest.param(PubrecV3Packet, b"\x50\x02\x00\x0a", id="PUBREC"), + pytest.param(PubrelV3Packet, b"\x62\x02\x00\x0a", id="PUBREL"), + pytest.param(PubcompV3Packet, b"\x70\x02\x00\x0a", id="PUBCOMP"), +] + + +@pytest.mark.parametrize(("packet_cls", "packet_name", "fixed_header"), ACK_PACKET_CASES) +def test_decode_spec_example_short_form(packet_cls, packet_name: str, fixed_header: bytes, make_reader) -> None: + data = fixed_header + b"\x02\x00\x0a" + + packet = asyncio.run(packet_cls.from_stream(make_reader(data))) + + assert packet.packet_id == 10 + assert packet.reason_code is ReasonCode.SUCCESS + assert packet.properties == Properties(packet_name=packet_name) + assert packet.to_bytes() == data + + +@pytest.mark.parametrize(("packet_cls", "packet_name", "fixed_header"), ACK_PACKET_CASES) +def test_short_form_round_trip(packet_cls, packet_name: str, fixed_header: bytes, make_reader) -> None: + expected = fixed_header + b"\x02\x00\x0a" + + packet = packet_cls.build(10) + decoded = asyncio.run(packet_cls.from_stream(make_reader(packet.to_bytes()))) + + assert packet.to_bytes() == expected + assert decoded.packet_id == 10 + assert decoded.reason_code is ReasonCode.SUCCESS + assert decoded.properties == Properties(packet_name=packet_name) + assert decoded.to_bytes() == expected + + +@pytest.mark.parametrize(("packet_cls", "packet_name", "fixed_header"), ACK_PACKET_CASES) +def test_full_form_round_trip(packet_cls, packet_name: str, fixed_header: bytes, make_reader) -> None: + properties = Properties(packet_name=packet_name) + properties.set(REASON_STRING, "nope") + expected = fixed_header + b"\x0b\x00\x0a\x92\x07\x1f\x00\x04nope" + + packet = packet_cls.build(10, ReasonCode.PACKET_IDENTIFIER_NOT_FOUND, properties) + decoded = asyncio.run(packet_cls.from_stream(make_reader(packet.to_bytes()))) + + assert packet.to_bytes() == expected + assert decoded.packet_id == 10 + assert decoded.reason_code is ReasonCode.PACKET_IDENTIFIER_NOT_FOUND + assert decoded.reason_code.is_error() is True + assert decoded.properties.get(REASON_STRING) == "nope" + assert decoded.properties == properties + assert decoded.to_bytes() == expected + + +@pytest.mark.parametrize(("packet_cls", "packet_name", "fixed_header"), ACK_PACKET_CASES) +def test_reason_code_without_properties_round_trip(packet_cls, packet_name: str, fixed_header: bytes, make_reader) -> None: + expected = fixed_header + b"\x03\x00\x0a\x10" + + packet = packet_cls.build(10, ReasonCode.NO_MATCHING_SUBSCRIBERS) + decoded = asyncio.run(packet_cls.from_stream(make_reader(packet.to_bytes()))) + + assert packet.to_bytes() == expected + assert decoded.packet_id == 10 + assert decoded.reason_code is ReasonCode.NO_MATCHING_SUBSCRIBERS + assert decoded.properties == Properties(packet_name=packet_name) + assert decoded.to_bytes() == expected + + +@pytest.mark.parametrize(("packet_cls", "packet_name", "fixed_header"), ACK_PACKET_CASES) +def test_parser_accepts_full_success_form(packet_cls, packet_name: str, fixed_header: bytes, make_reader) -> None: + data = fixed_header + b"\x04\x00\x0a\x00\x00" + + packet = asyncio.run(packet_cls.from_stream(make_reader(data))) + + assert packet.packet_id == 10 + assert packet.reason_code is ReasonCode.SUCCESS + assert packet.properties == Properties(packet_name=packet_name) + assert packet.to_bytes() == fixed_header + b"\x02\x00\x0a" + + +@pytest.mark.parametrize(("packet_cls", "data"), V3_ACK_PACKET_CASES) +def test_mqtt3_ack_parsing_is_unchanged(packet_cls, data: bytes, make_reader) -> None: + packet = asyncio.run(packet_cls.from_stream(make_reader(data))) + + assert packet.packet_id == 10 + assert packet.to_bytes() == data + + +@pytest.mark.parametrize( + ("packet_cls", "data"), + [ + pytest.param(PubackPacket, b"\x40\x01\x00", id="short-body"), + pytest.param(PubackPacket, b"\x40\x02\x00\x00", id="zero-packet-id"), + pytest.param(PubackPacket, b"\x40\x03\x00\x01\x03", id="unknown-reason-code"), + pytest.param(PubackPacket, b"\x40\x05\x00\x01\x80\x05\x1f", id="malformed-properties"), + pytest.param(PubrelPacket, b"\x60\x02\x00\x01", id="pubrel-invalid-flags"), + ], +) +def test_malformed_ack_input_raises(packet_cls, data: bytes, make_reader) -> None: + with pytest.raises(MQTTError): + asyncio.run(packet_cls.from_stream(make_reader(data))) + + +def test_incorrect_fixed_header_raises() -> None: + with pytest.raises(AMQTTError): + PubackPacket(fixed=PubrecPacket.build(10).fixed_header) + + +def test_build_rejects_non_ack_properties() -> None: + properties = Properties(packet_name=PACKET_CONNECT) + properties.set(SESSION_EXPIRY_INTERVAL, 60) + + with pytest.raises(MQTTError): + PubackPacket.build(10, properties=properties) + + +@pytest.mark.parametrize("prop", ["packet_id", "reason_code", "properties"]) +def test_empty_variable_header(prop: str) -> None: + packet = PubackPacket() + + with pytest.raises(ValueError): + assert getattr(packet, prop) is not None + + +@pytest.mark.parametrize("prop", ["packet_id", "reason_code"]) +def test_empty_variable_header_setter(prop: str) -> None: + packet = PubackPacket() + + with pytest.raises(ValueError): + setattr(packet, prop, 10) + + +@pytest.mark.parametrize(("packet_cls", "_packet_name", "_fixed_header"), ACK_PACKET_CASES) +@given(data=st.binary()) +def test_ack_decode_never_crashes(packet_cls, _packet_name: str, _fixed_header: bytes, make_reader, data: bytes) -> None: + try: + asyncio.run(packet_cls.from_stream(make_reader(data))) + except (AMQTTError, MQTTError): + pass