diff --git a/tests/handlers/test_appservice.py b/tests/handlers/test_appservice.py
index 9918ff6807..6e0ec37963 100644
--- a/tests/handlers/test_appservice.py
+++ b/tests/handlers/test_appservice.py
@@ -16,17 +16,25 @@ from typing import Dict, Iterable, List, Optional
from unittest.mock import Mock
from twisted.internet import defer
+from twisted.test.proto_helpers import MemoryReactor
import synapse.rest.admin
import synapse.storage
-from synapse.appservice import ApplicationService
+from synapse.appservice import (
+ ApplicationService,
+ TransactionOneTimeKeyCounts,
+ TransactionUnusedFallbackKeys,
+)
from synapse.handlers.appservice import ApplicationServicesHandler
-from synapse.rest.client import login, receipts, room, sendtodevice
+from synapse.rest.client import login, receipts, register, room, sendtodevice
+from synapse.server import HomeServer
from synapse.types import RoomStreamToken
+from synapse.util import Clock
from synapse.util.stringutils import random_string
from tests import unittest
from tests.test_utils import make_awaitable, simple_async_mock
+from tests.unittest import override_config
from tests.utils import MockClock
@@ -428,7 +436,14 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
#
# The uninterested application service should not have been notified at all.
self.send_mock.assert_called_once()
- service, _events, _ephemeral, to_device_messages = self.send_mock.call_args[0]
+ (
+ service,
+ _events,
+ _ephemeral,
+ to_device_messages,
+ _otks,
+ _fbks,
+ ) = self.send_mock.call_args[0]
# Assert that this was the same to-device message that local_user sent
self.assertEqual(service, interested_appservice)
@@ -540,7 +555,7 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
service_id_to_message_count: Dict[str, int] = {}
for call in self.send_mock.call_args_list:
- service, _events, _ephemeral, to_device_messages = call[0]
+ service, _events, _ephemeral, to_device_messages, _otks, _fbks = call[0]
# Check that this was made to an interested service
self.assertIn(service, interested_appservices)
@@ -582,3 +597,174 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
self._services.append(appservice)
return appservice
+
+
+class ApplicationServicesHandlerOtkCountsTestCase(unittest.HomeserverTestCase):
+ # Argument indices for pulling out arguments from a `send_mock`.
+ ARG_OTK_COUNTS = 4
+ ARG_FALLBACK_KEYS = 5
+
+ servlets = [
+ synapse.rest.admin.register_servlets_for_client_rest_resource,
+ login.register_servlets,
+ register.register_servlets,
+ room.register_servlets,
+ sendtodevice.register_servlets,
+ receipts.register_servlets,
+ ]
+
+ def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+ # Mock the ApplicationServiceScheduler's _TransactionController's send method so that
+ # we can track what's going out
+ self.send_mock = simple_async_mock()
+ hs.get_application_service_handler().scheduler.txn_ctrl.send = self.send_mock # type: ignore[assignment] # We assign to a method.
+
+ # Define an application service for the tests
+ self._service_token = "VERYSECRET"
+ self._service = ApplicationService(
+ self._service_token,
+ "as1.invalid",
+ "as1",
+ "@as.sender:test",
+ namespaces={
+ "users": [
+ {"regex": "@_as_.*:test", "exclusive": True},
+ {"regex": "@as.sender:test", "exclusive": True},
+ ]
+ },
+ msc3202_transaction_extensions=True,
+ )
+ self.hs.get_datastores().main.services_cache = [self._service]
+
+ # Register some appservice users
+ self._sender_user, self._sender_device = self.register_appservice_user(
+ "as.sender", self._service_token
+ )
+ self._namespaced_user, self._namespaced_device = self.register_appservice_user(
+ "_as_user1", self._service_token
+ )
+
+ # Register a real user as well.
+ self._real_user = self.register_user("real.user", "meow")
+ self._real_user_token = self.login("real.user", "meow")
+
+ async def _add_otks_for_device(
+ self, user_id: str, device_id: str, otk_count: int
+ ) -> None:
+ """
+ Add some dummy keys. It doesn't matter if they're not a real algorithm;
+ that should be opaque to the server anyway.
+ """
+ await self.hs.get_datastores().main.add_e2e_one_time_keys(
+ user_id,
+ device_id,
+ self.clock.time_msec(),
+ [("algo", f"k{i}", "{}") for i in range(otk_count)],
+ )
+
+ async def _add_fallback_key_for_device(
+ self, user_id: str, device_id: str, used: bool
+ ) -> None:
+ """
+ Adds a fake fallback key to a device, optionally marking it as used
+ right away.
+ """
+ store = self.hs.get_datastores().main
+ await store.set_e2e_fallback_keys(user_id, device_id, {"algo:fk": "fall back!"})
+ if used is True:
+ # Mark the key as used
+ await store.db_pool.simple_update_one(
+ table="e2e_fallback_keys_json",
+ keyvalues={
+ "user_id": user_id,
+ "device_id": device_id,
+ "algorithm": "algo",
+ "key_id": "fk",
+ },
+ updatevalues={"used": True},
+ desc="_get_fallback_key_set_used",
+ )
+
+ def _set_up_devices_and_a_room(self) -> str:
+ """
+ Helper to set up devices for all the users
+ and a room for the users to talk in.
+ """
+
+ async def preparation():
+ await self._add_otks_for_device(self._sender_user, self._sender_device, 42)
+ await self._add_fallback_key_for_device(
+ self._sender_user, self._sender_device, used=True
+ )
+ await self._add_otks_for_device(
+ self._namespaced_user, self._namespaced_device, 36
+ )
+ await self._add_fallback_key_for_device(
+ self._namespaced_user, self._namespaced_device, used=False
+ )
+
+ # Register a device for the real user, too, so that we can later ensure
+ # that we don't leak information to the AS about the non-AS user.
+ await self.hs.get_datastores().main.store_device(
+ self._real_user, "REALDEV", "UltraMatrix 3000"
+ )
+ await self._add_otks_for_device(self._real_user, "REALDEV", 50)
+
+ self.get_success(preparation())
+
+ room_id = self.helper.create_room_as(
+ self._real_user, is_public=True, tok=self._real_user_token
+ )
+ self.helper.join(
+ room_id,
+ self._namespaced_user,
+ tok=self._service_token,
+ appservice_user_id=self._namespaced_user,
+ )
+
+ # Check it was called for sanity. (This was to send the join event to the AS.)
+ self.send_mock.assert_called()
+ self.send_mock.reset_mock()
+
+ return room_id
+
+ @override_config(
+ {"experimental_features": {"msc3202_transaction_extensions": True}}
+ )
+ def test_application_services_receive_otk_counts_and_fallback_key_usages_with_pdus(
+ self,
+ ) -> None:
+ """
+ Tests that:
+ - the AS receives one-time key counts and unused fallback keys for:
+ - the specified sender; and
+ - any user who is in receipt of the PDUs
+ """
+
+ room_id = self._set_up_devices_and_a_room()
+
+ # Send a message into the AS's room
+ self.helper.send(room_id, "woof woof", tok=self._real_user_token)
+
+ # Capture what was sent as an AS transaction.
+ self.send_mock.assert_called()
+ last_args, _last_kwargs = self.send_mock.call_args
+ otks: Optional[TransactionOneTimeKeyCounts] = last_args[self.ARG_OTK_COUNTS]
+ unused_fallbacks: Optional[TransactionUnusedFallbackKeys] = last_args[
+ self.ARG_FALLBACK_KEYS
+ ]
+
+ self.assertEqual(
+ otks,
+ {
+ "@as.sender:test": {self._sender_device: {"algo": 42}},
+ "@_as_user1:test": {self._namespaced_device: {"algo": 36}},
+ },
+ )
+ self.assertEqual(
+ unused_fallbacks,
+ {
+ "@as.sender:test": {self._sender_device: []},
+ "@_as_user1:test": {self._namespaced_device: ["algo"]},
+ },
+ )
|
