Limit to_device EDU size to 65536

Author: MatMaul

synapse/api/constants.py

@@ -28,6 +28,7 @@
 
 # the max size of a (canonical-json-encoded) event
 MAX_PDU_SIZE = 65536
+MAX_EDU_SIZE = 65536
 
 # Max/min size of ints in canonical JSON
 CANONICALJSON_MAX_INT = (2**53) - 1

synapse/handlers/devicemessage.py

@@ -20,11 +20,19 @@
 #
 
 import logging
+from copy import deepcopy
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
-from synapse.api.constants import EduTypes, EventContentFields, ToDeviceEventTypes
-from synapse.api.errors import Codes, SynapseError
+from canonicaljson import encode_canonical_json
+
+from synapse.api.constants import (
+    MAX_EDU_SIZE,
+    EduTypes,
+    EventContentFields,
+    ToDeviceEventTypes,
+)
+from synapse.api.errors import Codes, EventSizeError, SynapseError
 from synapse.api.ratelimiting import Ratelimiter
 from synapse.logging.context import run_in_background
 from synapse.logging.opentracing import (
@@ -293,18 +301,18 @@ async def send_device_message(
 
         remote_edu_contents = {}
         for destination, messages in remote_messages.items():
-            # The EDU contains a "message_id" property which is used for
-            # idempotence. Make up a random one.
-            message_id = random_string(16)
-            log_kv({"destination": destination, "message_id": message_id})
-
-            remote_edu_contents[destination] = {
-                "messages": messages,
-                "sender": sender_user_id,
-                "type": message_type,
-                "message_id": message_id,
-                "org.matrix.opentracing_context": json_encoder.encode(context),
-            }
+            edu_contents = get_device_message_edu_contents(
+                sender_user_id, message_type, messages, context
+            )
+            remote_edu_contents[destination] = edu_contents
+            log_kv(
+                {
+                    "destination": destination,
+                    "message_ids": [
+                        edu_content["message_id"] for edu_content in edu_contents
+                    ],
+                }
+            )
 
         # Add messages to the database.
         # Retrieve the stream id of the last-processed to-device message.
@@ -409,3 +417,63 @@ async def get_events_for_dehydrated_device(
             "events": messages,
             "next_batch": f"d{stream_id}",
         }
+
+
+def get_device_message_edu_contents(
+    sender_user_id: str,
+    message_type: str,
+    messages: Dict[str, Dict[str, JsonDict]],
+    context: Dict[str, Any],
+) -> List[JsonDict]:
+    """
+    This function takes a dictionary of messages and splits them into several EDUs if needed.
+
+    It will raise an EventSizeError if a single message is too large to fit into an EDU.
+    """
+
+    base_edu_content = {
+        "messages": {},
+        "sender": sender_user_id,
+        "type": message_type,
+        "message_id": random_string(16),
+    }
+    # This is the size of the full EDU without any messages and without the opentracing context
+    base_edu_size = len(
+        encode_canonical_json(
+            {
+                "edu_type": "m.direct_to_device",
+                "content": base_edu_content,
+            }
+        )
+    )
+    base_edu_content["org.matrix.opentracing_context"] = json_encoder.encode(context)
+
+    edu_contents = []
+
+    current_edu_content: JsonDict = deepcopy(base_edu_content)
+    current_edu_size = base_edu_size
+
+    for recipient, message in messages.items():
+        # We remove 2 for the curly braces and add 2 for the colon and comma
+        # We may overshoot by 1 for single message EDUs because of the comma, but that's fine
+        message_entry_size = len(encode_canonical_json({recipient: message}))
+
+        if current_edu_size + message_entry_size > MAX_EDU_SIZE:
+            if len(current_edu_content["messages"]) == 0:
+                raise EventSizeError("device message too large", unpersistable=True)
+
+            edu_contents.append(current_edu_content)
+
+            current_edu_content = deepcopy(base_edu_content)
+            current_edu_content["message_id"] = random_string(16)
+
+            current_edu_size = base_edu_size
+        else:
+            current_edu_size += message_entry_size
+
+        current_edu_content["messages"][recipient] = message
+
+    if len(current_edu_content["messages"]) > 0:
+        edu_contents.append(current_edu_content)
+
+    return edu_contents

synapse/storage/databases/main/deviceinbox.py

@@ -718,15 +718,15 @@ def get_all_new_device_messages_txn(
     async def add_messages_to_device_inbox(
         self,
         local_messages_by_user_then_device: Dict[str, Dict[str, JsonDict]],
-        remote_messages_by_destination: Dict[str, JsonDict],
+        remote_edu_contents: Dict[str, List[JsonDict]],
     ) -> int:
         """Used to send messages from this server.
 
         Args:
             local_messages_by_user_then_device:
                 Dictionary of recipient user_id to recipient device_id to message.
-            remote_messages_by_destination:
-                Dictionary of destination server_name to the EDU JSON to send.
+            remote_edu_contents:
+                Dictionary of destination server_name to the list of EDU contents to send.
 
         Returns:
             The new stream_id.
@@ -760,42 +760,53 @@ def add_messages_txn(
                         destination,
                         stream_id,
                         now_ms,
-                        json_encoder.encode(edu),
+                        json_encoder.encode(edu_content),
                         self._instance_name,
                     )
-                    for destination, edu in remote_messages_by_destination.items()
+                    for destination, edu_contents in remote_edu_contents.items()
+                    for edu_content in edu_contents
                 ],
             )
 
-            for destination, edu in remote_messages_by_destination.items():
-                if issue9533_logger.isEnabledFor(logging.DEBUG):
-                    issue9533_logger.debug(
-                        "Queued outgoing to-device messages with "
-                        "stream_id %i, EDU message_id %s, type %s for %s: %s",
-                        stream_id,
-                        edu["message_id"],
-                        edu["type"],
-                        destination,
-                        [
-                            f"{user_id}/{device_id} (msgid "
-                            f"{msg.get(EventContentFields.TO_DEVICE_MSGID)})"
-                            for (user_id, messages_by_device) in edu["messages"].items()
-                            for (device_id, msg) in messages_by_device.items()
-                        ],
-                    )
+            for destination, edu_contents in remote_edu_contents.items():
+                for edu_content in edu_contents:
+                    if issue9533_logger.isEnabledFor(logging.DEBUG):
+                        issue9533_logger.debug(
+                            "Queued outgoing to-device messages with "
+                            "stream_id %i, EDU message_id %s, type %s for %s: %s",
+                            stream_id,
+                            edu_content["message_id"],
+                            edu_content["type"],
+                            destination,
+                            [
+                                f"{user_id}/{device_id} (msgid "
+                                f"{msg.get(EventContentFields.TO_DEVICE_MSGID)})"
+                                for (user_id, messages_by_device) in edu_content[
+                                    "messages"
+                                ].items()
+                                for (device_id, msg) in messages_by_device.items()
+                            ],
+                        )
 
-                for user_id, messages_by_device in edu["messages"].items():
-                    for device_id, msg in messages_by_device.items():
-                        with start_active_span("store_outgoing_to_device_message"):
-                            set_tag(SynapseTags.TO_DEVICE_EDU_ID, edu["sender"])
-                            set_tag(SynapseTags.TO_DEVICE_EDU_ID, edu["message_id"])
-                            set_tag(SynapseTags.TO_DEVICE_TYPE, edu["type"])
-                            set_tag(SynapseTags.TO_DEVICE_RECIPIENT, user_id)
-                            set_tag(SynapseTags.TO_DEVICE_RECIPIENT_DEVICE, device_id)
-                            set_tag(
-                                SynapseTags.TO_DEVICE_MSGID,
-                                msg.get(EventContentFields.TO_DEVICE_MSGID),
-                            )
+                    for user_id, messages_by_device in edu_content["messages"].items():
+                        for device_id, msg in messages_by_device.items():
+                            with start_active_span("store_outgoing_to_device_message"):
+                                set_tag(
+                                    SynapseTags.TO_DEVICE_EDU_ID, edu_content["sender"]
+                                )
+                                set_tag(
+                                    SynapseTags.TO_DEVICE_EDU_ID,
+                                    edu_content["message_id"],
+                                )
+                                set_tag(SynapseTags.TO_DEVICE_TYPE, edu_content["type"])
+                                set_tag(SynapseTags.TO_DEVICE_RECIPIENT, user_id)
+                                set_tag(
+                                    SynapseTags.TO_DEVICE_RECIPIENT_DEVICE, device_id
+                                )
+                                set_tag(
+                                    SynapseTags.TO_DEVICE_MSGID,
+                                    msg.get(EventContentFields.TO_DEVICE_MSGID),
+                                )
 
         async with self._to_device_msg_id_gen.get_next() as stream_id:
             now_ms = self._clock.time_msec()
@@ -804,7 +815,7 @@ def add_messages_txn(
             )
             for user_id in local_messages_by_user_then_device.keys():
                 self._device_inbox_stream_cache.entity_has_changed(user_id, stream_id)
-            for destination in remote_messages_by_destination.keys():
+            for destination in remote_edu_contents.keys():
                 self._device_federation_outbox_stream_cache.entity_has_changed(
                     destination, stream_id
                 )

tests/rest/client/test_sendtodevice.py

@@ -18,12 +18,20 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
+from unittest.mock import AsyncMock, Mock
+
 from parameterized import parameterized_class
 
-from synapse.api.constants import EduTypes
+from twisted.test.proto_helpers import MemoryReactor
+
+from synapse.api.constants import MAX_EDU_SIZE, EduTypes
+from synapse.api.errors import Codes
 from synapse.rest import admin
 from synapse.rest.client import login, sendtodevice, sync
+from synapse.server import HomeServer
 from synapse.types import JsonDict
+from synapse.util import Clock
+from synapse.util.stringutils import random_string
 
 from tests.unittest import HomeserverTestCase, override_config
 
@@ -61,8 +69,18 @@ class SendToDeviceTestCase(HomeserverTestCase):
     def default_config(self) -> JsonDict:
         config = super().default_config()
         config["experimental_features"] = self.experimental_features
+        config["federation_sender_instances"] = None
         return config
 
+    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
+        self.federation_transport_client = Mock(spec=["send_transaction"])
+        self.federation_transport_client.send_transaction = AsyncMock()
+        hs = self.setup_test_homeserver(
+            federation_transport_client=self.federation_transport_client,
+        )
+
+        return hs
+
     def test_user_to_user(self) -> None:
         """A to-device message from one user to another should get delivered"""
 
@@ -113,6 +131,82 @@ def test_user_to_user(self) -> None:
         self.assertEqual(channel.code, 200, channel.result)
         self.assertEqual(channel.json_body.get("to_device", {}).get("events", []), [])
 
+    def test_large_remote_todevice(self) -> None:
+        """A to-device message needs to fit in the EDU size limit"""
+        _ = self.register_user("u1", "pass")
+        user1_tok = self.login("u1", "pass", "d1")
+
+        # send the message
+        test_msg = {"foo": random_string(MAX_EDU_SIZE)}
+        channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/sendToDevice/m.test/12345",
+            content={"messages": {"@remote_user:secondserver": {"device": test_msg}}},
+            access_token=user1_tok,
+        )
+        self.assertEqual(channel.code, 413, channel.result)
+        self.assertEqual(Codes.TOO_LARGE, channel.json_body["errcode"])
+
+    def test_edu_splitting(self) -> None:
+        """Test that a bunch of to-device messages are split into multiple EDUs if they are too large"""
+        mock_send_transaction: AsyncMock = (
+            self.federation_transport_client.send_transaction
+        )
+        mock_send_transaction.return_value = {}
+
+        sender = self.hs.get_federation_sender()
+
+        _ = self.register_user("u1", "pass")
+        user1_tok = self.login("u1", "pass", "d1")
+        destination = "secondserver"
+        messages = {}
+
+        # 2 small messages that should fit in a single EDU
+        for i in range(2):
+            messages[f"@remote_user{i}:" + destination] = {
+                "device": {"foo": random_string(100)}
+            }
+
+        channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/sendToDevice/m.test/123456",
+            content={"messages": messages},
+            access_token=user1_tok,
+        )
+        self.assertEqual(channel.code, 200, channel.result)
+
+        self.get_success(sender.send_device_messages([destination]))
+
+        self.pump()
+
+        json_cb = mock_send_transaction.call_args[0][1]
+        data = json_cb()
+        self.assertEqual(len(data["edus"]), 1)
+
+        mock_send_transaction.reset_mock()
+
+        # 2 messages, each just big enough to fit in an EDU
+        for i in range(2):
+            messages[f"@remote_user{i}:" + destination] = {
+                "device": {"foo": random_string(MAX_EDU_SIZE - 1000)}
+            }
+
+        channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/sendToDevice/m.test/1234567",
+            content={"messages": messages},
+            access_token=user1_tok,
+        )
+        self.assertEqual(channel.code, 200, channel.result)
+
+        self.get_success(sender.send_device_messages([destination]))
+
+        self.pump()
+
+        json_cb = mock_send_transaction.call_args[0][1]
+        data = json_cb()
+        self.assertEqual(len(data["edus"]), 2)
+
     @override_config({"rc_key_requests": {"per_second": 10, "burst_count": 2}})
     def test_local_room_key_request(self) -> None:
         """m.room_key_request has special-casing; test from local user"""