Support MSC3814: Dehydrated Devices Part 2

changelog.d/16010.misc

@@ -0,0 +1 @@
+Update dehydrated devices implementation.

synapse/handlers/device.py

@@ -41,6 +41,7 @@
     run_as_background_process,
     wrap_as_background_process,
 )
+from synapse.replication.http.devices import ReplicationUploadKeysForUserRestServlet
 from synapse.types import (
     JsonDict,
     StrCollection,
@@ -656,15 +657,17 @@ async def store_dehydrated_device(
         device_id: Optional[str],
         device_data: JsonDict,
         initial_device_display_name: Optional[str] = None,
+        device_keys: Optional[JsonDict] = None,
     ) -> str:
-        """Store a dehydrated device for a user.  If the user had a previous
-        dehydrated device, it is removed.
+        """Store a dehydrated device for a user, optionally storing the keys associated with
+        it as well.  If the user had a previous dehydrated device, it is removed.
 
         Args:
             user_id: the user that we are storing the device for
             device_id: device id supplied by client
             device_data: the dehydrated device information
             initial_device_display_name: The display name to use for the device
+            device_keys: keys for the dehydrated device
         Returns:
             device id of the dehydrated device
         """
@@ -678,6 +681,18 @@ async def store_dehydrated_device(
         )
         if old_device_id is not None:
             await self.delete_devices(user_id, [old_device_id])
+
+        # we do this here to avoid a circular import
+        if self.hs.config.worker.worker_app is None:
+            # if main process
+            key_uploader = self.hs.get_e2e_keys_handler().upload_keys_for_user
+        else:
+            # if worker process
+            key_uploader = ReplicationUploadKeysForUserRestServlet.make_client(self.hs)
+
+        # if keys are provided store them
+        if device_keys:
+            await key_uploader(user_id=user_id, device_id=device_id, keys=device_keys)
         return device_id
 
     async def rehydrate_device(

synapse/handlers/devicemessage.py

@@ -367,19 +367,6 @@ async def get_events_for_dehydrated_device(
                     errcode=Codes.INVALID_PARAM,
                 )
 
-            # if we have a since token, delete any to-device messages before that token
-            # (since we now know that the device has received them)
-            deleted = await self.store.delete_messages_for_device(
-                user_id, device_id, since_stream_id
-            )
-            logger.debug(
-                "Deleted %d to-device messages up to %d for user_id %s device_id %s",
-                deleted,
-                since_stream_id,
-                user_id,
-                device_id,
-            )
-
         to_token = self.event_sources.get_current_token().to_device_key
 
         messages, stream_id = await self.store.get_messages_for_device(

synapse/rest/client/devices.py

@@ -536,7 +536,6 @@ class Config:
     async def on_PUT(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
         submission = parse_and_validate_json_object_from_request(request, self.PutBody)
         requester = await self.auth.get_user_by_req(request)
-        user_id = requester.user.to_string()
 
         device_info = submission.dict()
         if "device_keys" not in device_info.keys():
@@ -545,18 +544,12 @@ async def on_PUT(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
                 "Device key(s) not found, these must be provided.",
             )
 
-        # TODO: Those two operations, creating a device and storing the
-        # device's keys should be atomic.
         device_id = await self.device_handler.store_dehydrated_device(
             requester.user.to_string(),
             submission.device_id,
             submission.device_data.dict(),
             submission.initial_device_display_name,
-        )
-
-        # TODO: Do we need to do something with the result here?
-        await self.key_uploader(
-            user_id=user_id, device_id=submission.device_id, keys=submission.dict()
+            device_info,
         )
 
         return 200, {"device_id": device_id}

tests/handlers/test_device.py

@@ -566,15 +566,16 @@ def test_dehydrate_v2_and_fetch_events(self) -> None:
         self.assertEqual(len(res["events"]), 1)
         self.assertEqual(res["events"][0]["content"]["body"], "foo")
 
-        # Fetch the message of the dehydrated device again, which should return nothing
-        # and delete the old messages
+        # Fetch the message of the dehydrated device again, which should return
+        # the same message as it has not been deleted
         res = self.get_success(
             self.message_handler.get_events_for_dehydrated_device(
                 requester=requester,
                 device_id=stored_dehydrated_device_id,
-                since_token=res["next_batch"],
+                since_token=None,
                 limit=10,
             )
         )
         self.assertTrue(len(res["next_batch"]) > 1)
-        self.assertEqual(len(res["events"]), 0)
+        self.assertEqual(len(res["events"]), 1)
+        self.assertEqual(res["events"][0]["content"]["body"], "foo")

tests/rest/client/test_devices.py

@@ -312,6 +312,23 @@ def test_dehydrate_msc3814(self) -> None:
         }
         self.assertEqual(device_data, expected_device_data)
 
+        # test that the keys are correctly uploaded
+        channel = self.make_request(
+            "POST",
+            "/_matrix/client/r0/keys/query",
+            {
+                "device_keys": {
+                    user: ["device1"],
+                },
+            },
+            token,
+        )
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(
+            channel.json_body["device_keys"][user][device_id]["keys"],
+            content["device_keys"]["keys"],
+        )
+
         # create another device for the user
         (
             new_device_id,
@@ -348,10 +365,21 @@ def test_dehydrate_msc3814(self) -> None:
         self.assertEqual(channel.code, 200)
         expected_content = {"body": "test_message"}
         self.assertEqual(channel.json_body["events"][0]["content"], expected_content)
+
+        # fetch messages again and make sure that the message was not deleted
+        channel = self.make_request(
+            "POST",
+            f"_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device/{device_id}/events",
+            content={},
+            access_token=token,
+            shorthand=False,
+        )
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(channel.json_body["events"][0]["content"], expected_content)
         next_batch_token = channel.json_body.get("next_batch")
 
-        # fetch messages again and make sure that the message was deleted and we are returned an
-        # empty array
+        # make sure fetching messages with next batch token works - there are no unfetched
+        # messages so we should receive an empty array
         content = {"next_batch": next_batch_token}
         channel = self.make_request(
             "POST",