diff --git a/synapse/appservice/scheduler.py b/synapse/appservice/scheduler.py
index b4e602e880..72417151ba 100644
--- a/synapse/appservice/scheduler.py
+++ b/synapse/appservice/scheduler.py
@@ -54,12 +54,19 @@ from typing import (
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 @@ class ApplicationServiceScheduler:
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 @@ class _ServiceQueuer:
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 @@ class _ServiceQueuer:
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 @@ class _TransactionController:
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,6 +332,10 @@ class _TransactionController:
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(
@@ -255,6 +343,8 @@ class _TransactionController:
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:
|