fix: validate encoded telemetry message size

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

@@ -343,6 +343,7 @@ async def send_message(self, message: Union[Message, str]) -> None:
             connected (and there is no auto-connect enabled)
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: TypeError if the message data type is not supported by the MQTT transport.
         :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
@@ -643,16 +644,17 @@ async def send_message_to_output(self, message: Union[Message, str], output_name
             connected (and there is no auto-connect enabled)
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: TypeError if the message data type is not supported by the MQTT transport.
         :raises: ValueError if the message fails size validation.
         """
         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

@@ -338,6 +338,7 @@ def send_message(self, message: Union[Message, str]) -> None:
             connected (and there is no auto-connect enabled)
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: TypeError if the message data type is not supported by the MQTT transport.
         :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
@@ -664,16 +665,17 @@ def send_message_to_output(self, message: Union[Message, str], output_name: str)
             connected (and there is no auto-connect enabled)
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: TypeError if the message data type is not supported by the MQTT transport.
         :raises: ValueError if the message fails size validation.
         """
         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/e2e/iothub_e2e/aio/test_send_message.py

@@ -6,7 +6,7 @@
 import logging
 import json
 import dev_utils
-from azure.iot.device.exceptions import OperationCancelled, ClientError
+from azure.iot.device.exceptions import OperationCancelled
 
 logger = logging.getLogger(__name__)
 logger.setLevel(level=logging.INFO)
@@ -53,9 +53,8 @@ async def test_bad_payload_raises(self, client, leak_tracker):
         def thing_that_cant_serialize():
             pass
 
-        with pytest.raises(ClientError) as e_info:
+        with pytest.raises(TypeError):
             await client.send_message(thing_that_cant_serialize)
-        assert isinstance(e_info.value.__cause__, TypeError)
 
         # TODO: investigate leak
         # leak_tracker.check_for_leaks()

tests/e2e/iothub_e2e/sync/test_sync_send_message.py

@@ -6,7 +6,7 @@
 import json
 import time
 import dev_utils
-from azure.iot.device.exceptions import OperationCancelled, ClientError
+from azure.iot.device.exceptions import OperationCancelled
 
 logger = logging.getLogger(__name__)
 logger.setLevel(level=logging.INFO)
@@ -50,9 +50,8 @@ def test_sync_bad_payload_raises(self, client, leak_tracker):
         def thing_that_cant_serialize():
             pass
 
-        with pytest.raises(ClientError) as e_info:
+        with pytest.raises(TypeError):
             client.send_message(thing_that_cant_serialize)
-        assert isinstance(e_info.value.__cause__, TypeError)
 
         # TODO; investigate this leak
         # leak_tracker.check_for_leaks()