Query missing cross-signing keys on local sig upload

Add changelog

Save retrieved keys to the db

lint

Fix and de-brittle remote result dict processing

Use query_user_devices instead, assume only master, self_signing key types

Make changelog more useful

Remove very specific exception handling

Wrap get_verify_key_from_cross_signing_key in a try/except

Note that _get_e2e_cross_signing_verify_key can raise a SynapseError

lint

Add comment explaining why this is useful

Only fetch master and self_signing key types

Fix log statements, docstrings

Remove extraneous items from remote query try/except

lint

Factor key retrieval out into a separate function

Send device updates, modeled after SigningKeyEduUpdater._handle_signing_key_updates

Update method docstring

2 files changed, 140 insertions, 12 deletions

diff --git a/synapse/federation/transport/client.py b/synapse/federation/transport/client.py
index dc563538de..c35637a571 100644
--- a/synapse/federation/transport/client.py
+++ b/synapse/federation/transport/client.py
@@ -406,13 +406,19 @@ class TransportLayerClient(object):
               "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 containg the device keys.
+            A dict containing device and cross-signing keys.
         """
         path = _create_v1_path("/user/keys/query")
 
@@ -429,14 +435,16 @@ class TransportLayerClient(object):
         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 8d7075f2eb..afc173ab2f 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import logging
+from typing import Dict, Optional, Tuple
 
 from six import iteritems
 
@@ -23,6 +24,7 @@ import attr
 from canonicaljson import encode_canonical_json, json
 from signedjson.key import decode_verify_key_bytes
 from signedjson.sign import SignatureVerifyException, verify_signed_json
+from signedjson.types import VerifyKey
 from unpaddedbase64 import decode_base64
 
 from twisted.internet import defer
@@ -174,8 +176,8 @@ class E2eKeysHandler(object):
             """This is called when we are querying the device list of a user on
             a remote homeserver and their device list is not in the device list
             cache. If we share a room with this user and we're not querying for
-            specific user we will update the cache
-            with their device list."""
+            specific user we will update the cache with their device list.
+            """
 
             destination_query = remote_queries_not_in_cache[destination]
 
@@ -961,13 +963,19 @@ class E2eKeysHandler(object):
         return signature_list, failures
 
     @defer.inlineCallbacks
-    def _get_e2e_cross_signing_verify_key(self, user_id, key_type, from_user_id=None):
-        """Fetch the cross-signing public key from storage and interpret it.
+    def _get_e2e_cross_signing_verify_key(
+        self, user_id: str, key_type: str, from_user_id: str = None
+    ):
+        """Fetch locally or remotely query for a cross-signing public key.
+
+        First, attempt to fetch the cross-signing public key from storage.
+        If that fails, query the keys from the homeserver they belong to
+        and update our local copy.
 
         Args:
-            user_id (str): the user whose key should be fetched
-            key_type (str): the type of key to fetch
-            from_user_id (str): the user that we are fetching the keys for.
+            user_id: the user whose key should be fetched
+            key_type: the type of key to fetch
+            from_user_id: the user that we are fetching the keys for.
                 This affects what signatures are fetched.
 
         Returns:
@@ -976,16 +984,128 @@ class E2eKeysHandler(object):
 
         Raises:
             NotFoundError: if the key is not found
+            SynapseError: if `user_id` is invalid
         """
+        user = UserID.from_string(user_id)
+        key_id = None
+        verify_key = None
+
         key = yield self.store.get_e2e_cross_signing_key(
             user_id, key_type, from_user_id
         )
+
+        # If we couldn't find the key locally, and we're looking for keys of
+        # another user then attempt to fetch the missing key from the remote
+        # user's server.
+        #
+        # We may run into this in possible edge cases where a user tries to
+        # cross-sign a remote user, but does not share any rooms with them yet.
+        # Thus, we would not have their key list yet. We fetch the key here,
+        # store it and notify clients of new, associated device IDs.
+        if (
+            key is None
+            and not self.is_mine(user)
+            # We only get "master" and "self_signing" keys from remote servers
+            and key_type in ["master", "self_signing"]
+        ):
+            key = yield self._retrieve_cross_signing_keys_for_remote_user(
+                user, key_type
+            )
+
         if key is None:
-            logger.debug("no %s key found for %s", key_type, user_id)
+            logger.debug("No %s key found for %s", key_type, user_id)
             raise NotFoundError("No %s key found for %s" % (key_type, user_id))
-        key_id, verify_key = get_verify_key_from_cross_signing_key(key)
+
+        # If we retrieved the keys remotely, these values will already be set
+        if key_id is None or verify_key is None:
+            try:
+                key_id, verify_key = get_verify_key_from_cross_signing_key(key)
+            except ValueError as e:
+                logger.debug(
+                    "Invalid %s key retrieved: %s - %s %s", key_type, key, type(e), e,
+                )
+                raise SynapseError(
+                    502, "Invalid %s key retrieved from remote server", key_type
+                )
+
         return key, key_id, verify_key
 
+    @defer.inlineCallbacks
+    def _retrieve_cross_signing_keys_for_remote_user(
+        self, user: UserID, desired_key_type: str,
+    ) -> Tuple[Optional[Dict], Optional[str], Optional[VerifyKey]]:
+        """Queries cross-signing keys for a remote user and saves them to the database
+
+        Only the key specified by `key_type` will be returned, while all retrieved keys
+        will be saved regardless
+
+        Args:
+            user: The user to query remote keys for
+            desired_key_type: The type of key to receive. One of "master", "self_signing"
+
+        Returns:
+            A tuple of the retrieved key content, the key's ID and the matching VerifyKey.
+            If the key cannot be retrieved, all values in the tuple will instead be None.
+        """
+        try:
+            remote_result = yield self.federation.query_user_devices(
+                user.domain, user.to_string()
+            )
+        except Exception as e:
+            logger.warning(
+                "Unable to query %s for cross-signing keys of user %s: %s %s",
+                user.domain,
+                user.to_string(),
+                type(e),
+                e,
+            )
+            return None
+
+        # Process each of the retrieved cross-signing keys
+        final_key = None
+        final_key_id = None
+        final_verify_key = None
+        device_ids = []
+        for key_type in ["master", "self_signing"]:
+            key_content = remote_result.get(key_type + "_key")
+            if not key_content:
+                continue
+
+            # At the same time, store this key in the db for
+            # subsequent queries
+            yield self.store.set_e2e_cross_signing_key(
+                user.to_string(), key_type, key_content
+            )
+
+            # Note down the device ID attached to this key
+            try:
+                # verify_key is a VerifyKey from signedjson, which uses
+                # .version to denote the portion of the key ID after the
+                # algorithm and colon, which is the device ID
+                key_id, verify_key = get_verify_key_from_cross_signing_key(key_content)
+            except ValueError as e:
+                logger.debug(
+                    "Invalid %s key retrieved: %s - %s %s",
+                    key_type,
+                    key_content,
+                    type(e),
+                    e,
+                )
+                continue
+            device_ids.append(verify_key.version)
+
+            # If this is the desired key type, save it and it's ID/VerifyKey
+            if key_type == desired_key_type:
+                final_key = key_content
+                final_verify_key = verify_key
+                final_key_id = key_id
+
+        # Notify clients that new devices for this user have been discovered
+        if device_ids:
+            yield self.device_handler.notify_device_update(user.to_string(), device_ids)
+
+        return final_key, final_key_id, final_verify_key
+
 
 def _check_cross_signing_key(key, user_id, key_type, signing_key=None):
     """Check a cross-signing key uploaded by a user.  Performs some basic sanity