fix(azure-iot-device): Message size checks (#447)

fix(azure-iot-device): Message size checks

azure-iot-device/azure/iot/device/constant.py

@@ -11,3 +11,4 @@
 IOTHUB_API_VERSION = "2018-06-30"
 PROVISIONING_API_VERSION = "2019-03-31"
 SECURITY_MESSAGE_INTERFACE_ID = "urn:azureiot:Security:SecurityAgent:1"
+TELEMETRY_MESSAGE_SIZE_LIMIT = 262144

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

@@ -20,6 +20,7 @@
 from azure.iot.device import exceptions
 from azure.iot.device.iothub.inbox_manager import InboxManager
 from .async_inbox import AsyncClientInbox
+from azure.iot.device import constant as device_constant
 
 logger = logging.getLogger(__name__)
 
@@ -133,10 +134,14 @@ async def send_message(self, message):
             during execution.
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
             message = Message(message)
 
+        if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
+            raise ValueError("Size of telemetry message can not exceed 256 KB.")
+
         logger.info("Sending message to Hub...")
         send_message_async = async_adapter.emulate_async(self._iothub_pipeline.send_message)
 
@@ -413,10 +418,14 @@ async def send_message_to_output(self, message, output_name):
             during execution.
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
             message = Message(message)
 
+        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 + "...")

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

@@ -6,6 +6,7 @@
 """This module contains a class representing messages that are sent or received.
 """
 from azure.iot.device import constant
+import sys
 
 
 # TODO: Revise this class. Does all of this REALLY need to be here?
@@ -75,3 +76,16 @@ def set_as_security_message(self):
 
     def __str__(self):
         return str(self.data)
+
+    def get_size(self):
+        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
+        )
+        if self.custom_properties:
+            total = total + sum(
+                sys.getsizeof(v) for v in self.custom_properties.values() if v is not None
+            )
+        return total

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

@@ -21,6 +21,8 @@
 from azure.iot.device import exceptions
 from azure.iot.device.common.evented_callback import EventedCallback
 from azure.iot.device.common.callable_weak_method import CallableWeakMethod
+from azure.iot.device import constant as device_constant
+
 
 logger = logging.getLogger(__name__)
 
@@ -145,10 +147,14 @@ def send_message(self, message):
             during execution.
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
             message = Message(message)
 
+        if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
+            raise ValueError("Size of telemetry message can not exceed 256 KB.")
+
         logger.info("Sending message to Hub...")
 
         callback = EventedCallback()
@@ -451,9 +457,14 @@ def send_message_to_output(self, message, output_name):
             during execution.
         :raises: :class:`azure.iot.device.exceptions.ClientError` if there is an unexpected failure
             during execution.
+        :raises: ValueError if the message fails size validation.
         """
         if not isinstance(message, Message):
             message = Message(message)
+
+        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 + "...")

azure-iot-device/tests/iothub/aio/test_async_clients.py

@@ -19,6 +19,8 @@
 from azure.iot.device.iothub.aio.async_inbox import AsyncClientInbox
 from azure.iot.device.common import async_adapter
 from azure.iot.device.iothub.auth import IoTEdgeError
+import sys
+from azure.iot.device import constant as device_constant
 
 pytestmark = pytest.mark.asyncio
 logging.basicConfig(level=logging.DEBUG)
@@ -496,6 +498,41 @@ 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, iothub_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 iothub_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, iothub_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 iothub_pipeline.send_message.call_count == 0
+
+    @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, iothub_pipeline):
+        data_input = "a" * 261976
+        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)
+
+        assert iothub_pipeline.send_message.call_count == 1
+        sent_message = iothub_pipeline.send_message.call_args[0][0]
+        assert isinstance(sent_message, Message)
+        assert sent_message.data == data_input
+
 
 class SharedClientReceiveMethodRequestTests(object):
     @pytest.mark.it("Implicitly enables methods feature if not already enabled")
@@ -2266,6 +2303,50 @@ 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, iothub_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 iothub_pipeline.send_output_event.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, iothub_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 iothub_pipeline.send_output_event.call_count == 0
+
+    @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, iothub_pipeline
+    ):
+        output_name = "some_output"
+        data_input = "a" * 261976
+        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)
+
+        assert iothub_pipeline.send_output_event.call_count == 1
+        sent_message = iothub_pipeline.send_output_event.call_args[0][0]
+        assert isinstance(sent_message, Message)
+        assert sent_message.data == data_input
+
 
 @pytest.mark.describe("IoTHubModuleClient (Asynchronous) - .receive_message_on_input()")
 class TestIoTHubModuleClientReceiveInputMessage(IoTHubModuleClientTestsConfig):

azure-iot-device/tests/iothub/test_sync_clients.py

@@ -18,6 +18,7 @@
 from azure.iot.device.iothub.models import Message, MethodRequest
 from azure.iot.device.iothub.sync_inbox import SyncClientInbox
 from azure.iot.device.iothub.auth import IoTEdgeError
+from azure.iot.device import constant as device_constant
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -492,6 +493,42 @@ 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")
+    def test_raises_error_when_message_data_greater_than_256(self, client, iothub_pipeline):
+        data_input = "serpensortia" * 25600
+        message = Message(data_input)
+        with pytest.raises(ValueError) as e_info:
+            client.send_message(message)
+        assert "256 KB" in e_info.value.args[0]
+        assert iothub_pipeline.send_message.call_count == 0
+
+    @pytest.mark.it("Raises error when message size is greater than 256 KB")
+    def test_raises_error_when_message_size_greater_than_256(self, client, iothub_pipeline):
+        data_input = "serpensortia"
+        message = Message(data_input)
+        message.custom_properties["spell"] = data_input * 25600
+        with pytest.raises(ValueError) as e_info:
+            client.send_message(message)
+        assert "256 KB" in e_info.value.args[0]
+        assert iothub_pipeline.send_message.call_count == 0
+
+    @pytest.mark.it("Does not raises error when message data size is equal to 256 KB")
+    def test_raises_error_when_message_data_equal_to_256(self, client, iothub_pipeline):
+        data_input = "a" * 261976
+        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.
+        # Had to do greater than check for python 2. Ideally should be not equal check
+        if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
+            assert False
+
+        client.send_message(message)
+
+        assert iothub_pipeline.send_message.call_count == 1
+        sent_message = iothub_pipeline.send_message.call_args[0][0]
+        assert isinstance(sent_message, Message)
+        assert sent_message.data == data_input
+
 
 class SharedClientReceiveMethodRequestTests(object):
     @pytest.mark.it("Implicitly enables methods feature if not already enabled")
@@ -2462,6 +2499,49 @@ 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")
+    def test_raises_error_when_message_to_output_data_greater_than_256(
+        self, client, iothub_pipeline
+    ):
+        output_name = "some_output"
+        data_input = "serpensortia" * 256000
+        message = Message(data_input)
+        with pytest.raises(ValueError) as e_info:
+            client.send_message_to_output(message, output_name)
+        assert "256 KB" in e_info.value.args[0]
+        assert iothub_pipeline.send_output_event.call_count == 0
+
+    @pytest.mark.it("Raises error when message size is greater than 256 KB")
+    def test_raises_error_when_message_to_output_size_greater_than_256(
+        self, client, iothub_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:
+            client.send_message_to_output(message, output_name)
+        assert "256 KB" in e_info.value.args[0]
+        assert iothub_pipeline.send_output_event.call_count == 0
+
+    @pytest.mark.it("Does not raises error when message data size is equal to 256 KB")
+    def test_raises_error_when_message_to_output_data_equal_to_256(self, client, iothub_pipeline):
+        output_name = "some_output"
+        data_input = "a" * 261976
+        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.
+        # Had to do greater than check for python 2. Ideally should be not equal check
+        if message.get_size() > device_constant.TELEMETRY_MESSAGE_SIZE_LIMIT:
+            assert False
+
+        client.send_message_to_output(message, output_name)
+
+        assert iothub_pipeline.send_output_event.call_count == 1
+        sent_message = iothub_pipeline.send_output_event.call_args[0][0]
+        assert isinstance(sent_message, Message)
+        assert sent_message.data == data_input
+
 
 @pytest.mark.describe("IoTHubModuleClient (Synchronous) - .receive_message_on_input()")
 class TestIoTHubModuleClientReceiveInputMessage(IoTHubModuleClientTestsConfig):