diff --git a/synapse/handlers/appservice.py b/synapse/handlers/appservice.py
index 953df4d9cd..da887647d4 100644
--- a/synapse/handlers/appservice.py
+++ b/synapse/handlers/appservice.py
@@ -18,6 +18,7 @@ from typing import (
Dict,
Iterable,
List,
+ Mapping,
Optional,
Tuple,
Union,
@@ -846,6 +847,10 @@ class ApplicationServicesHandler:
]:
"""Claim one time keys from application services.
+ Users which are exclusively owned by an application service are sent a
+ key claim request to check if the application service provides keys
+ directly.
+
Args:
query: An iterable of tuples of (user ID, device ID, algorithm).
@@ -901,3 +906,59 @@ class ApplicationServicesHandler:
missing.extend(result[1])
return claimed_keys, missing
+
+ async def query_keys(
+ self, query: Mapping[str, Optional[List[str]]]
+ ) -> Dict[str, Dict[str, Dict[str, JsonDict]]]:
+ """Query application services for device keys.
+
+ Users which are exclusively owned by an application service are queried
+ for keys to check if the application service provides keys directly.
+
+ Args:
+ query: map from user_id to a list of devices to query
+
+ Returns:
+ A map from user_id -> device_id -> device details
+ """
+ services = self.store.get_app_services()
+
+ # Partition the users by appservice.
+ query_by_appservice: Dict[str, Dict[str, List[str]]] = {}
+ for user_id, device_ids in query.items():
+ if not self.store.get_if_app_services_interested_in_user(user_id):
+ continue
+
+ # Find the associated appservice.
+ for service in services:
+ if service.is_exclusive_user(user_id):
+ query_by_appservice.setdefault(service.id, {})[user_id] = (
+ device_ids or []
+ )
+ continue
+
+ # Query each service in parallel.
+ results = await make_deferred_yieldable(
+ defer.DeferredList(
+ [
+ run_in_background(
+ self.appservice_api.query_keys,
+ # We know this must be an app service.
+ self.store.get_app_service_by_id(service_id), # type: ignore[arg-type]
+ service_query,
+ )
+ for service_id, service_query in query_by_appservice.items()
+ ],
+ consumeErrors=True,
+ )
+ )
+
+ # Patch together the results -- they are all independent (since they
+ # require exclusive control over the users). They get returned as a single
+ # dictionary.
+ key_queries: Dict[str, Dict[str, Dict[str, JsonDict]]] = {}
+ for success, result in results:
+ if success:
+ key_queries.update(result)
+
+ return key_queries
diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
index 9e7c2c45b5..0073667470 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -91,6 +91,9 @@ class E2eKeysHandler:
self._query_appservices_for_otks = (
hs.config.experimental.msc3983_appservice_otk_claims
)
+ self._query_appservices_for_keys = (
+ hs.config.experimental.msc3984_appservice_key_query
+ )
@trace
@cancellable
@@ -497,6 +500,19 @@ class E2eKeysHandler:
local_query, include_displaynames
)
+ # Check if the application services have any additional results.
+ if self._query_appservices_for_keys:
+ # Query the appservices for any keys.
+ appservice_results = await self._appservice_handler.query_keys(query)
+
+ # Merge results, overriding with what the appservice returned.
+ for user_id, devices in appservice_results.get("device_keys", {}).items():
+ # Copy the appservice device info over the homeserver device info, but
+ # don't completely overwrite it.
+ results.setdefault(user_id, {}).update(devices)
+
+ # TODO Handle cross-signing keys.
+
# Build the result structure
for user_id, device_keys in results.items():
for device_id, device_info in device_keys.items():
|
