Get the end-to-end key federation working

4 files changed, 17 insertions, 21 deletions

diff --git a/synapse/federation/federation_client.py b/synapse/federation/federation_client.py
index 21a86a4c6d..44e4d0755a 100644
--- a/synapse/federation/federation_client.py
+++ b/synapse/federation/federation_client.py
@@ -135,7 +135,7 @@ class FederationClient(FederationBase):
         )
 
     @log_function
-    def query_client_keys(self, destination, content, retry_on_dns_fail=True):
+    def query_client_keys(self, destination, content):
         """Query device keys for a device hosted on a remote server.
 
         Args:
@@ -147,12 +147,10 @@ class FederationClient(FederationBase):
             response
         """
         sent_queries_counter.inc("client_device_keys")
-        return self.transport_layer.query_client_keys(
-            destination, content, retry_on_dns_fail=retry_on_dns_fail
-        )
+        return self.transport_layer.query_client_keys(destination, content)
 
     @log_function
-    def claim_client_keys(self, destination, content, retry_on_dns_fail=True):
+    def claim_client_keys(self, destination, content):
         """Claims one-time keys for a device hosted on a remote server.
 
         Args:
@@ -164,9 +162,7 @@ class FederationClient(FederationBase):
             response
         """
         sent_queries_counter.inc("client_one_time_keys")
-        return self.transport_layer.claim_client_keys(
-            destination, content, retry_on_dns_fail=retry_on_dns_fail
-        )
+        return self.transport_layer.claim_client_keys(destination, content)
 
     @defer.inlineCallbacks
     @log_function
diff --git a/synapse/federation/transport/client.py b/synapse/federation/transport/client.py
index df5083dd22..ced703364b 100644
--- a/synapse/federation/transport/client.py
+++ b/synapse/federation/transport/client.py
@@ -247,7 +247,7 @@ class TransportLayerClient(object):
         Returns:
             A dict containg the device keys.
         """
-        path = PREFIX + "/client_keys/query"
+        path = PREFIX + "/user/keys/query"
 
         content = yield self.client.post_json(
             destination=destination,
@@ -283,7 +283,7 @@ class TransportLayerClient(object):
             A dict containg the one-time keys.
         """
 
-        path = PREFIX + "/client_keys/claim"
+        path = PREFIX + "/user/keys/claim"
 
         content = yield self.client.post_json(
             destination=destination,
diff --git a/synapse/federation/transport/server.py b/synapse/federation/transport/server.py
index fb59383ecd..36f250e1a3 100644
--- a/synapse/federation/transport/server.py
+++ b/synapse/federation/transport/server.py
@@ -326,20 +326,20 @@ class FederationInviteServlet(BaseFederationServlet):
 
 
 class FederationClientKeysQueryServlet(BaseFederationServlet):
-    PATH = "/client_keys/query"
+    PATH = "/user/keys/query"
 
     @defer.inlineCallbacks
-    def on_POST(self, origin, content):
-        response = yield self.handler.on_client_key_query(origin, content)
+    def on_POST(self, origin, content, query):
+        response = yield self.handler.on_query_client_keys(origin, content)
         defer.returnValue((200, response))
 
 
 class FederationClientKeysClaimServlet(BaseFederationServlet):
-    PATH = "/client_keys/claim"
+    PATH = "/user/keys/claim"
 
     @defer.inlineCallbacks
-    def on_POST(self, origin, content):
-        response = yield self.handler.on_client_key_claim(origin, content)
+    def on_POST(self, origin, content, query):
+        response = yield self.handler.on_claim_client_keys(origin, content)
         defer.returnValue((200, response))
 
 
diff --git a/synapse/rest/client/v2_alpha/keys.py b/synapse/rest/client/v2_alpha/keys.py
index 739a08ada8..718928eedd 100644
--- a/synapse/rest/client/v2_alpha/keys.py
+++ b/synapse/rest/client/v2_alpha/keys.py
@@ -202,7 +202,7 @@ class KeyQueryServlet(RestServlet):
                     for device_id in device_ids:
                         local_query.append((user_id, device_id))
             else:
-                remote_queries.set_default(user.domain, {})[user_id] = list(
+                remote_queries.setdefault(user.domain, {})[user_id] = list(
                     device_ids
                 )
         results = yield self.store.get_e2e_device_keys(local_query)
@@ -218,7 +218,7 @@ class KeyQueryServlet(RestServlet):
             remote_result = yield self.federation.query_client_keys(
                 destination, {"device_keys": device_keys}
             )
-            for user_id, keys in remote_result.items():
+            for user_id, keys in remote_result["device_keys"].items():
                 if user_id in device_keys:
                     json_result[user_id] = keys
         defer.returnValue((200, {"device_keys": json_result}))
@@ -286,7 +286,7 @@ class OneTimeKeyServlet(RestServlet):
                 for device_id, algorithm in device_keys.items():
                     local_query.append((user_id, device_id, algorithm))
             else:
-                remote_queries.set_default(user.domain, {})[user_id] = (
+                remote_queries.setdefault(user.domain, {})[user_id] = (
                     device_keys
                 )
         results = yield self.store.claim_e2e_one_time_keys(local_query)
@@ -300,10 +300,10 @@ class OneTimeKeyServlet(RestServlet):
                     }
 
         for destination, device_keys in remote_queries.items():
-            remote_result = yield self.federation.query_client_keys(
+            remote_result = yield self.federation.claim_client_keys(
                 destination, {"one_time_keys": device_keys}
             )
-            for user_id, keys in remote_result.items():
+            for user_id, keys in remote_result["one_time_keys"].items():
                 if user_id in device_keys:
                     json_result[user_id] = keys