From 5f0726af8a825905df31459ff8804d4f582387b6 Mon Sep 17 00:00:00 2001 From: Jon Chambers <63609320+jon-signal@users.noreply.github.com> Date: Tue, 5 Dec 2023 12:18:09 -0500 Subject: [PATCH] Perform cleanup operations before overwriting an existing account record --- .../textsecuregcm/WhisperServerService.java | 8 ++- .../textsecuregcm/storage/Accounts.java | 11 +++- .../storage/AccountsManager.java | 51 ++++++++++--------- .../workers/AssignUsernameCommand.java | 5 +- .../workers/CommandDependencies.java | 5 +- .../AccountCreationIntegrationTest.java | 7 +++ ...ntsManagerChangeNumberIntegrationTest.java | 7 +++ ...ConcurrentModificationIntegrationTest.java | 1 + .../storage/AccountsManagerTest.java | 42 ++++++++++----- ...ccountsManagerUsernameIntegrationTest.java | 4 +- .../textsecuregcm/storage/AccountsTest.java | 31 ++++++++++- 11 files changed, 128 insertions(+), 44 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index e2d718000..5cfc83c39 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -468,6 +468,11 @@ public class WhisperServerService extends Application> additionalWriteItemsFunction) { + public boolean create(final Account account, + final Function> additionalWriteItemsFunction, + final BiFunction> existingAccountCleanupOperation) { return CREATE_TIMER.record(() -> { try { @@ -239,7 +243,10 @@ public class Accounts extends AbstractDynamoDbStore { account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid)); final Account existingAccount = getByAccountIdentifier(account.getUuid()).orElseThrow(); account.setNumber(existingAccount.getNumber(), existingAccount.getPhoneNumberIdentifier()); - joinAndUnwrapUpdateFuture(reclaimAccount(existingAccount, account, additionalWriteItemsFunction.apply(account))); + + existingAccountCleanupOperation.apply(existingAccount.getIdentifier(IdentityType.ACI), existingAccount.getIdentifier(IdentityType.PNI)) + .thenCompose(ignored -> reclaimAccount(existingAccount, account, additionalWriteItemsFunction.apply(account))) + .join(); return false; } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index df953afab..9d73cb056 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -114,6 +114,7 @@ public class AccountsManager { private final ExperimentEnrollmentManager experimentEnrollmentManager; private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager; private final Executor accountLockExecutor; + private final Executor clientPresenceExecutor; private final Clock clock; private static final ObjectWriter ACCOUNT_REDIS_JSON_WRITER = SystemMapper.jsonMapper() @@ -159,6 +160,7 @@ public class AccountsManager { final ExperimentEnrollmentManager experimentEnrollmentManager, final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager, final Executor accountLockExecutor, + final Executor clientPresenceExecutor, final Clock clock) { this.accounts = accounts; this.phoneNumberIdentifiers = phoneNumberIdentifiers; @@ -173,6 +175,7 @@ public class AccountsManager { this.experimentEnrollmentManager = experimentEnrollmentManager; this.registrationRecoveryPasswordsManager = requireNonNull(registrationRecoveryPasswordsManager); this.accountLockExecutor = accountLockExecutor; + this.clientPresenceExecutor = clientPresenceExecutor; this.clock = requireNonNull(clock); } @@ -243,14 +246,33 @@ public class AccountsManager { aciSignedPreKey, pniSignedPreKey, aciPqLastResortPreKey, - pniPqLastResortPreKey)); + pniPqLastResortPreKey), + (aci, pni) -> CompletableFuture.allOf( + keysManager.delete(aci), + keysManager.delete(pni), + messagesManager.clear(aci), + profilesManager.deleteAll(aci) + ).thenRunAsync(() -> clientPresenceManager.disconnectAllPresencesForUuid(aci), clientPresenceExecutor)); - // create() sometimes updates the UUID, if there was a number conflict. - // for metrics, we want secondary to run with the same original UUID - final UUID actualUuid = account.getUuid(); + if (!account.getUuid().equals(originalUuid)) { + // If the UUID changed, then we overwrote an existing account. We should have cleared all messages before + // overwriting the old account, but more may have arrived while we were working. Similarly, the old account + // holder could have added keys or profiles. We'll largely repeat the cleanup process after creating the + // account to make sure we really REALLY got everything. + // + // We exclude the primary device's repeated-use keys from deletion because new keys were provided as + // part of the account creation process, and we don't want to delete the keys that just got added. + CompletableFuture.allOf(keysManager.delete(account.getIdentifier(IdentityType.ACI), true), + keysManager.delete(account.getIdentifier(IdentityType.PNI), true), + messagesManager.clear(account.getIdentifier(IdentityType.ACI)), + profilesManager.deleteAll(account.getIdentifier(IdentityType.ACI))) + .join(); + } redisSet(account); + final Tags tags; + // In terms of previously-existing accounts, there are three possible cases: // // 1. This is a completely new account; there was no pre-existing account and no recently-deleted account @@ -259,27 +281,6 @@ public class AccountsManager { // instance to match the stored account record (i.e. originalUuid != actualUuid). // 3. This is a re-registration of a recently-deleted account, in which case maybeRecentlyDeletedUuid is // present. - // - // All cases are mutually-exclusive. In the first case, we don't need to do anything. In the third, we can be - // confident that everything has already been deleted. In the second case, though, we're taking over an existing - // account and need to clear out messages and keys that may have been stored for the old account. - if (!originalUuid.equals(actualUuid)) { - // We exclude the primary device's repeated-use keys from deletion because new keys were provided as part of - // the account creation process, and we don't want to delete the keys that just got added. - final CompletableFuture deleteKeysFuture = CompletableFuture.allOf( - keysManager.delete(actualUuid, true), - keysManager.delete(account.getPhoneNumberIdentifier(), true)); - - messagesManager.clear(actualUuid).join(); - profilesManager.deleteAll(actualUuid).join(); - - deleteKeysFuture.join(); - - clientPresenceManager.disconnectAllPresencesForUuid(actualUuid); - } - - final Tags tags; - if (freshUser) { tags = Tags.of("type", maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new"); } else { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java index 88c79c61d..a76bdce91 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java @@ -116,6 +116,8 @@ public class AssignUsernameCommand extends EnvironmentCommand { + final Runnable runnable = invocation.getArgument(0); + runnable.run(); + + return null; + }).when(clientPresenceExecutor).execute(any()); + //noinspection unchecked commands = mock(RedisAdvancedClusterCommands.class); @@ -224,6 +234,7 @@ class AccountsManagerTest { enrollmentManager, registrationRecoveryPasswordsManager, mock(Executor.class), + clientPresenceExecutor, mock(Clock.class)); } @@ -856,7 +867,7 @@ class AccountsManagerTest { when(commands.get(eq("Account3::" + uuid))).thenReturn(null); when(accounts.getByAccountIdentifier(uuid)).thenReturn(Optional.empty()) .thenReturn(Optional.of(account)); - when(accounts.create(any(), any())).thenThrow(ContestedOptimisticLockException.class); + when(accounts.create(any(), any(), any())).thenThrow(ContestedOptimisticLockException.class); accountsManager.update(account, a -> { }); @@ -974,14 +985,14 @@ class AccountsManagerTest { @Test void testCreateFreshAccount() throws InterruptedException { - when(accounts.create(any(), any())).thenReturn(true); + when(accounts.create(any(), any(), any())).thenReturn(true); final String e164 = "+18005550123"; final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null); createAccount(e164, attributes); - verify(accounts).create(argThat(account -> e164.equals(account.getNumber())), any()); + verify(accounts).create(argThat(account -> e164.equals(account.getNumber())), any(), any()); verifyNoInteractions(messagesManager); verifyNoInteractions(profilesManager); @@ -991,25 +1002,31 @@ class AccountsManagerTest { void testReregisterAccount() throws InterruptedException { final UUID existingUuid = UUID.randomUUID(); - when(accounts.create(any(), any())).thenAnswer(invocation -> { - invocation.getArgument(0, Account.class).setUuid(existingUuid); - return false; - }); - final String e164 = "+18005550123"; final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null); + when(accounts.create(any(), any(), any())).thenAnswer(invocation -> { + invocation.getArgument(0, Account.class).setUuid(existingUuid); + + final BiFunction> cleanupOperation = invocation.getArgument(2); + cleanupOperation.apply(existingUuid, phoneNumberIdentifiersByE164.get(e164)); + + return false; + }); + createAccount(e164, attributes); assertTrue(phoneNumberIdentifiersByE164.containsKey(e164)); verify(accounts) - .create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())), any()); + .create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())), any(), any()); + verify(keysManager).delete(existingUuid); + verify(keysManager).delete(phoneNumberIdentifiersByE164.get(e164)); verify(keysManager).delete(existingUuid, true); verify(keysManager).delete(phoneNumberIdentifiersByE164.get(e164), true); - verify(messagesManager).clear(existingUuid); - verify(profilesManager).deleteAll(existingUuid); + verify(messagesManager, times(2)).clear(existingUuid); + verify(profilesManager, times(2)).deleteAll(existingUuid); verify(clientPresenceManager).disconnectAllPresencesForUuid(existingUuid); } @@ -1018,7 +1035,7 @@ class AccountsManagerTest { final UUID recentlyDeletedUuid = UUID.randomUUID(); when(accounts.findRecentlyDeletedAccountIdentifier(anyString())).thenReturn(Optional.of(recentlyDeletedUuid)); - when(accounts.create(any(), any())).thenReturn(true); + when(accounts.create(any(), any(), any())).thenReturn(true); final String e164 = "+18005550123"; final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null); @@ -1027,6 +1044,7 @@ class AccountsManagerTest { verify(accounts).create( argThat(account -> e164.equals(account.getNumber()) && recentlyDeletedUuid.equals(account.getUuid())), + any(), any()); verifyNoInteractions(keysManager); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java index 9c5becc3e..ff515ce7a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java @@ -149,6 +149,7 @@ class AccountsManagerUsernameIntegrationTest { experimentEnrollmentManager, mock(RegistrationRecoveryPasswordsManager.class), mock(Executor.class), + mock(Executor.class), mock(Clock.class)); } @@ -312,7 +313,8 @@ class AccountsManagerUsernameIntegrationTest { final Account account = AccountsHelper.createAccount(accountsManager, "+18005551111"); account.setUsernameHash(TestRandomUtil.nextBytes(16)); - accounts.create(account, ignored -> Collections.emptyList()); + accounts.create(account, ignored -> Collections.emptyList(), + (ignoredAci, ignoredPni) -> CompletableFuture.completedFuture(null)); final UUID linkHandle = UUID.randomUUID(); final byte[] encryptedUsername = TestRandomUtil.nextBytes(32); 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 a6ef7faeb..748cb3327 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -209,6 +209,34 @@ class AccountsTest { assertThat(accounts.findRecentlyDeletedAccountIdentifier(account.getNumber())).isEmpty(); } + @Test + void testStoreCleanupFailure() { + final Account existingAccount = nextRandomAccount(); + createAccount(existingAccount); + + verifyStoredState(existingAccount.getNumber(), + existingAccount.getUuid(), + existingAccount.getPhoneNumberIdentifier(), + existingAccount.getUsernameHash().orElse(null), + existingAccount, + true); + + final CompletionException completionException = assertThrows(CompletionException.class, + () -> accounts.create(generateAccount(existingAccount.getNumber(), UUID.randomUUID(), UUID.randomUUID()), + ignored -> Collections.emptyList(), + (aci, pni) -> CompletableFuture.failedFuture(new RuntimeException("OH NO")))); + + assertTrue(completionException.getCause() instanceof RuntimeException); + + // If the existing account cleanup task failed, we should not overwrite the existing account record + verifyStoredState(existingAccount.getNumber(), + existingAccount.getUuid(), + existingAccount.getPhoneNumberIdentifier(), + existingAccount.getUsernameHash().orElse(null), + existingAccount, + true); + } + @Test void testStoreMulti() { final List devices = List.of(generateDevice(DEVICE_ID_1), generateDevice(DEVICE_ID_2)); @@ -1113,7 +1141,8 @@ class AccountsTest { } private boolean createAccount(final Account account) { - return accounts.create(account, ignored -> Collections.emptyList()); + return accounts.create(account, ignored -> Collections.emptyList(), + (ignoredAci, ignoredPni) -> CompletableFuture.completedFuture(null)); } private static Account nextRandomAccount() {