From fd8d154e6da6830b1cc828b7e6dfcddf53b9cad2 Mon Sep 17 00:00:00 2001
From: Andrew Morgan <andrew@amorgan.xyz>
Date: Thu, 16 Apr 2020 17:13:32 +0100
Subject: [PATCH] Use query_user_devices instead, assume only master,
 self_signing key types

---
 synapse/federation/transport/client.py | 14 ++++--
 synapse/handlers/e2e_keys.py           | 69 ++++++++++++--------------
 2 files changed, 42 insertions(+), 41 deletions(-)

diff --git a/synapse/federation/transport/client.py b/synapse/federation/transport/client.py
index d25db5daae2f..c35637a571dd 100644
--- a/synapse/federation/transport/client.py
+++ b/synapse/federation/transport/client.py
@@ -406,13 +406,19 @@ def query_client_keys(self, destination, query_content, timeout):
               "device_keys": {
                 "<user_id>": {
                   "<device_id>": {...}
+              } }
+              "master_keys": {
+                "<user_id>": {...}
+              } }
+              "self_signing_keys": {
+                "<user_id>": {...}
             } } }
 
         Args:
             destination(str): The server to query.
             query_content(dict): The user ids to query.
         Returns:
-            A dict containing the device keys.
+            A dict containing device and cross-signing keys.
         """
         path = _create_v1_path("/user/keys/query")
 
@@ -429,14 +435,16 @@ def query_user_devices(self, destination, user_id, timeout):
         Response:
             {
               "stream_id": "...",
-              "devices": [ { ... } ]
+              "devices": [ { ... } ],
+              "master_key": { ... },
+              "self_signing_key: { ... }
             }
 
         Args:
             destination(str): The server to query.
             query_content(dict): The user ids to query.
         Returns:
-            A dict containg the device keys.
+            A dict containing device and cross-signing keys.
         """
         path = _create_v1_path("/user/devices/%s", user_id)
 
diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
index c10eec8baa9e..f92544c7316f 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -969,12 +969,14 @@ def _process_other_signatures(self, user_id, signatures):
         return signature_list, failures
 
     @defer.inlineCallbacks
-    def _get_e2e_cross_signing_verify_key(self, user_id, key_type, from_user_id=None):
+    def _get_e2e_cross_signing_verify_key(
+        self, user_id, desired_key_type, from_user_id=None
+    ):
         """Fetch the cross-signing public key from storage and interpret it.
 
         Args:
             user_id (str): the user whose key should be fetched
-            key_type (str): the type of key to fetch
+            desired_key_type (str): the type of key to fetch
             from_user_id (str): the user that we are fetching the keys for.
                 This affects what signatures are fetched.
 
@@ -987,47 +989,38 @@ def _get_e2e_cross_signing_verify_key(self, user_id, key_type, from_user_id=None
         """
         user = UserID.from_string(user_id)
         key = yield self.store.get_e2e_cross_signing_key(
-            user_id, key_type, from_user_id
+            user_id, desired_key_type, from_user_id
         )
 
-        if key is None and self.is_mine(user):
-            # Attempt to fetch the missing key from the remote user's server
+        # If we still can't find the key, and we're looking for keys of another user,
+        # then attempt to fetch the missing key from the remote user's server.
+        #
+        # We don't get "user_signing" keys from remote servers, so disallow that here
+        if (
+            key is None
+            and not self.is_mine(user)
+            and desired_key_type != "user_signing"
+        ):
             try:
-                remote_result = yield self.federation.query_client_keys(
-                    user.domain, {"device_keys": {user_id: []}}, timeout=10 * 1000
+                remote_result = yield self.federation.query_user_devices(
+                    user.domain, user_id
                 )
 
-                # Process the result
-                for remote_key_type, remote_user_dict in remote_result.items():
-                    # The key_type variable passed to this function is in the form
-                    # "self_signing","master" etc. whereas the results returned from
-                    # the remote server use "self_signing_keys", "master_keys" etc.
-                    # Remove the "_keys" from the key type
-                    if remote_key_type.endswith("_keys"):
-                        remote_key_type = remote_key_type[:-5]
-
-                    # remote_user_dict is a dictionary in the form of
-                    # {
-                    #   "user_id": {
-                    #     "master_keys": ...
-                    #   },
-                    #   ...
-                    # }
-
-                    # Only extract the keys that pertain to the requested user
-                    key_content_list = remote_user_dict.get(user_id, {}).values()
-
-                    for key_content in key_content_list:
-                        # If the key_type here matches the key we're requesting,
-                        # then this is the key we want to return
-                        if remote_key_type == key_type:
-                            key = key_content
-
-                        # At the same time, save the key to the database for subsequent
-                        # queries
-                        yield self.store.set_e2e_cross_signing_key(
-                            user_id, remote_key_type, key_content
-                        )
+                # Process each of the retrieved cross-signing keys
+                for key_type in ["master", "self_signing"]:
+                    key_content = remote_result.get(key_type + "_key")
+                    if not key_content:
+                        continue
+
+                    # If this is the desired key type, return it
+                    if key_type == desired_key_type:
+                        key = key_content
+
+                    # At the same time, store this key in the db for
+                    # subsequent queries
+                    yield self.store.set_e2e_cross_signing_key(
+                        user_id, key_type, key_content
+                    )
             except (
                 HttpResponseException,
                 NotRetryingDestination,