Update recently-deleted accounts table transactionally as part of account mutations

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

@@ -49,11 +49,9 @@ import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
 import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
 import software.amazon.awssdk.services.dynamodb.model.Delete;
-import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.Put;
-import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
@@ -200,8 +198,12 @@ public class Accounts extends AbstractDynamoDbStore {
 
         final TransactWriteItem accountPut = buildAccountPut(account, uuidAttr, numberAttr, pniUuidAttr);
 
+        // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
+        // the newly-created account.
+        final TransactWriteItem deletedAccountDelete = buildRemoveDeletedAccount(account.getNumber());
+
         final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
-            .transactItems(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut)
+            .transactItems(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut, deletedAccountDelete)
             .build();
 
         try {
@@ -270,7 +272,11 @@ public class Accounts extends AbstractDynamoDbStore {
    * @param account the account for which to change the phone number
    * @param number the new phone number
    */
-  public void changeNumber(final Account account, final String number, final UUID phoneNumberIdentifier) {
+  public void changeNumber(final Account account,
+      final String number,
+      final UUID phoneNumberIdentifier,
+      final Optional<UUID> maybeDisplacedAccountIdentifier) {
+
     CHANGE_NUMBER_TIMER.record(() -> {
       final String originalNumber = account.getNumber();
       final UUID originalPni = account.getPhoneNumberIdentifier();
@@ -289,6 +295,10 @@ public class Accounts extends AbstractDynamoDbStore {
         writeItems.add(buildConstraintTablePut(phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr));
         writeItems.add(buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, originalPni));
         writeItems.add(buildConstraintTablePut(phoneNumberIdentifierConstraintTableName, uuidAttr, ATTR_PNI_UUID, pniAttr));
+        writeItems.add(buildRemoveDeletedAccount(number));
+        maybeDisplacedAccountIdentifier.ifPresent(displacedAccountIdentifier ->
+            writeItems.add(buildPutDeletedAccount(displacedAccountIdentifier, originalNumber)));
+
         writeItems.add(
             TransactWriteItem.builder()
                 .update(Update.builder()
@@ -717,21 +727,25 @@ public class Accounts extends AbstractDynamoDbStore {
             .map(Accounts::fromItem)));
   }
 
-  public void putRecentlyDeletedAccount(final UUID uuid, final String e164) {
+  private TransactWriteItem buildPutDeletedAccount(final UUID uuid, final String e164) {
-    db().putItem(PutItemRequest.builder()
+    return TransactWriteItem.builder()
-        .tableName(deletedAccountsTableName)
+        .put(Put.builder()
-        .item(Map.of(
+            .tableName(deletedAccountsTableName)
-            DELETED_ACCOUNTS_KEY_ACCOUNT_E164, AttributeValues.fromString(e164),
+            .item(Map.of(
-            DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
+                DELETED_ACCOUNTS_KEY_ACCOUNT_E164, AttributeValues.fromString(e164),
-            DELETED_ACCOUNTS_ATTR_EXPIRES, AttributeValues.fromLong(Instant.now().plus(DELETED_ACCOUNTS_TIME_TO_LIVE).getEpochSecond())))
+                DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
-        .build());
+                DELETED_ACCOUNTS_ATTR_EXPIRES, AttributeValues.fromLong(Instant.now().plus(DELETED_ACCOUNTS_TIME_TO_LIVE).getEpochSecond())))
+            .build())
+        .build();
   }
 
-  public void removeRecentlyDeletedAccount(final String e164) {
+  private TransactWriteItem buildRemoveDeletedAccount(final String e164) {
-    db().deleteItem(DeleteItemRequest.builder()
+    return TransactWriteItem.builder()
-        .tableName(deletedAccountsTableName)
+        .delete(Delete.builder()
-        .key(Map.of(DELETED_ACCOUNTS_KEY_ACCOUNT_E164, AttributeValues.fromString(e164)))
+            .tableName(deletedAccountsTableName)
-        .build());
+            .key(Map.of(DELETED_ACCOUNTS_KEY_ACCOUNT_E164, AttributeValues.fromString(e164)))
+            .build())
+        .build();
   }
 
   @Nonnull
@@ -778,7 +792,8 @@ public class Accounts extends AbstractDynamoDbStore {
       final List<TransactWriteItem> transactWriteItems = new ArrayList<>(List.of(
           buildDelete(phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, account.getNumber()),
           buildDelete(accountsTableName, KEY_ACCOUNT_UUID, uuid),
-          buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, account.getPhoneNumberIdentifier())
+          buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, account.getPhoneNumberIdentifier()),
+          buildPutDeletedAccount(uuid, account.getNumber())
       ));
 
       account.getUsernameHash().ifPresent(usernameHash -> transactWriteItems.add(

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

@@ -255,10 +255,6 @@ public class AccountsManager {
 
         accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
             registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword));
-
-        // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
-        // the newly-created account.
-        accounts.removeRecentlyDeletedAccount(number);
       });
 
       return account;
@@ -311,8 +307,6 @@ public class AccountsManager {
         maybeDisplacedUuid = recentlyDeletedAci;
       }
 
-      maybeDisplacedUuid.ifPresent(displacedUuid -> accounts.putRecentlyDeletedAccount(displacedUuid, originalNumber));
-
       final UUID uuid = account.getUuid();
       final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber);
 
@@ -324,7 +318,7 @@ public class AccountsManager {
             setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds);
             return true;
           },
-          a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier),
+          a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier, maybeDisplacedUuid),
           () -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
           AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
 
@@ -855,14 +849,7 @@ public class AccountsManager {
 
   public void delete(final Account account, final DeletionReason deletionReason) throws InterruptedException {
     try (final Timer.Context ignored = deleteTimer.time()) {
-      accountLockManager.withLock(List.of(account.getNumber()), () -> {
+      accountLockManager.withLock(List.of(account.getNumber()), () -> delete(account));
-        final UUID accountIdentifier = account.getUuid();
-        final String e164 = account.getNumber();
-
-        delete(account);
-
-        accounts.putRecentlyDeletedAccount(accountIdentifier, e164);
-      });
     } catch (final RuntimeException | InterruptedException e) {
       logger.warn("Failed to delete account", e);
       throw e;

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

@@ -148,7 +148,7 @@ class AccountsManagerTest {
       account.setNumber(number, phoneNumberIdentifier);
 
       return null;
-    }).when(accounts).changeNumber(any(), anyString(), any());
+    }).when(accounts).changeNumber(any(), anyString(), any(), any());
 
     final SecureStorageClient storageClient = mock(SecureStorageClient.class);
     when(storageClient.deleteStoredData(any())).thenReturn(CompletableFuture.completedFuture(null));
@@ -1031,7 +1031,6 @@ class AccountsManagerTest {
     account = accountsManager.changeNumber(account, number, null, null, null, null);
 
     assertEquals(number, account.getNumber());
-    verify(accounts, never()).putRecentlyDeletedAccount(any(), any());
     verify(keysManager, never()).delete(any());
   }
 

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

@@ -39,12 +39,15 @@ import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.apache.commons.lang3.RandomUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
@@ -228,6 +231,34 @@ class AccountsTest {
     assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier(), account.getUuid());
   }
 
+  @Test
+  void testStoreRecentlyDeleted() {
+    final UUID originalUuid = UUID.randomUUID();
+
+    Device device = generateDevice(1);
+    Account account = generateAccount("+14151112222", originalUuid, UUID.randomUUID(), List.of(device));
+
+    boolean freshUser = accounts.create(account);
+
+    assertThat(freshUser).isTrue();
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier(), null, account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier(), account.getUuid());
+
+    accounts.delete(originalUuid);
+    assertThat(accounts.findRecentlyDeletedAccountIdentifier(account.getNumber())).hasValue(originalUuid);
+
+    freshUser = accounts.create(account);
+    assertThat(freshUser).isTrue();
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier(), null, account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier(), account.getUuid());
+
+    assertThat(accounts.findRecentlyDeletedAccountIdentifier(account.getNumber())).isEmpty();
+  }
+
   @Test
   void testStoreMulti() {
     final List<Device> devices = List.of(generateDevice(1), generateDevice(2));
@@ -536,6 +567,8 @@ class AccountsTest {
     accounts.create(deletedAccount);
     accounts.create(retainedAccount);
 
+    assertThat(accounts.findRecentlyDeletedAccountIdentifier(deletedAccount.getNumber())).isEmpty();
+
     assertPhoneNumberConstraintExists("+14151112222", deletedAccount.getUuid());
     assertPhoneNumberIdentifierConstraintExists(deletedAccount.getPhoneNumberIdentifier(), deletedAccount.getUuid());
     assertPhoneNumberConstraintExists("+14151112345", retainedAccount.getUuid());
@@ -547,6 +580,7 @@ class AccountsTest {
     accounts.delete(deletedAccount.getUuid());
 
     assertThat(accounts.getByAccountIdentifier(deletedAccount.getUuid())).isNotPresent();
+    assertThat(accounts.findRecentlyDeletedAccountIdentifier(deletedAccount.getNumber())).hasValue(deletedAccount.getUuid());
 
     assertPhoneNumberConstraintDoesNotExist(deletedAccount.getNumber());
     assertPhoneNumberIdentifierConstraintDoesNotExist(deletedAccount.getPhoneNumberIdentifier());
@@ -637,8 +671,10 @@ class AccountsTest {
     verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier(), null, account, false);
   }
 
-  @Test
+  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
-  public void testChangeNumber() {
+  @ParameterizedTest
+  @MethodSource
+  public void testChangeNumber(final Optional<UUID> maybeDisplacedAccountIdentifier) {
     final String originalNumber = "+14151112222";
     final String targetNumber = "+14151113333";
 
@@ -662,7 +698,7 @@ class AccountsTest {
       verifyStoredState(originalNumber, account.getUuid(), account.getPhoneNumberIdentifier(), null, retrieved.get(), account);
     }
 
-    accounts.changeNumber(account, targetNumber, targetPni);
+    accounts.changeNumber(account, targetNumber, targetPni, maybeDisplacedAccountIdentifier);
 
     assertThat(accounts.getByE164(originalNumber)).isEmpty();
     assertThat(accounts.getByAccountIdentifier(originalPni)).isEmpty();
@@ -681,6 +717,15 @@ class AccountsTest {
       assertThat(retrieved.get().getPhoneNumberIdentifier()).isEqualTo(targetPni);
       assertThat(accounts.getByPhoneNumberIdentifier(targetPni)).isPresent();
     }
+
+    assertThat(accounts.findRecentlyDeletedAccountIdentifier(originalNumber)).isEqualTo(maybeDisplacedAccountIdentifier);
+  }
+
+  private static Stream<Arguments> testChangeNumber() {
+    return Stream.of(
+        Arguments.of(Optional.empty()),
+        Arguments.of(Optional.of(UUID.randomUUID()))
+    );
   }
 
   @Test
@@ -700,7 +745,7 @@ class AccountsTest {
     accounts.create(account);
     accounts.create(existingAccount);
 
-    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, targetPni));
+    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, targetPni, Optional.of(existingAccount.getUuid())));
 
     assertPhoneNumberConstraintExists(originalNumber, account.getUuid());
     assertPhoneNumberIdentifierConstraintExists(originalPni, account.getUuid());
@@ -736,7 +781,7 @@ class AccountsTest {
             Map.of(":uuid", AttributeValues.fromUUID(existingAccountIdentifier)))
         .build());
 
-    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, existingPhoneNumberIdentifier));
+    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, existingPhoneNumberIdentifier, Optional.empty()));
   }
 
   @Test
@@ -992,59 +1037,6 @@ class AccountsTest {
     assertThat(account.getUsernameHash()).isEmpty();
   }
 
-  @Test
-  void testPutFindRecentlyDeletedAccount() {
-    final UUID uuid = UUID.randomUUID();
-    final String e164 = "+18005551234";
-
-    assertEquals(Optional.empty(), accounts.findRecentlyDeletedAccountIdentifier(e164));
-
-    accounts.putRecentlyDeletedAccount(uuid, e164);
-
-    assertEquals(Optional.of(uuid), accounts.findRecentlyDeletedAccountIdentifier(e164));
-  }
-
-  @Test
-  void testRemoveRecentlyDeletedAccount() {
-    final UUID uuid = UUID.randomUUID();
-    final String e164 = "+18005551234";
-
-    assertEquals(Optional.empty(), accounts.findRecentlyDeletedAccountIdentifier(e164));
-
-    accounts.putRecentlyDeletedAccount(uuid, e164);
-
-    assertEquals(Optional.of(uuid), accounts.findRecentlyDeletedAccountIdentifier(e164));
-
-    accounts.removeRecentlyDeletedAccount(e164);
-
-    assertEquals(Optional.empty(), accounts.findRecentlyDeletedAccountIdentifier(e164));
-  }
-
-  @Test
-  void testFindRecentlyDeletedE164() {
-    assertEquals(Optional.empty(), accounts.findRecentlyDeletedE164(UUID.randomUUID()));
-
-    final UUID uuid = UUID.randomUUID();
-    final String e164 = "+18005551234";
-
-    accounts.putRecentlyDeletedAccount(uuid, e164);
-
-    assertEquals(Optional.of(e164), accounts.findRecentlyDeletedE164(uuid));
-  }
-
-  @Test
-  void testFindRecentlyDeletedUUID() {
-    final String e164 = "+18005551234";
-
-    assertEquals(Optional.empty(), accounts.findRecentlyDeletedAccountIdentifier(e164));
-
-    final UUID uuid = UUID.randomUUID();
-
-    accounts.putRecentlyDeletedAccount(uuid, e164);
-
-    assertEquals(Optional.of(uuid), accounts.findRecentlyDeletedAccountIdentifier(e164));
-  }
-
   @Test
   public void testIgnoredFieldsNotAddedToDataAttribute() throws Exception {
     final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());