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

Co-authored-by: Erik Johnston <erik@matrix.org>

2 files changed, 143 insertions, 2 deletions

diff --git a/synapse/storage/databases/main/appservice.py b/synapse/storage/databases/main/appservice.py
index 304814af5d..0694446558 100644
--- a/synapse/storage/databases/main/appservice.py
+++ b/synapse/storage/databases/main/appservice.py
@@ -20,14 +20,18 @@ from synapse.appservice import (
     ApplicationService,
     ApplicationServiceState,
     AppServiceTransaction,
+    TransactionOneTimeKeyCounts,
+    TransactionUnusedFallbackKeys,
 )
 from synapse.config.appservice import load_appservices
 from synapse.events import EventBase
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
+from synapse.storage.databases.main.roommember import RoomMemberWorkerStore
 from synapse.types import JsonDict
 from synapse.util import json_encoder
+from synapse.util.caches.descriptors import _CacheContext, cached
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -56,7 +60,7 @@ def _make_exclusive_regex(
     return exclusive_user_pattern
 
 
-class ApplicationServiceWorkerStore(SQLBaseStore):
+class ApplicationServiceWorkerStore(RoomMemberWorkerStore):
     def __init__(
         self,
         database: DatabasePool,
@@ -124,6 +128,18 @@ class ApplicationServiceWorkerStore(SQLBaseStore):
                 return service
         return None
 
+    @cached(iterable=True, cache_context=True)
+    async def get_app_service_users_in_room(
+        self,
+        room_id: str,
+        app_service: "ApplicationService",
+        cache_context: _CacheContext,
+    ) -> List[str]:
+        users_in_room = await self.get_users_in_room(
+            room_id, on_invalidate=cache_context.invalidate
+        )
+        return list(filter(app_service.is_interested_in_user, users_in_room))
+
 
 class ApplicationServiceStore(ApplicationServiceWorkerStore):
     # This is currently empty due to there not being any AS storage functions
@@ -199,6 +215,8 @@ class ApplicationServiceTransactionWorkerStore(
         events: List[EventBase],
         ephemeral: List[JsonDict],
         to_device_messages: List[JsonDict],
+        one_time_key_counts: TransactionOneTimeKeyCounts,
+        unused_fallback_keys: TransactionUnusedFallbackKeys,
     ) -> AppServiceTransaction:
         """Atomically creates a new transaction for this application service
         with the given list of events. Ephemeral events are NOT persisted to the
@@ -209,6 +227,10 @@ class ApplicationServiceTransactionWorkerStore(
             events: A list of persistent events to put in the transaction.
             ephemeral: A list of ephemeral events to put in the transaction.
             to_device_messages: A list of to-device messages to put 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.
 
         Returns:
             A new transaction.
@@ -244,6 +266,8 @@ class ApplicationServiceTransactionWorkerStore(
                 events=events,
                 ephemeral=ephemeral,
                 to_device_messages=to_device_messages,
+                one_time_key_counts=one_time_key_counts,
+                unused_fallback_keys=unused_fallback_keys,
             )
 
         return await self.db_pool.runInteraction(
@@ -335,12 +359,17 @@ class ApplicationServiceTransactionWorkerStore(
 
         events = await self.get_events_as_list(event_ids)
 
+        # TODO: to-device messages, one-time key counts and unused fallback keys
+        #       are not yet populated for catch-up transactions.
+        #       We likely want to populate those for reliability.
         return AppServiceTransaction(
             service=service,
             id=entry["txn_id"],
             events=events,
             ephemeral=[],
             to_device_messages=[],
+            one_time_key_counts={},
+            unused_fallback_keys={},
         )
 
     def _get_last_txn(self, txn, service_id: Optional[str]) -> int:
diff --git a/synapse/storage/databases/main/end_to_end_keys.py b/synapse/storage/databases/main/end_to_end_keys.py
index 1f8447b507..9b293475c8 100644
--- a/synapse/storage/databases/main/end_to_end_keys.py
+++ b/synapse/storage/databases/main/end_to_end_keys.py
@@ -29,6 +29,10 @@ import attr
 from canonicaljson import encode_canonical_json
 
 from synapse.api.constants import DeviceKeyAlgorithms
+from synapse.appservice import (
+    TransactionOneTimeKeyCounts,
+    TransactionUnusedFallbackKeys,
+)
 from synapse.logging.opentracing import log_kv, set_tag, trace
 from synapse.storage._base import SQLBaseStore, db_to_json
 from synapse.storage.database import (
@@ -439,6 +443,114 @@ class EndToEndKeyWorkerStore(EndToEndKeyBackgroundStore, CacheInvalidationWorker
             "count_e2e_one_time_keys", _count_e2e_one_time_keys
         )
 
+    async def count_bulk_e2e_one_time_keys_for_as(
+        self, user_ids: Collection[str]
+    ) -> TransactionOneTimeKeyCounts:
+        """
+        Counts, in bulk, the one-time keys for all the users specified.
+        Intended to be used by application services for populating OTK counts in
+        transactions.
+
+        Return structure is of the shape:
+          user_id -> device_id -> algorithm -> count
+          Empty algorithm -> count dicts are created if needed to represent a
+          lack of unused one-time keys.
+        """
+
+        def _count_bulk_e2e_one_time_keys_txn(
+            txn: LoggingTransaction,
+        ) -> TransactionOneTimeKeyCounts:
+            user_in_where_clause, user_parameters = make_in_list_sql_clause(
+                self.database_engine, "user_id", user_ids
+            )
+            sql = f"""
+                SELECT user_id, device_id, algorithm, COUNT(key_id)
+                FROM devices
+                LEFT JOIN e2e_one_time_keys_json USING (user_id, device_id)
+                WHERE {user_in_where_clause}
+                GROUP BY user_id, device_id, algorithm
+            """
+            txn.execute(sql, user_parameters)
+
+            result: TransactionOneTimeKeyCounts = {}
+
+            for user_id, device_id, algorithm, count in txn:
+                # We deliberately construct empty dictionaries for
+                # users and devices without any unused one-time keys.
+                # We *could* omit these empty dicts if there have been no
+                # changes since the last transaction, but we currently don't
+                # do any change tracking!
+                device_count_by_algo = result.setdefault(user_id, {}).setdefault(
+                    device_id, {}
+                )
+                if algorithm is not None:
+                    # algorithm will be None if this device has no keys.
+                    device_count_by_algo[algorithm] = count
+
+            return result
+
+        return await self.db_pool.runInteraction(
+            "count_bulk_e2e_one_time_keys", _count_bulk_e2e_one_time_keys_txn
+        )
+
+    async def get_e2e_bulk_unused_fallback_key_types(
+        self, user_ids: Collection[str]
+    ) -> TransactionUnusedFallbackKeys:
+        """
+        Finds, in bulk, the types of unused fallback keys for all the users specified.
+        Intended to be used by application services for populating unused fallback
+        keys in transactions.
+
+        Return structure is of the shape:
+          user_id -> device_id -> algorithms
+          Empty lists are created for devices if there are no unused fallback
+          keys. This matches the response structure of MSC3202.
+        """
+        if len(user_ids) == 0:
+            return {}
+
+        def _get_bulk_e2e_unused_fallback_keys_txn(
+            txn: LoggingTransaction,
+        ) -> TransactionUnusedFallbackKeys:
+            user_in_where_clause, user_parameters = make_in_list_sql_clause(
+                self.database_engine, "devices.user_id", user_ids
+            )
+            # We can't use USING here because we require the `.used` condition
+            # to be part of the JOIN condition so that we generate empty lists
+            # when all keys are used (as opposed to just when there are no keys at all).
+            sql = f"""
+                SELECT devices.user_id, devices.device_id, algorithm
+                FROM devices
+                LEFT JOIN e2e_fallback_keys_json AS fallback_keys
+                    ON devices.user_id = fallback_keys.user_id
+                    AND devices.device_id = fallback_keys.device_id
+                    AND NOT fallback_keys.used
+                WHERE
+                    {user_in_where_clause}
+            """
+            txn.execute(sql, user_parameters)
+
+            result: TransactionUnusedFallbackKeys = {}
+
+            for user_id, device_id, algorithm in txn:
+                # We deliberately construct empty dictionaries and lists for
+                # users and devices without any unused fallback keys.
+                # We *could* omit these empty dicts if there have been no
+                # changes since the last transaction, but we currently don't
+                # do any change tracking!
+                device_unused_keys = result.setdefault(user_id, {}).setdefault(
+                    device_id, []
+                )
+                if algorithm is not None:
+                    # algorithm will be None if this device has no keys.
+                    device_unused_keys.append(algorithm)
+
+            return result
+
+        return await self.db_pool.runInteraction(
+            "_get_bulk_e2e_unused_fallback_keys", _get_bulk_e2e_unused_fallback_keys_txn
+        )
+
     async def set_e2e_fallback_keys(
         self, user_id: str, device_id: str, fallback_keys: JsonDict
     ) -> None: