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 1b3a406b5..06e280c73 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java @@ -455,7 +455,8 @@ public class Accounts extends AbstractDynamoDbStore { /** * Reserve a username hash under the account UUID * @return a future that completes once the username hash has been reserved; may fail with an - * {@link ContestedOptimisticLockException} if the account has been updated or + * {@link ContestedOptimisticLockException} if the account has been updated or there are concurrent updates to the + * account or constraint records, and with an * {@link UsernameHashNotAvailableException} if the username was taken by someone else */ public CompletableFuture reserveUsernameHash( @@ -523,11 +524,12 @@ public class Accounts extends AbstractDynamoDbStore { .exceptionally(throwable -> { if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { // If the constraint table update failed the condition check, the username's taken and we should stop - // trying. However if it was only in the accounts table that the condition check update failed, it's an - // optimistic locking failure (the account was concurrently updated) and we should try again. + // trying. However, if the accounts table fails the conditional check or + // either table was concurrently updated, it's an optimistic locking failure and we should try again. if (conditionalCheckFailed(e.cancellationReasons().get(0))) { throw ExceptionUtils.wrap(new UsernameHashNotAvailableException()); - } else if (conditionalCheckFailed(e.cancellationReasons().get(1))) { + } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) || + e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { throw new ContestedOptimisticLockException(); } } @@ -551,8 +553,10 @@ public class Accounts extends AbstractDynamoDbStore { * * @param account to update * @param usernameHash believed to be available + * @param encryptedUsername the encrypted form of the previously reserved username; used for the username link * @return a future that completes once the username hash has been confirmed; may fail with an - * {@link ContestedOptimisticLockException} if the account has been updated or + * {@link ContestedOptimisticLockException} if the account has been updated or there are concurrent updates to the + * account or constraint records, and with an * {@link UsernameHashNotAvailableException} if the username was taken by someone else */ public CompletableFuture confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername) { @@ -585,13 +589,14 @@ public class Accounts extends AbstractDynamoDbStore { .exceptionally(throwable -> { if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { // If the constraint table update failed the condition check, the username's taken and we should stop - // trying. However if it was only in the accounts table that the condition check update failed, it's an - // optimistic locking failure (the account was concurrently updated) and we should try again. + // trying. However, if the accounts table fails the conditional check or + // either table was concurrently updated, it's an optimistic locking failure and we should try again. // NOTE: the fixed indices here must be kept in sync with the creation of the TransactWriteItems in // buildConfirmUsernameHashRequest! if (conditionalCheckFailed(e.cancellationReasons().get(0))) { throw ExceptionUtils.wrap(new UsernameHashNotAvailableException()); - } else if (conditionalCheckFailed(e.cancellationReasons().get(1))) { + } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) || + e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { throw new ContestedOptimisticLockException(); } } @@ -690,6 +695,14 @@ public class Accounts extends AbstractDynamoDbStore { .build(); } + /** + * Clear the username hash and link from the given account + * + * @param account to update + * @return a future that completes once the username data has been cleared; + * it can fail with a {@link ContestedOptimisticLockException} if there are concurrent updates + * to the account or username constraint records. + */ public CompletableFuture clearUsernameHash(final Account account) { return account.getUsernameHash().map(usernameHash -> { final Timer.Sample sample = Timer.start(); @@ -714,8 +727,9 @@ public class Accounts extends AbstractDynamoDbStore { account.setVersion(account.getVersion() + 1); }) .exceptionally(throwable -> { - if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException transactionCanceledException) { - if (conditionalCheckFailed(transactionCanceledException.cancellationReasons().get(0))) { + if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { + if (conditionalCheckFailed(e.cancellationReasons().get(0)) || + e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { throw new ContestedOptimisticLockException(); } } @@ -1371,4 +1385,8 @@ public class Accounts extends AbstractDynamoDbStore { private static boolean conditionalCheckFailed(final CancellationReason reason) { return CONDITIONAL_CHECK_FAILED.equals(reason.code()); } + + private static boolean isTransactionConflict(final CancellationReason reason) { + return TRANSACTION_CONFLICT.equals(reason.code()); + } } 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 7cc18ea96..6fe6e2bd0 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -73,6 +73,7 @@ import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest; +import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsResponse; import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException; import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException; import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest; @@ -960,7 +961,7 @@ class AccountsTest { } @Test - void testUsernameHashConflict() { + void testUsernameHashNotAvailable() { final Account firstAccount = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID()); final Account secondAccount = generateAccount("+18005559876", UUID.randomUUID(), UUID.randomUUID()); @@ -994,6 +995,92 @@ class AccountsTest { assertThat(secondAccount.getUsernameHash()).isEmpty(); } + @ParameterizedTest + @MethodSource + void testReserveUsernameHashTransactionConflict(final Optional constraintCancellationString, + final Optional accountsCancellationString, + final Class expectedException) { + final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); + + accounts = new Accounts(mock(DynamoDbClient.class), + dbAsyncClient, + Tables.ACCOUNTS.tableName(), + Tables.NUMBERS.tableName(), + Tables.PNI_ASSIGNMENTS.tableName(), + Tables.USERNAMES.tableName(), + Tables.DELETED_ACCOUNTS.tableName()); + final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID()); + createAccount(account); + + final CancellationReason constraintCancellationReason = constraintCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + final CancellationReason accountsCancellationReason = accountsCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class))) + .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder() + .cancellationReasons(constraintCancellationReason, accountsCancellationReason) + .build())); + + CompletableFutureTestUtil.assertFailsWithCause(expectedException, + accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1))); + } + + private static Stream testReserveUsernameHashTransactionConflict() { + return Stream.of( + Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class), + Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class), + Arguments.of(Optional.of("ConditionalCheckFailed"), Optional.of("TransactionConflict"), UsernameHashNotAvailableException.class) + ); + } + + @ParameterizedTest + @MethodSource + void testConfirmUsernameHashTransactionConflict(final Optional constraintCancellationString, + final Optional accountsCancellationString, + final Class expectedException) { + final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); + + accounts = new Accounts(mock(DynamoDbClient.class), + dbAsyncClient, + Tables.ACCOUNTS.tableName(), + Tables.NUMBERS.tableName(), + Tables.PNI_ASSIGNMENTS.tableName(), + Tables.USERNAMES.tableName(), + Tables.DELETED_ACCOUNTS.tableName()); + final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID()); + createAccount(account); + + final CancellationReason constraintCancellationReason = constraintCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + final CancellationReason accountsCancellationReason = accountsCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class))) + .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder() + .cancellationReasons(constraintCancellationReason, + accountsCancellationReason, + CancellationReason.builder().build()) + .build())); + + CompletableFutureTestUtil.assertFailsWithCause(expectedException, + accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1)); + } + + private static Stream testConfirmUsernameHashTransactionConflict() { + return Stream.of( + Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class), + Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class), + Arguments.of(Optional.of("ConditionalCheckFailed"), Optional.of("TransactionConflict"), UsernameHashNotAvailableException.class) + ); + } + @Test void testConfirmUsernameHashVersionMismatch() { final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID()); @@ -1051,6 +1138,55 @@ class AccountsTest { assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1); } + @ParameterizedTest + @MethodSource + void testClearUsernameTransactionConflict(final Optional constraintCancellationString, + final Optional accountsCancellationString) { + final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); + + accounts = new Accounts(mock(DynamoDbClient.class), + dbAsyncClient, + Tables.ACCOUNTS.tableName(), + Tables.NUMBERS.tableName(), + Tables.PNI_ASSIGNMENTS.tableName(), + Tables.USERNAMES.tableName(), + Tables.DELETED_ACCOUNTS.tableName()); + + final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID()); + createAccount(account); + + when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class))) + .thenReturn(CompletableFuture.completedFuture(mock(TransactWriteItemsResponse.class))); + + accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join(); + accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join(); + + final CancellationReason constraintCancellationReason = constraintCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + final CancellationReason accountsCancellationReason = accountsCancellationString.map( + reason -> CancellationReason.builder().code(reason).build() + ).orElse(CancellationReason.builder().build()); + + when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class))) + .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder() + .cancellationReasons(accountsCancellationReason, constraintCancellationReason) + .build())); + + CompletableFutureTestUtil.assertFailsWithCause(ContestedOptimisticLockException.class, + accounts.clearUsernameHash(account)); + + assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1); + } + + private static Stream testClearUsernameTransactionConflict() { + return Stream.of( + Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class), + Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class) + ); + } + @Test void testReservedUsernameHash() { final Account account1 = generateAccount("+18005551111", UUID.randomUUID(), UUID.randomUUID());