new /v2/accounts endpoint to distribute PNI key material without changing phone number

2026-04-21 05:58:00 +01:00 · 2023-04-21 12:20:57 -07:00
parent 4fb89360ce
commit 47ad5779ad
10 changed files with 658 additions and 68 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountControllerV2.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountControllerV2.java
@@ -40,6 +40,7 @@ import org.whispersystems.textsecuregcm.entities.ChangeNumberRequest;
 import org.whispersystems.textsecuregcm.entities.MismatchedDevices;
 import org.whispersystems.textsecuregcm.entities.PhoneNumberDiscoverabilityRequest;
 import org.whispersystems.textsecuregcm.entities.PhoneVerificationRequest;
+import org.whispersystems.textsecuregcm.entities.PhoneNumberIdentityKeyDistributionRequest;
 import org.whispersystems.textsecuregcm.entities.StaleDevices;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -138,6 +139,54 @@ public class AccountControllerV2 {
    }
  }

+  @Timed
+  @PUT
+  @Path("/phone_number_identity_key_distribution")
+  @Consumes(MediaType.APPLICATION_JSON)
+  @Produces(MediaType.APPLICATION_JSON)
+  @Operation(summary = "Updates key material for the phone-number identity for all devices and sends a synchronization message to companion devices")
+  public AccountIdentityResponse distributePhoneNumberIdentityKeys(@Auth final AuthenticatedAccount authenticatedAccount,
+      @NotNull @Valid final PhoneNumberIdentityKeyDistributionRequest request) {
+
+    if (!authenticatedAccount.getAuthenticatedDevice().isMaster()) {
+      throw new ForbiddenException();
+    }
+
+    final Account account = authenticatedAccount.getAccount();
+    if (!account.isPniSupported()) {
+      throw new WebApplicationException(Response.status(425).build());
+    }
+
+    try {
+      final Account updatedAccount = changeNumberManager.updatePNIKeys(
+          authenticatedAccount.getAccount(),
+          request.pniIdentityKey(),
+          request.devicePniSignedPrekeys(),
+          request.deviceMessages(),
+          request.pniRegistrationIds());
+
+      return new AccountIdentityResponse(
+          updatedAccount.getUuid(),
+          updatedAccount.getNumber(),
+          updatedAccount.getPhoneNumberIdentifier(),
+          updatedAccount.getUsernameHash().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
  @PUT
  @Path("/phone_number_discoverability")
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountIdentityResponse.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountIdentityResponse.java
@@ -8,14 +8,25 @@ package org.whispersystems.textsecuregcm.entities;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import org.whispersystems.textsecuregcm.util.ByteArrayBase64UrlAdapter;
+import io.swagger.v3.oas.annotations.media.Schema;
 import java.util.UUID;
 import javax.annotation.Nullable;

-public record AccountIdentityResponse(UUID uuid,
-                                      String number,
-                                      UUID pni,
-                                      @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
-                                      @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
-                                      @Nullable byte[] usernameHash,
-                                      boolean storageCapable) {
+public record AccountIdentityResponse(
+    @Schema(description="the account identifier for this account")
+     UUID uuid,
+
+    @Schema(description="the phone number associated with this account")
+    String number,
+
+    @Schema(description="the account identifier for this account's phone-number identity")
+    UUID pni,
+
+    @Schema(description="a hash of this account's username, if set")
+    @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
+    @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
+    @Nullable byte[] usernameHash,
+
+    @Schema(description="whether any of this account's devices support storage")
+    boolean storageCapable) {
 }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/PhoneNumberIdentityKeyDistributionRequest.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/PhoneNumberIdentityKeyDistributionRequest.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.entities;
+
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import io.swagger.v3.oas.annotations.media.Schema;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import javax.validation.Valid;
+import javax.validation.constraints.NotBlank;
+import javax.validation.constraints.NotNull;
+import org.whispersystems.textsecuregcm.util.ByteArrayAdapter;
+
+public record PhoneNumberIdentityKeyDistributionRequest(
+        @NotBlank
+        @Schema(description="the new identity key for this account's phone-number identity")
+        String pniIdentityKey,
+    
+        @NotNull
+        @Valid
+        @Schema(description="A message for each companion device to pass its new private keys")
+        List<@NotNull @Valid IncomingMessage> deviceMessages,
+    
+        @NotNull
+        @Valid
+        @Schema(description="The public key of a new signed elliptic-curve prekey pair for each device")
+        Map<Long, @NotNull @Valid SignedPreKey> devicePniSignedPrekeys,
+    
+        @NotNull
+        @Valid
+        @Schema(description="The new registration ID to use for the phone-number identity of each device")
+        Map<Long, Integer> pniRegistrationIds) {
+}
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
@@ -28,13 +28,17 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+
+import org.apache.commons.lang3.ObjectUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
@@ -255,24 +259,13 @@ public class AccountsManager {
    final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier();

    if (originalNumber.equals(number)) {
+      if (pniIdentityKey != null) {
+        throw new IllegalArgumentException("change number must supply a changed phone number; otherwise use updatePNIKeys");
+      }
      return account;
    }

-    if (pniSignedPreKeys != null && pniRegistrationIds != null) {
-      // Check that all including master ID are in signed pre-keys
-      DestinationDeviceValidator.validateCompleteDeviceList(
-          account,
-          pniSignedPreKeys.keySet(),
-          Collections.emptySet());
-
-      // Check that all devices are accounted for in the map of new PNI registration IDs
-      DestinationDeviceValidator.validateCompleteDeviceList(
-          account,
-          pniRegistrationIds.keySet(),
-          Collections.emptySet());
-    } else if (pniSignedPreKeys != null || pniRegistrationIds != null) {
-      throw new IllegalArgumentException("Signed pre-keys and registration IDs must both be null or both be non-null");
-    }
+    validateDevices(account, pniSignedPreKeys, pniRegistrationIds);

    final AtomicReference<Account> updatedAccount = new AtomicReference<>();

@@ -297,22 +290,7 @@ public class AccountsManager {

      numberChangedAccount = updateWithRetries(
          account,
-          a -> {
-            //noinspection ConstantConditions
-            if (pniSignedPreKeys != null && pniRegistrationIds != null) {
-              pniSignedPreKeys.forEach((deviceId, signedPreKey) ->
-                  a.getDevice(deviceId).ifPresent(device -> device.setPhoneNumberIdentitySignedPreKey(signedPreKey)));
-
-              pniRegistrationIds.forEach((deviceId, registrationId) ->
-                  a.getDevice(deviceId).ifPresent(device -> device.setPhoneNumberIdentityRegistrationId(registrationId)));
-            }
-
-            if (pniIdentityKey != null) {
-              a.setPhoneNumberIdentityKey(pniIdentityKey);
-            }
-
-            return true;
-          },
+          a -> setPNIKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds),
          a -> accounts.changeNumber(a, number, phoneNumberIdentifier),
          () -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
          AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
@@ -329,6 +307,58 @@ public class AccountsManager {
    return updatedAccount.get();
  }

+  public Account updatePNIKeys(final Account account,
+      final String pniIdentityKey,
+      final Map<Long, SignedPreKey> pniSignedPreKeys,
+      final Map<Long, Integer> pniRegistrationIds) throws MismatchedDevicesException {
+    validateDevices(account, pniSignedPreKeys, pniRegistrationIds);
+
+    return update(account, a -> { return setPNIKeys(a, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds); });
+  }
+
+  private boolean setPNIKeys(final Account account,
+      @Nullable final String pniIdentityKey,
+      @Nullable final Map<Long, SignedPreKey> pniSignedPreKeys,
+      @Nullable final Map<Long, Integer> pniRegistrationIds) {
+    if (ObjectUtils.allNull(pniIdentityKey, pniSignedPreKeys, pniRegistrationIds)) {
+      return true;
+    } else if (!ObjectUtils.allNotNull(pniIdentityKey, pniSignedPreKeys, pniRegistrationIds)) {
+      throw new IllegalArgumentException("PNI identity key, signed pre-keys, and registration IDs must be all null or all non-null");
+    }
+
+    pniSignedPreKeys.forEach((deviceId, signedPreKey) ->
+        account.getDevice(deviceId).ifPresent(device -> device.setPhoneNumberIdentitySignedPreKey(signedPreKey)));
+
+    pniRegistrationIds.forEach((deviceId, registrationId) ->
+        account.getDevice(deviceId).ifPresent(device -> device.setPhoneNumberIdentityRegistrationId(registrationId)));
+
+      account.setPhoneNumberIdentityKey(pniIdentityKey);
+
+    return true;
+  }
+
+  private void validateDevices(final Account account,
+      final Map<Long, SignedPreKey> pniSignedPreKeys,
+      final Map<Long, Integer> pniRegistrationIds) throws MismatchedDevicesException {
+    if (pniSignedPreKeys == null && pniRegistrationIds == null) {
+      return;
+    } else if (pniSignedPreKeys == null || pniRegistrationIds == null) {
+      throw new IllegalArgumentException("Signed pre-keys and registration IDs must both be null or both be non-null");
+    }
+
+    // Check that all including master ID are in signed pre-keys
+    DestinationDeviceValidator.validateCompleteDeviceList(
+        account,
+        pniSignedPreKeys.keySet(),
+        Collections.emptySet());
+
+    // Check that all devices are accounted for in the map of new PNI registration IDs
+    DestinationDeviceValidator.validateCompleteDeviceList(
+        account,
+        pniRegistrationIds.keySet(),
+        Collections.emptySet());
+  }
+
  public record UsernameReservation(Account account, byte[] reservedUsernameHash){}

  /**
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManager.java
@@ -5,6 +5,7 @@
 package org.whispersystems.textsecuregcm.storage;

 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
 import com.google.protobuf.ByteString;
 import org.apache.commons.lang3.ObjectUtils;
 import org.slf4j.Logger;
@@ -46,48 +47,70 @@ public class ChangeNumberManager {
      throws InterruptedException, MismatchedDevicesException, StaleDevicesException {

    if (ObjectUtils.allNotNull(pniIdentityKey, deviceSignedPreKeys, deviceMessages, pniRegistrationIds)) {
-      assert pniIdentityKey != null;
-      assert deviceSignedPreKeys != null;
-      assert deviceMessages != null;
-      assert pniRegistrationIds != null;
-
-      // Check that all except master ID are in device messages
-      DestinationDeviceValidator.validateCompleteDeviceList(
-          account,
-          deviceMessages.stream().map(IncomingMessage::destinationDeviceId).collect(Collectors.toSet()),
-          Set.of(Device.MASTER_ID));
-
-      DestinationDeviceValidator.validateRegistrationIds(
-          account,
-          deviceMessages,
-          IncomingMessage::destinationDeviceId,
-          IncomingMessage::destinationRegistrationId,
-          false);
+      // AccountsManager validates the device set on deviceSignedPreKeys and pniRegistrationIds
+      validateDeviceMessages(account, deviceMessages);
    } else if (!ObjectUtils.allNull(pniIdentityKey, deviceSignedPreKeys, deviceMessages, pniRegistrationIds)) {
      throw new IllegalArgumentException("PNI identity key, signed pre-keys, device messages, and registration IDs must be all null or all non-null");
    }

-    final Account updatedAccount;

    if (number.equals(account.getNumber())) {
-      // This may be a request that got repeated due to poor network conditions or other client error; take no action,
-      // but report success since the account is in the desired state
-      updatedAccount = account;
-    } else {
-      updatedAccount = accountsManager.changeNumber(account, number, pniIdentityKey, deviceSignedPreKeys, pniRegistrationIds);
+      // The client has gotten confused/desynchronized with us about their own phone number, most likely due to losing
+      // our OK response to an immediately preceding change-number request, and are sending a change they don't realize
+      // is a no-op change.
+      //
+      // We don't need to actually do a number-change operation in our DB, but we *do* need to accept their new key
+      // material and distribute the sync messages, to be sure all clients agree with us and each other about what their
+      // keys are. Pretend this change-number request was actually a PNI key distribution request.
+      return updatePNIKeys(account, pniIdentityKey, deviceSignedPreKeys, deviceMessages, pniRegistrationIds);
    }

-    // Whether the account already has this number or not, we resend messages. 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 even if the first time
-    // around the server crashed at/above this point.
+    final Account updatedAccount = accountsManager.changeNumber(account, number, pniIdentityKey, deviceSignedPreKeys, pniRegistrationIds);
+
    if (deviceMessages != null) {
-      deviceMessages.forEach(message ->
-          sendMessageToSelf(updatedAccount, updatedAccount.getDevice(message.destinationDeviceId()), message));
+      sendDeviceMessages(updatedAccount, deviceMessages);
    }

    return updatedAccount;
  }

+  public Account updatePNIKeys(final Account account,
+      final String pniIdentityKey,
+      final Map<Long, SignedPreKey> deviceSignedPreKeys,
+      final List<IncomingMessage> deviceMessages,
+      final Map<Long, Integer> pniRegistrationIds) throws MismatchedDevicesException, StaleDevicesException {
+    validateDeviceMessages(account, deviceMessages);
+
+    // Don't try to be smart about ignoring unnecessary retries. If we make literally no change we will skip the ddb
+    // write anyway. Linked devices can handle some wasted extra key rotations.
+    final Account updatedAccount = accountsManager.updatePNIKeys(account, pniIdentityKey, deviceSignedPreKeys, pniRegistrationIds);
+
+    sendDeviceMessages(updatedAccount, deviceMessages);
+    return updatedAccount;
+  }
+
+  private void validateDeviceMessages(final Account account,
+      final List<IncomingMessage> deviceMessages) throws MismatchedDevicesException, StaleDevicesException {
+    // Check that all except master ID are in device messages
+    DestinationDeviceValidator.validateCompleteDeviceList(
+        account,
+        deviceMessages.stream().map(IncomingMessage::destinationDeviceId).collect(Collectors.toSet()),
+        Set.of(Device.MASTER_ID));
+
+    // check that all sync messages are to the current registration ID for the matching device
+    DestinationDeviceValidator.validateRegistrationIds(
+        account,
+        deviceMessages,
+        IncomingMessage::destinationDeviceId,
+        IncomingMessage::destinationRegistrationId,
+        false);
+  }
+
+  private void sendDeviceMessages(final Account account, final List<IncomingMessage> deviceMessages) {
+    deviceMessages.forEach(message ->
+        sendMessageToSelf(account, account.getDevice(message.destinationDeviceId()), message));
+  }
+
  @VisibleForTesting
  void sendMessageToSelf(
      Account sourceAndDestinationAccount, Optional<Device> destinationDevice, IncomingMessage message) {