Add support for claiming multiple OTKs at once. (#15468)

MSC3983 provides a way to request multiple OTKs at once from appservices,
this extends this concept to the Client-Server API.

Note that this will likely be spit out into a separate MSC, but is currently part of
MSC3983.

2 files changed, 19 insertions, 24 deletions

diff --git a/tests/appservice/test_api.py b/tests/appservice/test_api.py
index 7deb923a28..15fce165b6 100644
--- a/tests/appservice/test_api.py
+++ b/tests/appservice/test_api.py
@@ -195,11 +195,11 @@ class ApplicationServiceApiTestCase(unittest.HomeserverTestCase):
 
         MISSING_KEYS = [
             # Known user, known device, missing algorithm.
-            ("@alice:example.org", "DEVICE_1", "signed_curve25519:DDDDHg"),
+            ("@alice:example.org", "DEVICE_2", "xyz", 1),
             # Known user, missing device.
-            ("@alice:example.org", "DEVICE_3", "signed_curve25519:EEEEHg"),
+            ("@alice:example.org", "DEVICE_3", "signed_curve25519", 1),
             # Unknown user.
-            ("@bob:example.org", "DEVICE_4", "signed_curve25519:FFFFHg"),
+            ("@bob:example.org", "DEVICE_4", "signed_curve25519", 1),
         ]
 
         claimed_keys, missing = self.get_success(
@@ -207,9 +207,8 @@ class ApplicationServiceApiTestCase(unittest.HomeserverTestCase):
                 self.service,
                 [
                     # Found devices
-                    ("@alice:example.org", "DEVICE_1", "signed_curve25519:AAAAHg"),
-                    ("@alice:example.org", "DEVICE_1", "signed_curve25519:BBBBHg"),
-                    ("@alice:example.org", "DEVICE_2", "signed_curve25519:CCCCHg"),
+                    ("@alice:example.org", "DEVICE_1", "signed_curve25519", 1),
+                    ("@alice:example.org", "DEVICE_2", "signed_curve25519", 1),
                 ]
                 + MISSING_KEYS,
             )
diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py
index 18edebd652..72d0584061 100644
--- a/tests/handlers/test_e2e_keys.py
+++ b/tests/handlers/test_e2e_keys.py
@@ -160,7 +160,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
         res2 = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -205,7 +205,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # key
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -224,7 +224,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # claiming an OTK again should return the same fallback key
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -273,7 +273,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -285,7 +285,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -306,7 +306,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -347,7 +347,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # return both.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )
@@ -369,7 +369,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # Claiming an OTK again should return only the fallback key.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id: "alg1"}}},
+                {local_user: {device_id: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )
@@ -1052,7 +1052,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
         # Setup a response, but only for device 2.
         self.appservice_api.claim_client_keys.return_value = make_awaitable(
-            ({local_user: {device_id_2: otk}}, [(local_user, device_id_1, "alg1")])
+            ({local_user: {device_id_2: otk}}, [(local_user, device_id_1, "alg1", 1)])
         )
 
         # we shouldn't have any unused fallback keys yet
@@ -1079,11 +1079,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # query the fallback keys.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {
-                    "one_time_keys": {
-                        local_user: {device_id_1: "alg1", device_id_2: "alg1"}
-                    }
-                },
+                {local_user: {device_id_1: {"alg1": 1}, device_id_2: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=False,
             )
@@ -1128,7 +1124,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # Claim OTKs, which will ask the appservice and do nothing else.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id_1: "alg1"}}},
+                {local_user: {device_id_1: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )
@@ -1172,7 +1168,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # uploaded fallback key.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id_1: "alg1"}}},
+                {local_user: {device_id_1: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )
@@ -1205,7 +1201,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # Claim OTKs, which will return information only from the database.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id_1: "alg1"}}},
+                {local_user: {device_id_1: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )
@@ -1232,7 +1228,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
         # Claim OTKs, which will return only the fallback key from the database.
         claim_res = self.get_success(
             self.handler.claim_one_time_keys(
-                {"one_time_keys": {local_user: {device_id_1: "alg1"}}},
+                {local_user: {device_id_1: {"alg1": 1}}},
                 timeout=None,
                 always_include_fallback_keys=True,
             )