Add support for setting PNI-associated registration IDs and identity keys when changing numbers

2026-04-22 09:18:05 +01:00 · 2022-07-26 15:19:27 -04:00
parent c252118cfc
commit dce391a248
26 changed files with 927 additions and 673 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
@@ -68,13 +68,11 @@ import org.whispersystems.textsecuregcm.entities.ApnRegistrationId;
 import org.whispersystems.textsecuregcm.entities.ChangePhoneNumberRequest;
 import org.whispersystems.textsecuregcm.entities.DeviceName;
 import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
-import org.whispersystems.textsecuregcm.entities.IncomingMessage;
 import org.whispersystems.textsecuregcm.entities.MismatchedDevices;
 import org.whispersystems.textsecuregcm.entities.RegistrationLock;
 import org.whispersystems.textsecuregcm.entities.RegistrationLockFailure;
 import org.whispersystems.textsecuregcm.entities.StaleDevices;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
-import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.APNSender;
 import org.whispersystems.textsecuregcm.push.ApnMessage;
@@ -95,7 +93,6 @@ import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.ForwardedIpUtil;
 import org.whispersystems.textsecuregcm.util.Hex;
 import org.whispersystems.textsecuregcm.util.ImpossiblePhoneNumberException;
-import org.whispersystems.textsecuregcm.util.MessageValidation;
 import org.whispersystems.textsecuregcm.util.NonNormalizedPhoneNumberException;
 import org.whispersystems.textsecuregcm.util.Username;
 import org.whispersystems.textsecuregcm.util.Util;
@@ -416,41 +413,9 @@ 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());
-      }
-      try {
-        // Checks that all except master ID are in device messages
-        MessageValidation.validateCompleteDeviceList(
-            authenticatedAccount.getAccount(), request.getDeviceMessages(),
-            IncomingMessage::getDestinationDeviceId, true, Optional.of(Device.MASTER_ID));
-        MessageValidation.validateRegistrationIds(
-            authenticatedAccount.getAccount(), request.getDeviceMessages(),
-            IncomingMessage::getDestinationDeviceId, IncomingMessage::getDestinationRegistrationId);
-        // Checks that all including master ID are in signed prekeys
-        MessageValidation.validateCompleteDeviceList(
-            authenticatedAccount.getAccount(), request.getDeviceSignedPrekeys().entrySet(),
-            e -> e.getKey(), false, Optional.empty());
-      } catch (MismatchedDevicesException e) {
-        throw new WebApplicationException(Response.status(409)
-            .type(MediaType.APPLICATION_JSON_TYPE)
-            .entity(new MismatchedDevices(e.getMissingDevices(),
-                e.getExtraDevices()))
-            .build());
-      } catch (StaleDevicesException e) {
-        throw new WebApplicationException(Response.status(410)
-            .type(MediaType.APPLICATION_JSON)
-            .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.number();

-    final String number = request.getNumber();
+    // Only "bill" for rate limiting if we think there's a change to be made...
    if (!authenticatedAccount.getAccount().getNumber().equals(number)) {
      Util.requireNormalizedNumber(number);

@@ -459,7 +424,7 @@ public class AccountController {
      final Optional<StoredVerificationCode> storedVerificationCode =
          pendingAccounts.getCodeForNumber(number);

-      if (storedVerificationCode.isEmpty() || !storedVerificationCode.get().isValid(request.getCode())) {
+      if (storedVerificationCode.isEmpty() || !storedVerificationCode.get().isValid(request.code())) {
        throw new ForbiddenException();
      }

@@ -469,24 +434,42 @@ public class AccountController {
      final Optional<Account> existingAccount = accounts.getByE164(number);

      if (existingAccount.isPresent()) {
-        verifyRegistrationLock(existingAccount.get(), request.getRegistrationLock());
+        verifyRegistrationLock(existingAccount.get(), request.registrationLock());
      }

      rateLimiters.getVerifyLimiter().clear(number);
    }

-    final Account updatedAccount = changeNumberManager.changeNumber(
-        authenticatedAccount.getAccount(),
-        request.getNumber(),
-        Optional.ofNullable(request.getDeviceSignedPrekeys()).orElse(Collections.emptyMap()),
-        Optional.ofNullable(request.getDeviceMessages()).orElse(Collections.emptyList()));
+    // ...but always attempt to make the change in case a client retries and needs to re-send messages
+    try {
+      final Account updatedAccount = changeNumberManager.changeNumber(
+          authenticatedAccount.getAccount(),
+          request.number(),
+          request.pniIdentityKey(),
+          Optional.ofNullable(request.devicePniSignedPrekeys()).orElse(Collections.emptyMap()),
+          Optional.ofNullable(request.deviceMessages()).orElse(Collections.emptyList()),
+          Optional.ofNullable(request.pniRegistrationIds()).orElse(Collections.emptyMap()));

-    return new AccountIdentityResponse(
-        updatedAccount.getUuid(),
-        updatedAccount.getNumber(),
-        updatedAccount.getPhoneNumberIdentifier(),
-        updatedAccount.getUsername().orElse(null),
-        updatedAccount.isStorageSupported());
+      return new AccountIdentityResponse(
+          updatedAccount.getUuid(),
+          updatedAccount.getNumber(),
+          updatedAccount.getPhoneNumberIdentifier(),
+          updatedAccount.getUsername().orElse(null),
+          updatedAccount.isStorageSupported());
+    } catch (MismatchedDevicesException e) {
+      throw new WebApplicationException(Response.status(409)
+          .type(MediaType.APPLICATION_JSON_TYPE)
+          .entity(new MismatchedDevices(e.getMissingDevices(),
+              e.getExtraDevices()))
+          .build());
+    } catch (StaleDevicesException e) {
+      throw new WebApplicationException(Response.status(410)
+          .type(MediaType.APPLICATION_JSON)
+          .entity(new StaleDevices(e.getStaleDevices()))
+          .build());
+    } catch (IllegalArgumentException e) {
+      throw new BadRequestException(e);
+    }
  }

  @Timed
@@ -625,6 +608,7 @@ public class AccountController {
        d.setLastSeen(Util.todayInMillis());
        d.setCapabilities(attributes.getCapabilities());
        d.setRegistrationId(attributes.getRegistrationId());
+        attributes.getPhoneNumberIdentityRegistrationId().ifPresent(d::setPhoneNumberIdentityRegistrationId);
        d.setUserAgent(userAgent);
      });

--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
@@ -198,6 +198,7 @@ public class DeviceController {
    device.setAuthenticationCredentials(new AuthenticationCredentials(password));
    device.setFetchesMessages(accountAttributes.getFetchesMessages());
    device.setRegistrationId(accountAttributes.getRegistrationId());
+    accountAttributes.getPhoneNumberIdentityRegistrationId().ifPresent(device::setPhoneNumberIdentityRegistrationId);
    device.setLastSeen(Util.todayInMillis());
    device.setCreated(System.currentTimeMillis());
    device.setCapabilities(accountAttributes.getCapabilities());
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -197,7 +197,11 @@ public class KeysController {
        PreKey       preKey       = preKeysByDeviceId.get(device.getId());

        if (signedPreKey != null || preKey != null) {
-          responseItems.add(new PreKeyResponseItem(device.getId(), device.getRegistrationId(), signedPreKey, preKey));
+          final int registrationId = usePhoneNumberIdentity ?
+              device.getPhoneNumberIdentityRegistrationId().orElse(device.getRegistrationId()) :
+              device.getRegistrationId();
+
+          responseItems.add(new PreKeyResponseItem(device.getId(), registrationId, signedPreKey, preKey));
        }
      }
    }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -21,7 +21,7 @@ import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -89,8 +89,7 @@ import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
-import org.whispersystems.textsecuregcm.util.MessageValidation;
-import org.whispersystems.textsecuregcm.util.Pair;
+import org.whispersystems.textsecuregcm.util.DestinationDeviceValidator;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
 import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
@@ -214,11 +213,23 @@ public class MessageController {
        checkRateLimit(source.get(), destination.get(), userAgent);
      }

-      MessageValidation.validateCompleteDeviceList(destination.get(), messages.getMessages(),
-          IncomingMessage::getDestinationDeviceId, isSyncMessage,
-          source.map(AuthenticatedAccount::getAuthenticatedDevice).map(Device::getId));
-      MessageValidation.validateRegistrationIds(destination.get(), messages.getMessages(),
-          IncomingMessage::getDestinationDeviceId, IncomingMessage::getDestinationRegistrationId);
+      final Set<Long> excludedDeviceIds;
+
+      if (isSyncMessage) {
+        excludedDeviceIds = Set.of(source.get().getAuthenticatedDevice().getId());
+      } else {
+        excludedDeviceIds = Collections.emptySet();
+      }
+
+      DestinationDeviceValidator.validateCompleteDeviceList(destination.get(),
+          messages.getMessages().stream().map(IncomingMessage::getDestinationDeviceId).collect(Collectors.toSet()),
+          excludedDeviceIds);
+
+      DestinationDeviceValidator.validateRegistrationIds(destination.get(),
+          messages.getMessages().stream().collect(Collectors.toMap(
+              IncomingMessage::getDestinationDeviceId,
+              IncomingMessage::getDestinationRegistrationId)),
+          destination.get().getPhoneNumberIdentifier().equals(destinationUuid));

      final List<Tag> tags = List.of(UserAgentTagUtil.getPlatformTag(userAgent),
          Tag.of(EPHEMERAL_TAG_NAME, String.valueOf(messages.isOnline())),
@@ -307,13 +318,25 @@ public class MessageController {
        checkRateLimit(source.get(), destination.get(), userAgent);
      }

-      final List<IncomingDeviceMessage> messagesAsList = Arrays.asList(messages);
-      MessageValidation.validateCompleteDeviceList(destination.get(), messagesAsList,
-          IncomingDeviceMessage::getDeviceId, isSyncMessage,
-          source.map(AuthenticatedAccount::getAuthenticatedDevice).map(Device::getId));
-      MessageValidation.validateRegistrationIds(destination.get(), messagesAsList,
-          IncomingDeviceMessage::getDeviceId,
-          IncomingDeviceMessage::getRegistrationId);
+      final Set<Long> excludedDeviceIds;
+
+      if (isSyncMessage) {
+        excludedDeviceIds = Set.of(source.get().getAuthenticatedDevice().getId());
+      } else {
+        excludedDeviceIds = Collections.emptySet();
+      }
+
+      DestinationDeviceValidator.validateCompleteDeviceList(
+          destination.get(),
+          Arrays.stream(messages).map(IncomingDeviceMessage::getDeviceId).collect(Collectors.toSet()),
+          excludedDeviceIds);
+
+      DestinationDeviceValidator.validateRegistrationIds(
+          destination.get(),
+          Arrays.stream(messages).collect(Collectors.toMap(
+              IncomingDeviceMessage::getDeviceId,
+              IncomingDeviceMessage::getRegistrationId)),
+          destination.get().getPhoneNumberIdentifier().equals(destinationUuid));

      final List<Tag> tags = List.of(UserAgentTagUtil.getPlatformTag(userAgent),
          Tag.of(EPHEMERAL_TAG_NAME, String.valueOf(online)),
@@ -372,27 +395,29 @@ public class MessageController {
        }));
    checkAccessKeys(accessKeys, uuidToAccountMap);

-    final Map<Account, HashSet<Pair<Long, Integer>>> accountToDeviceIdAndRegistrationIdMap =
-        Arrays
-            .stream(multiRecipientMessage.getRecipients())
-            .collect(Collectors.toMap(
-                recipient -> uuidToAccountMap.get(recipient.getUuid()),
-                recipient -> new HashSet<>(
-                    Collections.singletonList(new Pair<>(recipient.getDeviceId(), recipient.getRegistrationId()))),
-                (a, b) -> {
-                  a.addAll(b);
-                  return a;
-                }
-            ));
+    final Map<Account, Map<Long, Integer>> accountToDeviceIdAndRegistrationIdMap = Arrays.stream(multiRecipientMessage.getRecipients())
+        .collect(Collectors.toMap(
+            recipient -> uuidToAccountMap.get(recipient.getUuid()),
+            recipient -> Map.of(recipient.getDeviceId(), recipient.getRegistrationId()),
+            (a, b) -> {
+              final Map<Long, Integer> combined = new HashMap<>();
+              combined.putAll(a);
+              combined.putAll(b);
+
+              return combined;
+            }
+        ));

    Collection<AccountMismatchedDevices> accountMismatchedDevices = new ArrayList<>();
    Collection<AccountStaleDevices> accountStaleDevices = new ArrayList<>();
    uuidToAccountMap.values().forEach(account -> {
-      final Set<Pair<Long, Integer>> deviceIdAndRegistrationIdSet = accountToDeviceIdAndRegistrationIdMap.get(account);
-      final Set<Long> deviceIds = deviceIdAndRegistrationIdSet.stream().map(Pair::first).collect(Collectors.toSet());
+      final Set<Long> deviceIds = accountToDeviceIdAndRegistrationIdMap.get(account).keySet();
      try {
-        MessageValidation.validateCompleteDeviceList(account, deviceIds, false, Optional.empty());
-        MessageValidation.validateRegistrationIds(account, deviceIdAndRegistrationIdSet.stream());
+        DestinationDeviceValidator.validateCompleteDeviceList(account, deviceIds, Collections.emptySet());
+
+        // Multi-recipient messages are always sealed-sender messages, and so can never be sent to a phone number
+        // identity
+        DestinationDeviceValidator.validateRegistrationIds(account, accountToDeviceIdAndRegistrationIdMap.get(account), false);
      } catch (MismatchedDevicesException e) {
        accountMismatchedDevices.add(new AccountMismatchedDevices(account.getUuid(),
            new MismatchedDevices(e.getMissingDevices(), e.getExtraDevices())));