From cfac798f9eb4d2c172a5fb65cdf9c33fbfade466 Mon Sep 17 00:00:00 2001 From: Katherine Date: Wed, 1 Apr 2026 11:29:06 -0400 Subject: [PATCH] Handle account reclamation with equivalent phone numbers --- .../textsecuregcm/storage/Accounts.java | 142 +++++++++++++++-- ...nexpectedExistingPhoneNumberException.java | 9 ++ .../textsecuregcm/storage/AccountsTest.java | 147 +++++++++++++++++- 3 files changed, 285 insertions(+), 13 deletions(-) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/UnexpectedExistingPhoneNumberException.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java index 70906cb76..a7b28631b 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java +++ b/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 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 setClauses, List addClauses, List 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 key, Map attrNames, Map 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 attrValues = new HashMap<>(base.attrValues()); + attrValues.put(":number", AttributeValues.fromString(account.getNumber())); + + final UpdateExpression updateExpression = base.updateExpression(); + final List 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 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 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 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 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 values) { + public static MembershipExpression build(final Collection values) { + if (values.isEmpty()) { + throw new IllegalArgumentException("must have at least one value"); + } + + final Map expressionValues = new HashMap<>(); + final List 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); + } + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/UnexpectedExistingPhoneNumberException.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/UnexpectedExistingPhoneNumberException.java new file mode 100644 index 000000000..de704fc8c --- /dev/null +++ b/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 { +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java index 132425395..07321f471 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/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 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 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);