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:
|
