Handle account reclamation with equivalent phone numbers

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

@@ -6,6 +6,7 @@ package org.whispersystems.textsecuregcm.storage;
 
 import static java.util.Objects.requireNonNull;
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+import static org.whispersystems.textsecuregcm.util.Util.getAlternateForms;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectWriter;
@@ -371,7 +372,28 @@ public class Accounts {
                 .build())
             .build());
       }
-      writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, accountToCreate).transactItem());
+
+      // Phone number canonicalization means that a user can use a different phone number in the same equivalence class
+      // to reclaim the account.
+      if (!existingAccount.getNumber().equals(accountToCreate.getNumber())) {
+        if (getAlternateForms(existingAccount.getNumber()).contains(accountToCreate.getNumber())) {
+          final AttributeValue uuidAttr = AttributeValues.fromUUID(existingAccount.getUuid());
+          final AttributeValue numberAttr = AttributeValues.fromString(accountToCreate.getNumber());
+          final TransactWriteItem phoneNumberConstraintPut = buildConstraintTablePutIfAbsent(
+              phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr);
+
+          writeItems.add(buildDelete(phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, existingAccount.getNumber()));
+          writeItems.add(phoneNumberConstraintPut);
+        } else {
+          log.error("Reclaiming account with a non-equivalent phone number. Old account {}:{}:{}, new account {}:{}:{}",
+              existingAccount.getUuid(), existingAccount.getNumber(), existingAccount.getPhoneNumberIdentifier(),
+              accountToCreate.getUuid(), accountToCreate.getNumber(), accountToCreate.getPhoneNumberIdentifier());
+          throw new IllegalArgumentException("reclaimed accounts must have equivalent phone numbers");
+        }
+      }
+
+      final int updateAccountItemIndex = writeItems.size();
+      writeItems.add(UpdateAccountSpec.forReclaimedAccount(accountsTableName, accountToCreate, existingAccount.getNumber()).transactItem());
       writeItems.addAll(additionalWriteItems);
 
       return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
@@ -382,6 +404,16 @@ public class Accounts {
           .exceptionally(throwable -> {
             final Throwable unwrapped = ExceptionUtils.unwrap(throwable);
             if (unwrapped instanceof TransactionCanceledException te) {
+              if (Accounts.conditionalCheckFailed(te.cancellationReasons().get(updateAccountItemIndex))) {
+                final Map<String, AttributeValue> item = te.cancellationReasons().get(updateAccountItemIndex).item();
+                final String existingNumber = AttributeValues.getString(item, Accounts.ATTR_ACCOUNT_E164, null);
+                if (!existingAccount.getNumber().equals(existingNumber)) {
+                  log.error("Failed to update account due to unexpected existing phone number. Account {}. Expected {}, got {}",
+                      existingAccount.getUuid(), existingAccount.getNumber(), existingNumber);
+                  throw new UnexpectedExistingPhoneNumberException();
+                }
+              }
+
               if (te.cancellationReasons().stream().anyMatch(Accounts::conditionalCheckFailed)) {
                 throw new ContestedOptimisticLockException();
               }
@@ -889,6 +921,37 @@ public class Accounts {
     }
   }
 
+  record UpdateExpression (List<String> setClauses, List<String> addClauses, List<String> removeClauses) {
+    public String toExpressionString() {
+      final StringBuilder updateExpressionBuilder = new StringBuilder();
+
+      if (!setClauses.isEmpty()) {
+        updateExpressionBuilder.append("SET ");
+        updateExpressionBuilder.append(String.join(", ", setClauses));
+      }
+
+      if (!removeClauses.isEmpty()) {
+        if (!updateExpressionBuilder.isEmpty()) {
+          updateExpressionBuilder.append(" ");
+        }
+
+        updateExpressionBuilder.append("REMOVE ");
+        updateExpressionBuilder.append(String.join(", ", removeClauses));
+      }
+
+      if (!addClauses.isEmpty()) {
+        if (!updateExpressionBuilder.isEmpty()) {
+          updateExpressionBuilder.append(" ");
+        }
+
+        updateExpressionBuilder.append("ADD ");
+        updateExpressionBuilder.append(String.join(", ", addClauses));
+      }
+
+      return updateExpressionBuilder.toString();
+    }
+  }
+
   /**
    * A ddb update that can be used as part of a transaction or single-item update statement.
    */
@@ -897,13 +960,13 @@ public class Accounts {
       Map<String, AttributeValue> key,
       Map<String, String> attrNames,
       Map<String, AttributeValue> attrValues,
-      String updateExpression,
+      UpdateExpression updateExpression,
       String conditionExpression) {
     UpdateItemRequest updateItemRequest() {
       return UpdateItemRequest.builder()
           .tableName(tableName)
           .key(key)
-          .updateExpression(updateExpression)
+          .updateExpression(updateExpression.toExpressionString())
           .conditionExpression(conditionExpression)
           .expressionAttributeNames(attrNames)
           .expressionAttributeValues(attrValues)
@@ -914,17 +977,46 @@ public class Accounts {
       return TransactWriteItem.builder().update(Update.builder()
           .tableName(tableName)
           .key(key)
-          .updateExpression(updateExpression)
+          .updateExpression(updateExpression.toExpressionString())
           .conditionExpression(conditionExpression)
+          .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
           .expressionAttributeNames(attrNames)
           .expressionAttributeValues(attrValues)
           .build()).build();
     }
 
+    static UpdateAccountSpec forReclaimedAccount(
+        final String accountTableName,
+        final Account account,
+        final String expectedExistingE164) {
+      final UpdateAccountSpec base = forAccount(accountTableName, account);
+
+      final Map<String, AttributeValue> attrValues = new HashMap<>(base.attrValues());
+      attrValues.put(":number", AttributeValues.fromString(account.getNumber()));
+
+      final UpdateExpression updateExpression = base.updateExpression();
+      final List<String> setClauses = new ArrayList<>(updateExpression.setClauses());
+      setClauses.add("#number = :number");
+
+      final MembershipExpression membershipExpression = MembershipExpression.build(getAlternateForms(expectedExistingE164));
+      attrValues.putAll(membershipExpression.values());
+
+      // Defensive check: we should only update the e164 to another e164 in the same equivalence class
+      final String conditionExpression = base.conditionExpression() + " AND #number IN %s".formatted(membershipExpression.expression());
+
+      return new UpdateAccountSpec(
+          base.tableName(),
+          base.key(),
+          base.attrNames(),
+          attrValues,
+          new UpdateExpression(setClauses, updateExpression.addClauses(), updateExpression.removeClauses()),
+          conditionExpression
+      );
+    }
+
     static UpdateAccountSpec forAccount(
         final String accountTableName,
         final Account account) {
-      // username, e164, and pni cannot be modified through this method
       final Map<String, String> attrNames = new HashMap<>(Map.of(
           "#number", ATTR_ACCOUNT_E164,
           "#data", ATTR_ACCOUNT_DATA,
@@ -937,19 +1029,20 @@ public class Accounts {
           ":version", AttributeValues.fromInt(account.getVersion()),
           ":version_increment", AttributeValues.fromInt(1)));
 
-      final StringBuilder updateExpressionBuilder = new StringBuilder("SET #data = :data, #cds = :cds");
+      final List<String> setClauses = new ArrayList<>(List.of("#data = :data", "#cds = :cds"));
+
       if (account.getUnidentifiedAccessKey().isPresent()) {
         // if it's present in the account, also set the uak
         attrNames.put("#uak", ATTR_UAK);
         attrValues.put(":uak", AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get()));
-        updateExpressionBuilder.append(", #uak = :uak");
+        setClauses.add("#uak = :uak");
       }
 
       if (account.getUsernameHash().isPresent()) {
         // if it's present in the account, also set the username hash
         attrNames.put("#usernameHash", ATTR_USERNAME_HASH);
         attrValues.put(":usernameHash", AttributeValues.fromByteArray(account.getUsernameHash().get()));
-        updateExpressionBuilder.append(", #usernameHash = :usernameHash");
+        setClauses.add("#usernameHash = :usernameHash");
       }
 
       // If the account has a username/handle pair, we should add it to the top level attributes.
@@ -959,7 +1052,7 @@ public class Accounts {
       if (account.getEncryptedUsername().isPresent() && account.getUsernameLinkHandle() != null) {
         attrNames.put("#ul", ATTR_USERNAME_LINK_UUID);
         attrValues.put(":ul", AttributeValues.fromUUID(account.getUsernameLinkHandle()));
-        updateExpressionBuilder.append(", #ul = :ul");
+        setClauses.add("#ul = :ul");
       }
 
       // Some operations may remove the usernameLink or the usernameHash (re-registration, clear username link, and
@@ -974,17 +1067,20 @@ public class Accounts {
         attrNames.put("#username_hash", ATTR_USERNAME_HASH);
         removes.add("#username_hash");
       }
+
+      final List<String> removeClauses = new ArrayList<>();
       if (!removes.isEmpty()) {
-        updateExpressionBuilder.append(" REMOVE %s".formatted(String.join(",", removes)));
+        removeClauses.add(String.join(",", removes));
       }
-      updateExpressionBuilder.append(" ADD #version :version_increment");
+
+      final List<String> addClauses = List.of("#version :version_increment");
 
       return new UpdateAccountSpec(
           accountTableName,
           Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())),
           attrNames,
           attrValues,
-          updateExpressionBuilder.toString(),
+          new UpdateExpression(setClauses, addClauses, removeClauses),
           "attribute_exists(#number) AND #version = :version");
     }
   }
@@ -1584,4 +1680,26 @@ public class Accounts {
         : phoneNumber.substring(phoneNumber.length() - 2));
     return sb.toString();
   }
+
+  record MembershipExpression(String expression, Map<String, AttributeValue> values) {
+    public static MembershipExpression build(final Collection<String> values) {
+      if (values.isEmpty()) {
+        throw new IllegalArgumentException("must have at least one value");
+      }
+
+      final Map<String, AttributeValue> expressionValues = new HashMap<>();
+      final List<String> placeholders = new ArrayList<>();
+      final String prefix = "val";
+      int i = 0;
+      for (final String value : values) {
+        String key = ":" + prefix + i;
+        placeholders.add(key);
+        expressionValues.put(key, AttributeValues.fromString(value));
+        i++;
+      }
+
+      final String expression = "(" + String.join(", ", placeholders) + ")";
+      return new MembershipExpression(expression, expressionValues);
+    }
+  }
 }

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

@@ -0,0 +1,9 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+public class UnexpectedExistingPhoneNumberException extends RuntimeException {
+}

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java

@@ -20,8 +20,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil.assertFailsWithCause;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
+import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.time.Instant;
@@ -47,6 +49,7 @@ import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
+import com.google.i18n.phonenumbers.PhoneNumberUtil;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
@@ -455,7 +458,11 @@ class AccountsTest {
             () -> accounts.create(reregisteredAccount, Collections.emptyList()));
 
     reregisteredAccount.setUuid(accountAlreadyExistsException.getExistingAccount().getUuid());
-    reregisteredAccount.setNumber(accountAlreadyExistsException.getExistingAccount().getNumber(),
+
+    // Phone number canonicalization means that a user can re-register with a different phone number
+    // in the same equivalence class and get back the same phone number identifier.
+    // In that case, we favor the re-registering account's phone number.
+    reregisteredAccount.setNumber(reregisteredAccount.getNumber(),
         accountAlreadyExistsException.getExistingAccount().getPhoneNumberIdentifier());
 
     assertDoesNotThrow(() -> accounts.reclaimAccount(accountAlreadyExistsException.getExistingAccount(),
@@ -554,6 +561,144 @@ class AccountsTest {
     assertThatThrownBy(() -> createAccount(invalidAccount));
   }
 
+  @ParameterizedTest
+  @MethodSource
+  void testReclaimAccountEquivalentPhoneNumbers(final String firstNumber, final String secondNumber) throws IOException {
+    final UUID existingUuid = UUID.randomUUID();
+    final UUID pni = UUID.randomUUID();
+    final Account existingAccount = generateAccount(firstNumber, existingUuid, pni, List.of(generateDevice(DEVICE_ID_1)));
+
+    createAccount(existingAccount);
+
+    verifyStoredState(firstNumber, existingAccount.getUuid(), existingAccount.getPhoneNumberIdentifier(), null, existingAccount, true);
+
+    assertPhoneNumberConstraintExists(firstNumber, existingUuid);
+    assertPhoneNumberIdentifierConstraintExists(pni, existingUuid);
+
+    assertDoesNotThrow(() -> accounts.update(existingAccount));
+
+    final UUID secondUuid = UUID.randomUUID();
+
+    final Account secondAccount = generateAccount(secondNumber, secondUuid, pni, List.of(generateDevice(DEVICE_ID_1)));
+
+    reclaimAccount(secondAccount);
+
+    Map<String, AttributeValue> item = readAccount(existingUuid);
+    final Account account = SystemMapper.jsonMapper().readValue(item.get(Accounts.ATTR_ACCOUNT_DATA).b().asByteArray(), Account.class);
+
+    assertThat(AttributeValues.getString(item, Accounts.ATTR_ACCOUNT_E164, null))
+        .isEqualTo(secondNumber)
+        .isEqualTo(account.getNumber());
+    assertPhoneNumberConstraintDoesNotExist(firstNumber);
+    assertPhoneNumberConstraintExists(secondNumber, existingUuid);
+    assertPhoneNumberIdentifierConstraintExists(pni, existingUuid);
+  }
+
+  private static Stream<Arguments> testReclaimAccountEquivalentPhoneNumbers() {
+    final String newFormatBeninE164 = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("BJ"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final String oldFormatBeninE164 = newFormatBeninE164.replaceFirst("01", "");
+    return Stream.of(
+        Arguments.of(newFormatBeninE164, oldFormatBeninE164),
+        Arguments.of(oldFormatBeninE164, newFormatBeninE164)
+    );
+  }
+
+  @Test
+  void testReclaimAccountNonEquivalentPhoneNumbers() {
+    final String beninPhoneNumber = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("BJ"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final UUID existingUuid = UUID.randomUUID();
+    final UUID pni = UUID.randomUUID();
+    final Account existingAccount = generateAccount(beninPhoneNumber, existingUuid, pni, List.of(generateDevice(DEVICE_ID_1)));
+
+    createAccount(existingAccount);
+
+    verifyStoredState(beninPhoneNumber, existingAccount.getUuid(), existingAccount.getPhoneNumberIdentifier(), null, existingAccount, true);
+
+    assertPhoneNumberConstraintExists(beninPhoneNumber, existingUuid);
+    assertPhoneNumberIdentifierConstraintExists(pni, existingUuid);
+
+    assertDoesNotThrow(() -> accounts.update(existingAccount));
+
+    final String usPhoneNumber = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("US"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final UUID secondUuid = UUID.randomUUID();
+
+    // A non-equivalent phone number with the same PNI should fail reclamation
+    final Account secondAccount = generateAccount(usPhoneNumber, secondUuid, pni, List.of(generateDevice(DEVICE_ID_1)));
+
+    final AccountAlreadyExistsException accountAlreadyExistsException =
+        assertThrows(AccountAlreadyExistsException.class,
+            () -> accounts.create(secondAccount, Collections.emptyList()));
+
+    secondAccount.setUuid(accountAlreadyExistsException.getExistingAccount().getUuid());
+
+    assertThrows(IllegalArgumentException.class, () -> accounts.reclaimAccount(existingAccount,
+        secondAccount,
+        Collections.emptyList()).toCompletableFuture().join());
+  }
+
+  @Test
+  void testReclaimAccountUnexpectedDatabasePhoneNumber() {
+    final String beninPhoneNumber = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("BJ"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final UUID existingUuid = UUID.randomUUID();
+    final UUID existingPni = UUID.randomUUID();
+    final Account existingAccount = generateAccount(beninPhoneNumber, existingUuid, existingPni, List.of(generateDevice(DEVICE_ID_1)));
+
+    createAccount(existingAccount);
+
+    verifyStoredState(beninPhoneNumber, existingAccount.getUuid(), existingAccount.getPhoneNumberIdentifier(), null, existingAccount, true);
+
+    assertPhoneNumberConstraintExists(beninPhoneNumber, existingUuid);
+    assertPhoneNumberIdentifierConstraintExists(existingPni, existingUuid);
+
+    assertDoesNotThrow(() -> accounts.update(existingAccount));
+
+    final String usPhoneNumber = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("US"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final Account secondAccount = generateAccount(usPhoneNumber, existingUuid, existingPni, List.of(generateDevice(DEVICE_ID_1)));
+
+    // This scenario is very contrived but tests our error handling if we somehow use an existing account with a different
+    // phone number than what actually exists in the database.
+    assertFailsWithCause(UnexpectedExistingPhoneNumberException.class, accounts.reclaimAccount(secondAccount, secondAccount,
+        Collections.emptyList()).toCompletableFuture());
+
+  }
+
+  @Test
+  void testUpdateAccountWithMismatchedJsonDdbPhoneNumbers() {
+    // Test that fixing the DynamoDB/JSON phone number mismatch does not break account updates for existing accounts
+    // with bad data in the time after we ship this change and before we run the crawler to fix the mismatch.
+    final String newFormatBeninE164 = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("BJ"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final String oldFormatBeninE164 = newFormatBeninE164.replaceFirst("01", "");
+    final UUID existingUuid = UUID.randomUUID();
+    final UUID existingPni = UUID.randomUUID();
+    final Account existingAccount = generateAccount(newFormatBeninE164, existingUuid, existingPni, List.of(generateDevice(DEVICE_ID_1)));
+
+    createAccount(existingAccount);
+
+    verifyStoredState(newFormatBeninE164, existingAccount.getUuid(), existingAccount.getPhoneNumberIdentifier(), null, existingAccount, true);
+
+    assertPhoneNumberConstraintExists(newFormatBeninE164, existingUuid);
+    assertPhoneNumberIdentifierConstraintExists(existingPni, existingUuid);
+
+    // Mimic the current bad state
+    DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient().updateItem(UpdateItemRequest.builder()
+        .tableName(Tables.ACCOUNTS.tableName())
+        .key(Map.of(Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(existingUuid)))
+        .updateExpression("SET #number = :old_number")
+        .expressionAttributeNames(Map.of("#number", Accounts.ATTR_ACCOUNT_E164))
+        .expressionAttributeValues(
+            Map.of(":old_number", AttributeValues.fromString(oldFormatBeninE164)))
+        .build())
+        .join();
+
+    assertDoesNotThrow(() -> accounts.update(existingAccount));
+  }
+
   @Test
   void testUpdate() {
     Device device = generateDevice(DEVICE_ID_1);