fix: validate encoded telemetry message size

azure-iot-device/azure/iot/device/iothub/aio/async_clients.py

@@ -648,11 +648,11 @@ async def send_message_to_output(self, message: Union[Message, str], output_name
         if not isinstance(message, Message):
             message = Message(message)
 
+        message.output_name = output_name
+
         if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
             raise ValueError("Size of message can not exceed 256 KB.")
 
-        message.output_name = output_name
-
         logger.info("Sending message to output:" + output_name + "...")
         send_output_message_async = async_adapter.emulate_async(
             self._mqtt_pipeline.send_output_message

azure-iot-device/azure/iot/device/iothub/models/message.py

@@ -5,8 +5,64 @@
 # --------------------------------------------------------------------------
 """This module contains a class representing messages that are sent or received.
 """
+from datetime import date
+
 from azure.iot.device import constant
-import sys
+
+
+def _encode_message_data(data):
+    if isinstance(data, str):
+        return data.encode("utf-8")
+    if isinstance(data, (int, float)):
+        return str(data).encode("ascii")
+    if data is None:
+        return b""
+    if not isinstance(data, (bytes, bytearray)):
+        raise TypeError("Message data must be a string, bytes, bytearray, int, float, or None.")
+    return data
+
+
+def _get_system_properties(message):
+    properties = []
+    if message.output_name:
+        properties.append(("$.on", str(message.output_name)))
+    if message.message_id:
+        properties.append(("$.mid", str(message.message_id)))
+    if message.correlation_id:
+        properties.append(("$.cid", str(message.correlation_id)))
+    if message.user_id:
+        properties.append(("$.uid", str(message.user_id)))
+    if message.content_type:
+        properties.append(("$.ct", str(message.content_type)))
+    if message.content_encoding:
+        properties.append(("$.ce", str(message.content_encoding)))
+    if message.iothub_interface_id:
+        properties.append(("$.ifid", str(message.iothub_interface_id)))
+    if message.expiry_time_utc:
+        expiry_time = (
+            message.expiry_time_utc.isoformat()
+            if isinstance(message.expiry_time_utc, date)
+            else message.expiry_time_utc
+        )
+        properties.append(("$.exp", str(expiry_time)))
+    return properties
+
+
+def _get_custom_properties(message):
+    if not message.custom_properties:
+        return []
+
+    properties = [(str(key), str(value)) for key, value in message.custom_properties.items()]
+    properties.sort()
+
+    keys = [key for key, _ in properties]
+    if len(keys) != len(set(keys)):
+        raise ValueError("Duplicate keys in custom properties!")
+    return properties
+
+
+def _get_string_size(value):
+    return len(value.encode("utf-8"))
 
 
 class Message(object):
@@ -65,14 +121,22 @@ def __str__(self):
         return str(self.data)
 
     def get_size(self) -> int:
-        total = 0
-        total = total + sum(
-            sys.getsizeof(v)
-            for v in self.__dict__.values()
-            if v is not None and v is not self.custom_properties
+        """Return the message size in bytes as measured by IoT Hub.
+
+        The size is the encoded body plus system property values and application property names
+        and values. Strings are measured as UTF-8, matching the MQTT transport; bytes and
+        bytearrays are measured as-is. MQTT topic and packet overhead are not included.
+
+        :raises TypeError: If the message data is not a payload type supported by the MQTT
+            transport.
+        :raises ValueError: If custom property keys are duplicated after string conversion.
+        """
+        payload_size = len(_encode_message_data(self.data))
+        system_property_size = sum(
+            _get_string_size(value) for _, value in _get_system_properties(self)
+        )
+        application_property_size = sum(
+            _get_string_size(key) + _get_string_size(value)
+            for key, value in _get_custom_properties(self)
         )
-        if self.custom_properties:
-            total = total + sum(
-                sys.getsizeof(v) for v in self.custom_properties.values() if v is not None
-            )
-        return total
+        return payload_size + system_property_size + application_property_size

azure-iot-device/azure/iot/device/iothub/pipeline/mqtt_topic_iothub.py

@@ -5,9 +5,13 @@
 # --------------------------------------------------------------------------
 
 import logging
-from datetime import date
 import urllib
 
+from azure.iot.device.iothub.models.message import (
+    _get_custom_properties,
+    _get_system_properties,
+)
+
 logger = logging.getLogger(__name__)
 
 # NOTE: Whenever using standard URL encoding via the urllib.parse.quote() API
@@ -343,59 +347,17 @@ def encode_message_properties_in_topic(message_to_send, topic):
     "devices/<deviceId>/modules/<moduleId>/messages/events/
     :return: The topic which has been uri-encoded
     """
-    system_properties = []
-    if message_to_send.output_name:
-        system_properties.append(("$.on", str(message_to_send.output_name)))
-    if message_to_send.message_id:
-        system_properties.append(("$.mid", str(message_to_send.message_id)))
-
-    if message_to_send.correlation_id:
-        system_properties.append(("$.cid", str(message_to_send.correlation_id)))
-
-    if message_to_send.user_id:
-        system_properties.append(("$.uid", str(message_to_send.user_id)))
-
-    if message_to_send.content_type:
-        system_properties.append(("$.ct", str(message_to_send.content_type)))
-
-    if message_to_send.content_encoding:
-        system_properties.append(("$.ce", str(message_to_send.content_encoding)))
-
-    if message_to_send.iothub_interface_id:
-        system_properties.append(("$.ifid", str(message_to_send.iothub_interface_id)))
-
-    if message_to_send.expiry_time_utc:
-        system_properties.append(
-            (
-                "$.exp",
-                message_to_send.expiry_time_utc.isoformat()  # returns string
-                if isinstance(message_to_send.expiry_time_utc, date)
-                else message_to_send.expiry_time_utc,
-            )
-        )
-
+    system_properties = _get_system_properties(message_to_send)
     system_properties_encoded = urllib.parse.urlencode(
         system_properties, quote_via=urllib.parse.quote
     )
     topic += system_properties_encoded
 
-    if message_to_send.custom_properties and len(message_to_send.custom_properties) > 0:
+    custom_prop_seq = _get_custom_properties(message_to_send)
+    if custom_prop_seq:
         if system_properties and len(system_properties) > 0:
             topic += "&"
 
-        # Convert the custom properties to a sorted list in order to ensure the
-        # resulting ordering in the topic string is consistent across versions of Python.
-        # Convert to the properties to strings for safety.
-        custom_prop_seq = [
-            (str(i[0]), str(i[1])) for i in list(message_to_send.custom_properties.items())
-        ]
-        custom_prop_seq.sort()
-
-        # Validate that string conversion has not created duplicate keys
-        keys = [i[0] for i in custom_prop_seq]
-        if len(keys) != len(set(keys)):
-            raise ValueError("Duplicate keys in custom properties!")
-
         user_properties_encoded = urllib.parse.urlencode(
             custom_prop_seq, quote_via=urllib.parse.quote
         )

azure-iot-device/azure/iot/device/iothub/sync_clients.py

@@ -669,11 +669,11 @@ def send_message_to_output(self, message: Union[Message, str], output_name: str)
         if not isinstance(message, Message):
             message = Message(message)
 
+        message.output_name = output_name
+
         if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
             raise ValueError("Size of message can not exceed 256 KB.")
 
-        message.output_name = output_name
-
         logger.info("Sending message to output:" + output_name + "...")
 
         callback = EventedCallback()

tests/unit/iothub/aio/test_async_clients.py

@@ -9,7 +9,6 @@
 import asyncio
 import time
 import urllib
-import sys
 from azure.iot.device import exceptions as client_exceptions
 from azure.iot.device.common.auth import sastoken as st
 from azure.iot.device.iothub.aio import IoTHubDeviceClient, IoTHubModuleClient
@@ -659,10 +658,10 @@ def fail_send_message(message, callback):
         [
             pytest.param("message", id="String input"),
             pytest.param(222, id="Integer input"),
-            pytest.param(object(), id="Object input"),
+            pytest.param(1.5, id="Float input"),
+            pytest.param(b"message", id="Bytes input"),
+            pytest.param(bytearray(b"message"), id="Bytearray input"),
             pytest.param(None, id="None input"),
-            pytest.param([1, "str"], id="List input"),
-            pytest.param({"a": 2}, id="Dictionary input"),
         ],
     )
     async def test_wraps_data_in_message_and_calls_pipeline_send_message(
@@ -674,44 +673,27 @@ async def test_wraps_data_in_message_and_calls_pipeline_send_message(
         assert isinstance(sent_message, Message)
         assert sent_message.data == message_input
 
-    @pytest.mark.it("Raises error when message data size is greater than 256 KB")
-    async def test_raises_error_when_message_data_greater_than_256(self, client, mqtt_pipeline):
-        data_input = "serpensortia" * 256000
-        message = Message(data_input)
-        with pytest.raises(ValueError) as e_info:
-            await client.send_message(message)
-        assert "256 KB" in e_info.value.args[0]
-        assert mqtt_pipeline.send_message.call_count == 0
-
-    @pytest.mark.it("Raises error when message size is greater than 256 KB")
-    async def test_raises_error_when_message_size_greater_than_256(self, client, mqtt_pipeline):
-        data_input = "serpensortia"
-        message = Message(data_input)
-        message.custom_properties["spell"] = data_input * 256000
-        with pytest.raises(ValueError) as e_info:
-            await client.send_message(message)
-        assert "256 KB" in e_info.value.args[0]
-        assert mqtt_pipeline.send_message.call_count == 0
-
-    @pytest.mark.skipif(
-        sys.version_info >= (3, 12),
-        reason="Python 3.12 appears to have an issue. Investigate further.",
-    )
-    @pytest.mark.it("Does not raises error when message data size is equal to 256 KB")
-    async def test_raises_error_when_message_data_equal_to_256(self, client, mqtt_pipeline):
-        data_input = "a" * 262095
-        message = Message(data_input)
-        # This check was put as message class may undergo the default content type encoding change
-        # and the above calculation will change.
-        if message.get_size() != device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
-            assert False
-
-        await client.send_message(message)
+    @pytest.mark.it("Validates message size at the 256 KB boundary")
+    @pytest.mark.parametrize(
+        "size_delta, raises_error",
+        [
+            pytest.param(-1, False, id="Below limit"),
+            pytest.param(0, False, id="At limit"),
+            pytest.param(1, True, id="Above limit"),
+        ],
+    )
+    async def test_validates_message_size(self, client, mqtt_pipeline, size_delta, raises_error):
+        message_size = device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT + size_delta
+        message = Message(b"a" * message_size)
 
-        assert mqtt_pipeline.send_message.call_count == 1
-        sent_message = mqtt_pipeline.send_message.call_args[0][0]
-        assert isinstance(sent_message, Message)
-        assert sent_message.data == data_input
+        if raises_error:
+            with pytest.raises(ValueError, match="256 KB"):
+                await client.send_message(message)
+            assert mqtt_pipeline.send_message.call_count == 0
+        else:
+            await client.send_message(message)
+            assert mqtt_pipeline.send_message.call_count == 1
+            assert mqtt_pipeline.send_message.call_args[0][0] is message
 
 
 class SharedClientReceiveMethodRequestTests(object):
@@ -2061,10 +2043,10 @@ def fail_send_output_message(message, callback):
         [
             pytest.param("message", id="String input"),
             pytest.param(222, id="Integer input"),
-            pytest.param(object(), id="Object input"),
+            pytest.param(1.5, id="Float input"),
+            pytest.param(b"message", id="Bytes input"),
+            pytest.param(bytearray(b"message"), id="Bytearray input"),
             pytest.param(None, id="None input"),
-            pytest.param([1, "str"], id="List input"),
-            pytest.param({"a": 2}, id="Dictionary input"),
         ],
     )
     async def test_send_message_to_output_calls_pipeline_wraps_data_in_message(
@@ -2077,53 +2059,31 @@ async def test_send_message_to_output_calls_pipeline_wraps_data_in_message(
         assert isinstance(sent_message, Message)
         assert sent_message.data == message_input
 
-    @pytest.mark.it("Raises error when message data size is greater than 256 KB")
-    async def test_raises_error_when_message_to_output_data_greater_than_256(
-        self, client, mqtt_pipeline
-    ):
-        output_name = "some_output"
-        data_input = "serpensortia" * 256000
-        message = Message(data_input)
-        with pytest.raises(ValueError) as e_info:
-            await client.send_message_to_output(message, output_name)
-        assert "256 KB" in e_info.value.args[0]
-        assert mqtt_pipeline.send_output_message.call_count == 0
-
-    @pytest.mark.it("Raises error when message size is greater than 256 KB")
-    async def test_raises_error_when_message_to_output_size_greater_than_256(
-        self, client, mqtt_pipeline
-    ):
-        output_name = "some_output"
-        data_input = "serpensortia"
-        message = Message(data_input)
-        message.custom_properties["spell"] = data_input * 256000
-        with pytest.raises(ValueError) as e_info:
-            await client.send_message_to_output(message, output_name)
-        assert "256 KB" in e_info.value.args[0]
-        assert mqtt_pipeline.send_output_message.call_count == 0
-
-    @pytest.mark.skipif(
-        sys.version_info >= (3, 12),
-        reason="Python 3.12 appears to have an issue. Investigate further.",
+    @pytest.mark.it("Validates message size at the 256 KB boundary")
+    @pytest.mark.parametrize(
+        "size_delta, raises_error",
+        [
+            pytest.param(-1, False, id="Below limit"),
+            pytest.param(0, False, id="At limit"),
+            pytest.param(1, True, id="Above limit"),
+        ],
     )
-    @pytest.mark.it("Does not raises error when message data size is equal to 256 KB")
-    async def test_raises_error_when_message_to_output_data_equal_to_256(
-        self, client, mqtt_pipeline
-    ):
+    async def test_validates_message_size(self, client, mqtt_pipeline, size_delta, raises_error):
         output_name = "some_output"
-        data_input = "a" * 262095
-        message = Message(data_input)
-        # This check was put as message class may undergo the default content type encoding change
-        # and the above calculation will change.
-        if message.get_size() != device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
-            assert False
-
-        await client.send_message_to_output(message, output_name)
+        message_size = device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT + size_delta
+        output_name_size = len(output_name.encode("utf-8"))
+        message = Message(b"a" * (message_size - output_name_size))
+
+        if raises_error:
+            with pytest.raises(ValueError, match="256 KB"):
+                await client.send_message_to_output(message, output_name)
+            assert mqtt_pipeline.send_output_message.call_count == 0
+        else:
+            await client.send_message_to_output(message, output_name)
+            assert mqtt_pipeline.send_output_message.call_count == 1
+            assert mqtt_pipeline.send_output_message.call_args[0][0] is message
 
-        assert mqtt_pipeline.send_output_message.call_count == 1
-        sent_message = mqtt_pipeline.send_output_message.call_args[0][0]
-        assert isinstance(sent_message, Message)
-        assert sent_message.data == data_input
+        assert message.get_size() == message_size
 
 
 @pytest.mark.describe("IoTHubModuleClient (Asynchronous) - .receive_message_on_input()")