From 1167d0ac2e863ea750ff2d9d7cf136ed11fd6891 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 13 Dec 2023 17:32:56 -0500
Subject: [PATCH] Make key-setting methods synchronous

---
 .../textsecuregcm/controllers/KeysController.java   | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
index bfe530a68..d05f52cad 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -60,7 +60,6 @@
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.KeysManager;
-import org.whispersystems.textsecuregcm.util.Util;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -110,7 +109,7 @@ public CompletableFuture<PreKeyCount> getStatus(@Auth final AuthenticatedAccount
   @ApiResponse(responseCode = "401", description = "Account authentication check failed.")
   @ApiResponse(responseCode = "403", description = "Attempt to change identity key from a non-primary device.")
   @ApiResponse(responseCode = "422", description = "Invalid request format.")
-  public CompletableFuture<Response> setKeys(@Auth final DisabledPermittedAuthenticatedAccount disabledPermittedAuth,
+  public void setKeys(@Auth final DisabledPermittedAuthenticatedAccount disabledPermittedAuth,
       @RequestBody @NotNull @Valid final SetKeysRequest setKeysRequest,
 
       @Parameter(allowEmptyValue=true)
@@ -147,7 +146,7 @@ public CompletableFuture<Response> setKeys(@Auth final DisabledPermittedAuthenti
       updateAccountFuture = CompletableFuture.completedFuture(account);
     }
 
-    return updateAccountFuture.thenCompose(updatedAccount -> {
+    updateAccountFuture.thenCompose(updatedAccount -> {
           final List<CompletableFuture<Void>> storeFutures = new ArrayList<>(3);
 
           if (setKeysRequest.preKeys() != null && !setKeysRequest.preKeys().isEmpty()) {
@@ -165,7 +164,7 @@ public CompletableFuture<Response> setKeys(@Auth final DisabledPermittedAuthenti
 
           return CompletableFuture.allOf(storeFutures.toArray(EMPTY_FUTURE_ARRAY));
         })
-        .thenApply(Util.ASYNC_EMPTY_RESPONSE);
+        .join();
   }
 
   private void checkSignedPreKeySignatures(final SetKeysRequest setKeysRequest, final IdentityKey identityKey) {
@@ -292,14 +291,14 @@ public PreKeyResponse getDeviceKeys(@Auth Optional<AuthenticatedAccount> auth,
   @ApiResponse(responseCode = "200", description = "Indicates that new prekey was successfully stored.")
   @ApiResponse(responseCode = "401", description = "Account authentication check failed.")
   @ApiResponse(responseCode = "422", description = "Invalid request format.")
-  public CompletableFuture<Response> setSignedKey(@Auth final AuthenticatedAccount auth,
+  public void setSignedKey(@Auth final AuthenticatedAccount auth,
       @Valid final ECSignedPreKey signedPreKey,
       @QueryParam("identity") @DefaultValue("aci") final IdentityType identityType) {
 
     final UUID identifier = auth.getAccount().getIdentifier(identityType);
     final byte deviceId = auth.getAuthenticatedDevice().getId();
 
-    return accounts.updateDeviceTransactionallyAsync(auth.getAccount(),
+    accounts.updateDeviceTransactionallyAsync(auth.getAccount(),
             deviceId,
             d -> {
               switch (identityType) {
@@ -311,7 +310,7 @@ public CompletableFuture<Response> setSignedKey(@Auth final AuthenticatedAccount
                 .map(List::of)
                 .orElseGet(Collections::emptyList))
         .toCompletableFuture()
-        .thenApply(Util.ASYNC_EMPTY_RESPONSE);
+        .join();
   }
 
   private List<Device> parseDeviceId(String deviceId, Account account) {