Add support for MSC3202: sending one-time key counts and fallback key usage states to Application Services. (#11617)

Co-authored-by: Erik Johnston <[email protected]>

Author: reivilibre

changelog.d/11617.feature

@@ -0,0 +1 @@
+Add support for MSC3202: sending one-time key counts and fallback key usage states to Application Services.

synapse/appservice/__init__.py

@@ -31,6 +31,14 @@
 
 logger = logging.getLogger(__name__)
 
+# Type for the `device_one_time_key_counts` field in an appservice transaction
+#   user ID -> {device ID -> {algorithm -> count}}
+TransactionOneTimeKeyCounts = Dict[str, Dict[str, Dict[str, int]]]
+
+# Type for the `device_unused_fallback_keys` field in an appservice transaction
+#   user ID -> {device ID -> [algorithm]}
+TransactionUnusedFallbackKeys = Dict[str, Dict[str, List[str]]]
+
 
 class ApplicationServiceState(Enum):
     DOWN = "down"
@@ -72,6 +80,7 @@ def __init__(
         rate_limited: bool = True,
         ip_range_whitelist: Optional[IPSet] = None,
         supports_ephemeral: bool = False,
+        msc3202_transaction_extensions: bool = False,
     ):
         self.token = token
         self.url = (
@@ -84,6 +93,7 @@ def __init__(
         self.id = id
         self.ip_range_whitelist = ip_range_whitelist
         self.supports_ephemeral = supports_ephemeral
+        self.msc3202_transaction_extensions = msc3202_transaction_extensions
 
         if "|" in self.id:
             raise Exception("application service ID cannot contain '|' character")
@@ -339,12 +349,16 @@ def __init__(
         events: List[EventBase],
         ephemeral: List[JsonDict],
         to_device_messages: List[JsonDict],
+        one_time_key_counts: TransactionOneTimeKeyCounts,
+        unused_fallback_keys: TransactionUnusedFallbackKeys,
     ):
         self.service = service
         self.id = id
         self.events = events
         self.ephemeral = ephemeral
         self.to_device_messages = to_device_messages
+        self.one_time_key_counts = one_time_key_counts
+        self.unused_fallback_keys = unused_fallback_keys
 
     async def send(self, as_api: "ApplicationServiceApi") -> bool:
         """Sends this transaction using the provided AS API interface.
@@ -359,6 +373,8 @@ async def send(self, as_api: "ApplicationServiceApi") -> bool:
             events=self.events,
             ephemeral=self.ephemeral,
             to_device_messages=self.to_device_messages,
+            one_time_key_counts=self.one_time_key_counts,
+            unused_fallback_keys=self.unused_fallback_keys,
             txn_id=self.id,
         )
 

synapse/appservice/api.py

@@ -19,14 +19,18 @@
 
 from synapse.api.constants import EventTypes, Membership, ThirdPartyEntityKind
 from synapse.api.errors import CodeMessageException
+from synapse.appservice import (
+    ApplicationService,
+    TransactionOneTimeKeyCounts,
+    TransactionUnusedFallbackKeys,
+)
 from synapse.events import EventBase
 from synapse.events.utils import serialize_event
 from synapse.http.client import SimpleHttpClient
 from synapse.types import JsonDict, ThirdPartyInstanceID
 from synapse.util.caches.response_cache import ResponseCache
 
 if TYPE_CHECKING:
-    from synapse.appservice import ApplicationService
     from synapse.server import HomeServer
 
 logger = logging.getLogger(__name__)
@@ -219,6 +223,8 @@ async def push_bulk(
         events: List[EventBase],
         ephemeral: List[JsonDict],
         to_device_messages: List[JsonDict],
+        one_time_key_counts: TransactionOneTimeKeyCounts,
+        unused_fallback_keys: TransactionUnusedFallbackKeys,
         txn_id: Optional[int] = None,
     ) -> bool:
         """
@@ -252,7 +258,7 @@ async def push_bulk(
         uri = service.url + ("/transactions/%s" % urllib.parse.quote(str(txn_id)))
 
         # Never send ephemeral events to appservices that do not support it
-        body: Dict[str, List[JsonDict]] = {"events": serialized_events}
+        body: JsonDict = {"events": serialized_events}
         if service.supports_ephemeral:
             body.update(
                 {
@@ -262,6 +268,16 @@ async def push_bulk(
                 }
             )
 
+        if service.msc3202_transaction_extensions:
+            if one_time_key_counts:
+                body[
+                    "org.matrix.msc3202.device_one_time_key_counts"
+                ] = one_time_key_counts
+            if unused_fallback_keys:
+                body[
+                    "org.matrix.msc3202.device_unused_fallback_keys"
+                ] = unused_fallback_keys
+
         try:
             await self.put_json(
                 uri=uri,

synapse/appservice/scheduler.py

@@ -54,12 +54,19 @@
     Callable,
     Collection,
     Dict,
+    Iterable,
     List,
     Optional,
     Set,
+    Tuple,
 )
 
-from synapse.appservice import ApplicationService, ApplicationServiceState
+from synapse.appservice import (
+    ApplicationService,
+    ApplicationServiceState,
+    TransactionOneTimeKeyCounts,
+    TransactionUnusedFallbackKeys,
+)
 from synapse.appservice.api import ApplicationServiceApi
 from synapse.events import EventBase
 from synapse.logging.context import run_in_background
@@ -96,7 +103,7 @@ def __init__(self, hs: "HomeServer"):
         self.as_api = hs.get_application_service_api()
 
         self.txn_ctrl = _TransactionController(self.clock, self.store, self.as_api)
-        self.queuer = _ServiceQueuer(self.txn_ctrl, self.clock)
+        self.queuer = _ServiceQueuer(self.txn_ctrl, self.clock, hs)
 
     async def start(self) -> None:
         logger.info("Starting appservice scheduler")
@@ -153,7 +160,9 @@ class _ServiceQueuer:
     appservice at a given time.
     """
 
-    def __init__(self, txn_ctrl: "_TransactionController", clock: Clock):
+    def __init__(
+        self, txn_ctrl: "_TransactionController", clock: Clock, hs: "HomeServer"
+    ):
         # dict of {service_id: [events]}
         self.queued_events: Dict[str, List[EventBase]] = {}
         # dict of {service_id: [events]}
@@ -165,6 +174,10 @@ def __init__(self, txn_ctrl: "_TransactionController", clock: Clock):
         self.requests_in_flight: Set[str] = set()
         self.txn_ctrl = txn_ctrl
         self.clock = clock
+        self._msc3202_transaction_extensions_enabled: bool = (
+            hs.config.experimental.msc3202_transaction_extensions
+        )
+        self._store = hs.get_datastores().main
 
     def start_background_request(self, service: ApplicationService) -> None:
         # start a sender for this appservice if we don't already have one
@@ -202,15 +215,84 @@ async def _send_request(self, service: ApplicationService) -> None:
                 if not events and not ephemeral and not to_device_messages_to_send:
                     return
 
+                one_time_key_counts: Optional[TransactionOneTimeKeyCounts] = None
+                unused_fallback_keys: Optional[TransactionUnusedFallbackKeys] = None
+
+                if (
+                    self._msc3202_transaction_extensions_enabled
+                    and service.msc3202_transaction_extensions
+                ):
+                    # Compute the one-time key counts and fallback key usage states
+                    # for the users which are mentioned in this transaction,
+                    # as well as the appservice's sender.
+                    (
+                        one_time_key_counts,
+                        unused_fallback_keys,
+                    ) = await self._compute_msc3202_otk_counts_and_fallback_keys(
+                        service, events, ephemeral, to_device_messages_to_send
+                    )
+
                 try:
                     await self.txn_ctrl.send(
-                        service, events, ephemeral, to_device_messages_to_send
+                        service,
+                        events,
+                        ephemeral,
+                        to_device_messages_to_send,
+                        one_time_key_counts,
+                        unused_fallback_keys,
                     )
                 except Exception:
                     logger.exception("AS request failed")
         finally:
             self.requests_in_flight.discard(service.id)
 
+    async def _compute_msc3202_otk_counts_and_fallback_keys(
+        self,
+        service: ApplicationService,
+        events: Iterable[EventBase],
+        ephemerals: Iterable[JsonDict],
+        to_device_messages: Iterable[JsonDict],
+    ) -> Tuple[TransactionOneTimeKeyCounts, TransactionUnusedFallbackKeys]:
+        """
+        Given a list of the events, ephemeral messages and to-device messages,
+        - first computes a list of application services users that may have
+          interesting updates to the one-time key counts or fallback key usage.
+        - then computes one-time key counts and fallback key usages for those users.
+        Given a list of application service users that are interesting,
+        compute one-time key counts and fallback key usages for the users.
+        """
+
+        # Set of 'interesting' users who may have updates
+        users: Set[str] = set()
+
+        # The sender is always included
+        users.add(service.sender)
+
+        # All AS users that would receive the PDUs or EDUs sent to these rooms
+        # are classed as 'interesting'.
+        rooms_of_interesting_users: Set[str] = set()
+        # PDUs
+        rooms_of_interesting_users.update(event.room_id for event in events)
+        # EDUs
+        rooms_of_interesting_users.update(
+            ephemeral["room_id"] for ephemeral in ephemerals
+        )
+
+        # Look up the AS users in those rooms
+        for room_id in rooms_of_interesting_users:
+            users.update(
+                await self._store.get_app_service_users_in_room(room_id, service)
+            )
+
+        # Add recipients of to-device messages.
+        # device_message["user_id"] is the ID of the recipient.
+        users.update(device_message["user_id"] for device_message in to_device_messages)
+
+        # Compute and return the counts / fallback key usage states
+        otk_counts = await self._store.count_bulk_e2e_one_time_keys_for_as(users)
+        unused_fbks = await self._store.get_e2e_bulk_unused_fallback_key_types(users)
+        return otk_counts, unused_fbks
+
 
 class _TransactionController:
     """Transaction manager.
@@ -238,6 +320,8 @@ async def send(
         events: List[EventBase],
         ephemeral: Optional[List[JsonDict]] = None,
         to_device_messages: Optional[List[JsonDict]] = None,
+        one_time_key_counts: Optional[TransactionOneTimeKeyCounts] = None,
+        unused_fallback_keys: Optional[TransactionUnusedFallbackKeys] = None,
     ) -> None:
         """
         Create a transaction with the given data and send to the provided
@@ -248,13 +332,19 @@ async def send(
             events: The persistent events to include in the transaction.
             ephemeral: The ephemeral events to include in the transaction.
             to_device_messages: The to-device messages to include in the transaction.
+            one_time_key_counts: Counts of remaining one-time keys for relevant
+                appservice devices in the transaction.
+            unused_fallback_keys: Lists of unused fallback keys for relevant
+                appservice devices in the transaction.
         """
         try:
             txn = await self.store.create_appservice_txn(
                 service=service,
                 events=events,
                 ephemeral=ephemeral or [],
                 to_device_messages=to_device_messages or [],
+                one_time_key_counts=one_time_key_counts or {},
+                unused_fallback_keys=unused_fallback_keys or {},
             )
             service_is_up = await self._is_service_up(service)
             if service_is_up:

synapse/config/appservice.py

@@ -166,6 +166,16 @@ def _load_appservice(
 
     supports_ephemeral = as_info.get("de.sorunome.msc2409.push_ephemeral", False)
 
+    # Opt-in flag for the MSC3202-specific transactional behaviour.
+    # When enabled, appservice transactions contain the following information:
+    #  - device One-Time Key counts
+    #  - device unused fallback key usage states
+    msc3202_transaction_extensions = as_info.get("org.matrix.msc3202", False)
+    if not isinstance(msc3202_transaction_extensions, bool):
+        raise ValueError(
+            "The `org.matrix.msc3202` option should be true or false if specified."
+        )
+
     return ApplicationService(
         token=as_info["as_token"],
         hostname=hostname,
@@ -174,8 +184,9 @@ def _load_appservice(
         hs_token=as_info["hs_token"],
         sender=user_id,
         id=as_info["id"],
-        supports_ephemeral=supports_ephemeral,
         protocols=protocols,
         rate_limited=rate_limited,
         ip_range_whitelist=ip_range_whitelist,
+        supports_ephemeral=supports_ephemeral,
+        msc3202_transaction_extensions=msc3202_transaction_extensions,
     )

synapse/config/experimental.py

@@ -47,18 +47,24 @@ def read_config(self, config: JsonDict, **kwargs):
         # MSC3030 (Jump to date API endpoint)
         self.msc3030_enabled: bool = experimental.get("msc3030_enabled", False)
 
-        # The portion of MSC3202 which is related to device masquerading.
-        self.msc3202_device_masquerading_enabled: bool = experimental.get(
-            "msc3202_device_masquerading", False
-        )
-
         # MSC2409 (this setting only relates to optionally sending to-device messages).
         # Presence, typing and read receipt EDUs are already sent to application services that
         # have opted in to receive them. If enabled, this adds to-device messages to that list.
         self.msc2409_to_device_messages_enabled: bool = experimental.get(
             "msc2409_to_device_messages_enabled", False
         )
 
+        # The portion of MSC3202 which is related to device masquerading.
+        self.msc3202_device_masquerading_enabled: bool = experimental.get(
+            "msc3202_device_masquerading", False
+        )
+
+        # Portion of MSC3202 related to transaction extensions:
+        # sending one-time key counts and fallback key usage to application services.
+        self.msc3202_transaction_extensions: bool = experimental.get(
+            "msc3202_transaction_extensions", False
+        )
+
         # MSC3706 (server-side support for partial state in /send_join responses)
         self.msc3706_enabled: bool = experimental.get("msc3706_enabled", False)