Update ChangeNumber to allow reset of registration IDs.

service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java

@@ -25,6 +25,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.servlet.http.HttpServletRequest;
 import javax.validation.Valid;
@@ -86,6 +87,7 @@ import org.whispersystems.textsecuregcm.storage.AbusiveHostRules;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
+import org.whispersystems.textsecuregcm.storage.ChangeNumberManager.DeviceUpdate;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
@@ -408,22 +410,27 @@ public class AccountController {
       throw new ForbiddenException();
     }
 
-    if (request.getDeviceSignedPrekeys() != null && !request.getDeviceSignedPrekeys().isEmpty()) {
-      if (request.getDeviceMessages() == null || request.getDeviceMessages().size() != request.getDeviceSignedPrekeys().size() - 1) {
-        // device_messages should exist and be one shorter than device_signed_prekeys, since it doesn't have the primary's key.
-        throw new WebApplicationException(Response.status(400).build());
-      }
+    Map<Long, DeviceUpdate> devices = Collections.emptyMap();
+    if (request.getDeviceUpdates() != null && !request.getDeviceUpdates().isEmpty()) {
+      devices = request.getDeviceUpdates().entrySet().stream()
+          .collect(Collectors.toMap(
+              e -> e.getKey(),
+              e -> new DeviceUpdate(e.getValue().getSignedPhoneNumberIdentityPreKey(), e.getValue().getMessage(), e.getValue().getRegistrationID())));
       try {
         // Checks that all except master ID are in device messages
+        List<IncomingMessage> deviceMessages = devices.entrySet().stream()
+            .map(e -> e.getValue().message())
+            .filter(e -> e != null)
+            .collect(Collectors.toList());
         MessageValidation.validateCompleteDeviceList(
-            authenticatedAccount.getAccount(), request.getDeviceMessages(),
+            authenticatedAccount.getAccount(), deviceMessages,
             IncomingMessage::getDestinationDeviceId, true, Optional.of(Device.MASTER_ID));
         MessageValidation.validateRegistrationIds(
-            authenticatedAccount.getAccount(), request.getDeviceMessages(),
+            authenticatedAccount.getAccount(), deviceMessages,
             IncomingMessage::getDestinationDeviceId, IncomingMessage::getDestinationRegistrationId);
         // Checks that all including master ID are in signed prekeys
         MessageValidation.validateCompleteDeviceList(
-            authenticatedAccount.getAccount(), request.getDeviceSignedPrekeys().entrySet(),
+            authenticatedAccount.getAccount(), devices.entrySet(),
             e -> e.getKey(), false, Optional.empty());
       } catch (MismatchedDevicesException e) {
         throw new WebApplicationException(Response.status(409)
@@ -437,9 +444,6 @@ public class AccountController {
             .entity(new StaleDevices(e.getStaleDevices()))
             .build());
       }
-    } else if (request.getDeviceMessages() != null && !request.getDeviceMessages().isEmpty()) {
-      // device_messages shouldn't exist without device_signed_prekeys.
-      throw new WebApplicationException(Response.status(400).build());
     }
 
     final String number = request.getNumber();
@@ -470,8 +474,7 @@ public class AccountController {
     final Account updatedAccount = changeNumberManager.changeNumber(
         authenticatedAccount.getAccount(),
         request.getNumber(),
-        Optional.ofNullable(request.getDeviceSignedPrekeys()).orElse(Collections.emptyMap()),
-        Optional.ofNullable(request.getDeviceMessages()).orElse(Collections.emptyList()));
+        devices);
 
     return new AccountIdentityResponse(
         updatedAccount.getUuid(),

service/src/main/java/org/whispersystems/textsecuregcm/entities/ChangePhoneNumberRequest.java

@@ -9,7 +9,6 @@ import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import javax.annotation.Nullable;
 import javax.validation.constraints.NotBlank;
-import java.util.List;
 import java.util.Map;
 
 public class ChangePhoneNumberRequest {
@@ -26,26 +25,48 @@ public class ChangePhoneNumberRequest {
   @Nullable
   final String registrationLock;
 
-  @JsonProperty("device_messages")
+  @JsonProperty("deviceUpdates")
   @Nullable
-  final List<IncomingMessage> deviceMessages;
+  final Map<Long, DeviceUpdate> deviceUpdates;
 
-  @JsonProperty("device_signed_prekeys")
-  @Nullable
-  final Map<Long, SignedPreKey> deviceSignedPrekeys;
+  public static class DeviceUpdate {
+
+    private final IncomingMessage message;
+    private final SignedPreKey signedPhoneNumberIdentityPreKey;
+    private final Integer registrationID;
+
+    @JsonCreator
+    public DeviceUpdate(
+        @JsonProperty("message") final IncomingMessage message,
+        @JsonProperty("signedPhoneNumberIdentityPrekey") final SignedPreKey signedPhoneNumberIdentityPreKey,
+        @JsonProperty("registratonId") final Integer registrationID) {
+      this.message = message;
+      this.signedPhoneNumberIdentityPreKey = signedPhoneNumberIdentityPreKey;
+      this.registrationID = registrationID;
+    }
+
+    public IncomingMessage getMessage() {
+      return message;
+    }
+
+    public SignedPreKey getSignedPhoneNumberIdentityPreKey() {
+      return signedPhoneNumberIdentityPreKey;
+    }
+
+    public Integer getRegistrationID() {
+      return registrationID;
+    }
+  }
 
   @JsonCreator
   public ChangePhoneNumberRequest(@JsonProperty("number") final String number,
       @JsonProperty("code") final String code,
       @JsonProperty("reglock") @Nullable final String registrationLock,
-      @JsonProperty("device_messages") @Nullable final List<IncomingMessage> deviceMessages,
-      @JsonProperty("device_signed_prekeys") @Nullable final Map<Long, SignedPreKey> deviceSignedPrekeys) {
-
+      @JsonProperty("deviceUpdates") @Nullable final Map<Long, DeviceUpdate> deviceUpdates) {
     this.number = number;
     this.code = code;
     this.registrationLock = registrationLock;
-    this.deviceMessages = deviceMessages;
-    this.deviceSignedPrekeys = deviceSignedPrekeys;
+    this.deviceUpdates = deviceUpdates;
   }
 
   public String getNumber() {
@@ -62,12 +83,5 @@ public class ChangePhoneNumberRequest {
   }
 
   @Nullable
-  public List<IncomingMessage> getDeviceMessages() {
-    return deviceMessages;
-  }
-
-  @Nullable
-  public Map<Long, SignedPreKey> getDeviceSignedPrekeys() {
-    return deviceSignedPrekeys;
-  }
+  public Map<Long, DeviceUpdate> getDeviceUpdates() { return deviceUpdates; }
 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManager.java

@@ -16,7 +16,6 @@ import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.push.MessageSender;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
 import javax.validation.constraints.NotNull;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -32,11 +31,13 @@ public class ChangeNumberManager {
     this.accountsManager = accountsManager;
   }
 
+  public record DeviceUpdate(SignedPreKey signedPhoneNumberIdentityPreKey, IncomingMessage message, Integer registrationID) {
+  }
+
   public Account changeNumber(
       @NotNull Account account,
       @NotNull final String number,
-      @NotNull final Map<Long, SignedPreKey> deviceSignedPrekeys,
-      @NotNull final List<IncomingMessage> deviceMessages) throws InterruptedException {
+      @NotNull final Map<Long, DeviceUpdate> deviceUpdates) throws InterruptedException {
 
     final Account updatedAccount;
     if (number.equals(account.getNumber())) {
@@ -51,14 +52,17 @@ public class ChangeNumberManager {
     // This makes it so the client can resend a request they didn't get a response for (timeout, etc)
     // to make sure their messages sent and prekeys were updated, even if the first time around the
     // server crashed at/above this point.
-    if (deviceSignedPrekeys != null && !deviceSignedPrekeys.isEmpty()) {
-      for (Map.Entry<Long, SignedPreKey> entry : deviceSignedPrekeys.entrySet()) {
+    if (deviceUpdates != null && !deviceUpdates.isEmpty()) {
+      for (Map.Entry<Long, DeviceUpdate> entry : deviceUpdates.entrySet()) {
+        DeviceUpdate deviceUpdate = entry.getValue();
         accountsManager.updateDevice(updatedAccount, entry.getKey(),
-            d -> d.setPhoneNumberIdentitySignedPreKey(entry.getValue()));
-      }
-
-      for (IncomingMessage message : deviceMessages) {
-        sendMessageToSelf(updatedAccount, updatedAccount.getDevice(message.getDestinationDeviceId()), message);
+            device -> {
+              if (deviceUpdate.signedPhoneNumberIdentityPreKey() != null) device.setPhoneNumberIdentitySignedPreKey(deviceUpdate.signedPhoneNumberIdentityPreKey());
+              if (deviceUpdate.registrationID() != null) device.setRegistrationId(deviceUpdate.registrationID());
+            });
+        if (deviceUpdate.message() != null) {
+          sendMessageToSelf(updatedAccount, updatedAccount.getDevice(entry.getKey()), deviceUpdate.message());
+        }
       }
     }
     return updatedAccount;