Experimental: /sendToDevice with shared content and cached recipient sets

Behind experimental_features.photon_shared_to_device_enabled (advertised as
org.matrix.experimental.photon.shared_to_device). A body of
{content, recipients: {base, add, remove}} sends content unchanged to the
base set (looked up by hash, per sending device) minus remove plus add, and
caches the resulting set under its hash for later deltas. An unknown base
returns ORG.MATRIX.EXPERIMENTAL.UNKNOWN_RECIPIENT_SET so the client resends
the full list.

rust/src/config/mod.rs

@@ -76,4 +76,5 @@ pub struct ExperimentalConfig {
     pub msc4143_enabled: bool,
     pub msc4446_enabled: bool,
     pub msc4502_enabled: bool,
+    pub photon_shared_to_device_enabled: bool,
 }

rust/src/handlers/versions.rs

@@ -282,6 +282,9 @@ pub struct UnstableFeatureMap {
     /// MSC4502: Targeted and unrestricted room member queries
     #[serde(rename = "io.element.msc4502")]
     msc4502: bool,
+    /// Photon: shared to-device content with cached recipient sets
+    #[serde(rename = "org.matrix.experimental.photon.shared_to_device")]
+    photon_shared_to_device: bool,
 
     // Whether new rooms will be set to encrypted or not (based on presets).
     #[serde(rename = "io.element.e2ee_forced.public")]
@@ -336,6 +339,7 @@ pub fn synapse_config_to_global_unstable_feature_map(
         msc4143_enabled: config.experimental.msc4143_enabled,
         msc4446_enabled: config.experimental.msc4446_enabled,
         msc4502: config.experimental.msc4502_enabled,
+        photon_shared_to_device: config.experimental.photon_shared_to_device_enabled,
         e2ee_forced_public: config
             .room
             .encryption_enabled_by_default_for_room_presets

synapse/config/experimental.py

@@ -318,3 +318,9 @@ class ExperimentalConfig(Config):
 
         # MSC4512: Delegating parts of the C-S and S-S API to application services
         self.msc4512_enabled: bool = experimental.get("msc4512_enabled", False)
+
+        # Photon: `/sendToDevice` with the content once and a recipient set
+        # given as a delta against a server-cached set.
+        self.photon_shared_to_device_enabled: bool = experimental.get(
+            "photon_shared_to_device_enabled", False
+        )

synapse/rest/client/sendtodevice.py

@@ -19,9 +19,14 @@
 #
 #
 
+import hashlib
 import logging
 from typing import TYPE_CHECKING
 
+from canonicaljson import encode_canonical_json
+from unpaddedbase64 import encode_base64
+
+from synapse.api.errors import SynapseError
 from synapse.http import servlet
 from synapse.http.server import HttpServer
 from synapse.http.servlet import assert_params_in_dict, parse_json_object_from_request
@@ -29,6 +34,7 @@ from synapse.http.site import SynapseRequest
 from synapse.logging.opentracing import set_tag
 from synapse.rest.client.transactions import HttpTransactionCache
 from synapse.types import JsonDict, Requester
+from synapse.util.caches.lrucache import LruCache
 
 from ._base import client_patterns
 
@@ -37,6 +43,28 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+UNKNOWN_RECIPIENT_SET = "ORG.MATRIX.EXPERIMENTAL.UNKNOWN_RECIPIENT_SET"
+
+RecipientSet = frozenset[tuple[str, str]]
+
+
+def recipient_set_hash(recipients: RecipientSet) -> str:
+    """Unpadded base64 of the SHA-256 of the canonical JSON of the sorted
+    `[user_id, device_id]` pairs."""
+    pairs = [list(pair) for pair in sorted(recipients)]
+    return encode_base64(hashlib.sha256(encode_canonical_json(pairs)).digest())
+
+
+def _parse_pairs(value: object, name: str) -> RecipientSet:
+    if not isinstance(value, list) or not all(
+        isinstance(pair, list)
+        and len(pair) == 2
+        and all(isinstance(part, str) for part in pair)
+        for pair in value
+    ):
+        raise SynapseError(400, f"'{name}' must be a list of [user_id, device_id]")
+    return frozenset((user_id, device_id) for user_id, device_id in value)
+
 
 class SendToDeviceRestServlet(servlet.RestServlet):
     PATTERNS = client_patterns(
@@ -50,6 +78,17 @@ class SendToDeviceRestServlet(servlet.RestServlet):
         self.auth = hs.get_auth()
         self.txns = HttpTransactionCache(hs)
         self.device_message_handler = hs.get_device_message_handler()
+        self._shared_enabled = hs.config.experimental.photon_shared_to_device_enabled
+        # (user_id, device_id, set hash) -> set. Soft state: a miss makes the
+        # client resend the full list, so it can be evicted, lost on restart
+        # and differ between workers.
+        # ponytail: whole sets per entry; store deltas if memory bites.
+        self._recipient_sets: LruCache[tuple[str, str, str], RecipientSet] = LruCache(
+            max_size=1000,
+            clock=hs.get_clock(),
+            server_name=hs.hostname,
+            cache_name="photon_recipient_sets",
+        )
 
     async def on_PUT(
         self, request: SynapseRequest, message_type: str, txn_id: str
@@ -72,14 +111,52 @@ class SendToDeviceRestServlet(servlet.RestServlet):
         message_type: str,
     ) -> tuple[int, JsonDict]:
         content = parse_json_object_from_request(request)
-        assert_params_in_dict(content, ("messages",))
+        if self._shared_enabled and "recipients" in content:
+            messages = self._expand_shared(requester, content)
+        else:
+            assert_params_in_dict(content, ("messages",))
+            messages = content["messages"]
 
         await self.device_message_handler.send_device_message(
-            requester, message_type, content["messages"]
+            requester, message_type, messages
         )
 
         return 200, {}
 
+    def _expand_shared(self, requester: Requester, body: JsonDict) -> JsonDict:
+        """Turn `{"content", "recipients": {"base", "add", "remove"}}` into
+        ordinary per-device `messages`, caching the resulting recipient set."""
+        recipients = body["recipients"]
+        if "messages" in body or not isinstance(recipients, dict):
+            raise SynapseError(
+                400, "'recipients' must be an object, without 'messages'"
+            )
+        assert_params_in_dict(body, ("content",))
+        if not isinstance(body["content"], dict):
+            raise SynapseError(400, "'content' must be an object")
+
+        owner = (requester.user.to_string(), requester.device_id or "")
+        targets: RecipientSet = frozenset()
+        base = recipients.get("base")
+        if base is not None:
+            cached = self._recipient_sets.get((*owner, base))
+            if cached is None:
+                raise SynapseError(
+                    400,
+                    "Unknown recipient set; send the full list",
+                    UNKNOWN_RECIPIENT_SET,
+                )
+            targets = cached
+        targets = (targets - _parse_pairs(recipients.get("remove", []), "remove")) | (
+            _parse_pairs(recipients.get("add", []), "add")
+        )
+        self._recipient_sets.set((*owner, recipient_set_hash(targets)), targets)
+
+        messages: dict[str, dict[str, JsonDict]] = {}
+        for user_id, device_id in targets:
+            messages.setdefault(user_id, {})[device_id] = body["content"]
+        return messages
+
 
 def register_servlets(hs: "HomeServer", http_server: HttpServer) -> None:
     SendToDeviceRestServlet(hs).register(http_server)

tests/rest/client/test_sendtodevice.py

@@ -35,10 +35,12 @@ from synapse.handlers.devicemessage import (
 )
 from synapse.rest import admin
 from synapse.rest.client import login, sendtodevice, sync
+from synapse.rest.client.sendtodevice import UNKNOWN_RECIPIENT_SET, recipient_set_hash
 from synapse.server import HomeServer
 from synapse.types import JsonDict
 from synapse.util.clock import Clock
 
+from tests.server import FakeChannel
 from tests.unittest import HomeserverTestCase, override_config
 
 
@@ -586,3 +588,104 @@ class SendToDeviceTestCase(HomeserverTestCase):
         self.assertEqual(channel.code, 200, channel.result)
         messages = channel.json_body.get("to_device", {}).get("events", [])
         self.assertEqual(len(messages), 50)
+
+
+class SharedToDeviceTestCase(HomeserverTestCase):
+    """`/sendToDevice` with the content once and a cached recipient set."""
+
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        sendtodevice.register_servlets,
+        sync.register_servlets,
+    ]
+
+    def default_config(self) -> JsonDict:
+        config = super().default_config()
+        config.setdefault(
+            "experimental_features", {"photon_shared_to_device_enabled": True}
+        )
+        return config
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        self.sender = self.register_user("sender", "pass")
+        self.sender_tok = self.login("sender", "pass", "S")
+        self.tokens = {}
+        for name in ("a", "b", "c", "d"):
+            user = self.register_user(name, "pass")
+            self.tokens[(user, name.upper())] = self.login(name, "pass", name.upper())
+        self.txn = 0
+
+    def send(self, recipients: JsonDict, body: str) -> FakeChannel:
+        self.txn += 1
+        return self.make_request(
+            "PUT",
+            f"/_matrix/client/v3/sendToDevice/m.test/{self.txn}",
+            content={"content": {"body": body}, "recipients": recipients},
+            access_token=self.sender_tok,
+        )
+
+    def received(self) -> set[tuple[str, str]]:
+        """Which devices have a message they haven't seen yet, consuming it."""
+        got = set()
+        for (user, device), tok in self.tokens.items():
+            since = getattr(self, f"since_{device}", None)
+            path = "/sync" + (f"?since={since}" if since else "")
+            channel = self.make_request("GET", path, access_token=tok)
+            setattr(self, f"since_{device}", channel.json_body["next_batch"])
+            if channel.json_body.get("to_device", {}).get("events"):
+                got.add((user, device))
+        return got
+
+    def pair(self, name: str) -> tuple[str, str]:
+        return (f"@{name}:test", name.upper())
+
+    def test_base_plus_add_minus_remove(self) -> None:
+        a, b, c, d = (self.pair(n) for n in "abcd")
+        abc = recipient_set_hash(frozenset([a, b, c]))
+
+        # Full list seeds hash(a b c).
+        self.assertEqual(self.send({"add": [a, b, c]}, "1").code, 200)
+        self.assertEqual(self.received(), {a, b, c})
+
+        # hash(a b c) + d, which seeds hash(a b c d).
+        self.assertEqual(self.send({"base": abc, "add": [d]}, "2").code, 200)
+        self.assertEqual(self.received(), {a, b, c, d})
+        abcd = recipient_set_hash(frozenset([a, b, c, d]))
+
+        # hash(a b c d) - b.
+        self.assertEqual(self.send({"base": abcd, "remove": [b]}, "3").code, 200)
+        self.assertEqual(self.received(), {a, c, d})
+
+        # The older base is still usable.
+        self.assertEqual(self.send({"base": abc}, "4").code, 200)
+        self.assertEqual(self.received(), {a, b, c})
+
+    def test_unknown_base(self) -> None:
+        channel = self.send({"base": "nope", "add": [self.pair("a")]}, "x")
+        self.assertEqual(channel.code, 400)
+        self.assertEqual(channel.json_body["errcode"], UNKNOWN_RECIPIENT_SET)
+        self.assertEqual(self.received(), set())
+
+    def test_sets_are_per_sending_device(self) -> None:
+        a = self.pair("a")
+        self.assertEqual(self.send({"add": [a]}, "1").code, 200)
+        other_tok = self.login("sender", "pass", "S2")
+        channel = self.make_request(
+            "PUT",
+            "/_matrix/client/v3/sendToDevice/m.test/other",
+            content={
+                "content": {},
+                "recipients": {"base": recipient_set_hash(frozenset([a]))},
+            },
+            access_token=other_tok,
+        )
+        self.assertEqual(channel.json_body["errcode"], UNKNOWN_RECIPIENT_SET)
+
+    @override_config(
+        {"experimental_features": {"photon_shared_to_device_enabled": False}}
+    )
+    def test_disabled(self) -> None:
+        channel = self.send({"add": [self.pair("a")]}, "x")
+        self.assertEqual(channel.code, 400)
+        self.assertEqual(channel.json_body["errcode"], Codes.MISSING_PARAM)