create device and store keys in same call

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/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}