Skip to content

Commit

Permalink
Remove two-stage check of username availability in reserve/confirm
Browse files Browse the repository at this point in the history
  • Loading branch information
@jkt-signal
jkt-signal authored Jan 9, 2024
1 parent ed972a0 commit 184cdc0
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 107 deletions.
50 changes: 22 additions & 28 deletions service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,9 @@ public void changeNumber(final Account account,

/**
* 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 UsernameHashNotAvailableException} if the username was taken by someone else
*/
public CompletableFuture<Void> reserveUsernameHash(
final Account account,
Expand Down Expand Up @@ -504,7 +507,12 @@ public CompletableFuture<Void> reserveUsernameHash(
.build())
.exceptionally(throwable -> {
if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) {
if (e.cancellationReasons().stream().map(CancellationReason::code).anyMatch(CONDITIONAL_CHECK_FAILED::equals)) {
// 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.
if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
} else if (conditionalCheckFailed(e.cancellationReasons().get(1))) {
throw new ContestedOptimisticLockException();
}
}
Expand All @@ -529,7 +537,8 @@ public CompletableFuture<Void> reserveUsernameHash(
* @param account to update
* @param usernameHash believed to be available
* @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 the username has taken by someone else
* {@link ContestedOptimisticLockException} if the account has been updated or
* {@link UsernameHashNotAvailableException} if the username was taken by someone else
*/
public CompletableFuture<Void> confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername) {
final Timer.Sample sample = Timer.start();
Expand Down Expand Up @@ -559,10 +568,15 @@ public CompletableFuture<Void> confirmUsernameHash(final Account account, final
return (Void) null;
})
.exceptionally(throwable -> {
if (ExceptionUtils.unwrap(
throwable) instanceof TransactionCanceledException transactionCanceledException) {
if (transactionCanceledException.cancellationReasons().stream().map(CancellationReason::code)
.anyMatch(CONDITIONAL_CHECK_FAILED::equals)) {
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.
// 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))) {
throw new ContestedOptimisticLockException();
}
}
Expand Down Expand Up @@ -603,6 +617,8 @@ private TransactWriteItemsRequest buildConfirmUsernameHashRequest(final Account
final byte[] usernameHash = updatedAccount.getUsernameHash()
.orElseThrow(() -> new IllegalArgumentException("Account must have a username hash"));

// NOTE: the order in which writeItems are added to the list is significant, and must be kept in sync with the catch block in confirmUsernameHash!

// add the username hash to the constraint table, wiping out the ttl if we had already reserved the hash
writeItems.add(TransactWriteItem.builder()
.put(Put.builder()
Expand Down Expand Up @@ -908,28 +924,6 @@ public CompletionStage<Void> updateTransactionallyAsync(final Account account,
});
}

public CompletableFuture<Boolean> usernameHashAvailable(final byte[] username) {
return usernameHashAvailable(Optional.empty(), username);
}

public CompletableFuture<Boolean> usernameHashAvailable(final Optional<UUID> accountUuid, final byte[] usernameHash) {
return itemByKeyAsync(usernamesConstraintTableName, ATTR_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash))
.thenApply(maybeUsernameHashItem -> maybeUsernameHashItem
.map(item -> {
if (AttributeValues.getLong(item, ATTR_TTL, Long.MAX_VALUE) < clock.instant().getEpochSecond()) {
// username hash was reserved, but has expired
return true;
}

// username hash is reserved by us
return !AttributeValues.getBool(item, ATTR_CONFIRMED, true) && accountUuid
.map(AttributeValues.getUUID(item, KEY_ACCOUNT_UUID, new UUID(0, 0))::equals)
.orElse(false);
})
// If no item was found, then the username hash is free
.orElse(true));
}

@Nonnull
public Optional<Account> getByE164(final String number) {
return getByIndirectLookup(
Expand Down
37 changes: 12 additions & 25 deletions service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -579,23 +579,17 @@ public CompletableFuture<UsernameReservation> reserveUsernameHash(final Account
}

private CompletableFuture<byte[]> checkAndReserveNextUsernameHash(final Account account, final Queue<byte[]> requestedUsernameHashes) {
final byte[] usernameHash;

try {
usernameHash = requestedUsernameHashes.remove();
} catch (final NoSuchElementException e) {
return CompletableFuture.failedFuture(new UsernameHashNotAvailableException());
}

return accounts.usernameHashAvailable(usernameHash)
.thenCompose(usernameHashAvailable -> {
if (usernameHashAvailable) {
return accounts.reserveUsernameHash(account, usernameHash, USERNAME_HASH_RESERVATION_TTL_MINUTES)
.thenApply(ignored -> usernameHash);
} else {
return checkAndReserveNextUsernameHash(account, requestedUsernameHashes);
}
});
final byte[] usernameHash = requestedUsernameHashes.remove();

return accounts.reserveUsernameHash(account, usernameHash, USERNAME_HASH_RESERVATION_TTL_MINUTES)
.thenApply(ignored -> usernameHash)
.exceptionallyComposeAsync(
throwable -> {
if (ExceptionUtils.unwrap(throwable) instanceof UsernameHashNotAvailableException && !requestedUsernameHashes.isEmpty()) {
return checkAndReserveNextUsernameHash(account, requestedUsernameHashes);
}
return CompletableFuture.failedFuture(throwable);
});
}

/**
Expand Down Expand Up @@ -629,14 +623,7 @@ public CompletableFuture<Account> confirmReservedUsernameHash(final Account acco
.thenCompose(ignored -> updateWithRetriesAsync(
account,
a -> true,
a -> accounts.usernameHashAvailable(Optional.of(account.getUuid()), reservedUsernameHash)
.thenCompose(usernameHashAvailable -> {
if (!usernameHashAvailable) {
return CompletableFuture.failedFuture(new UsernameHashNotAvailableException());
}

return accounts.confirmUsernameHash(a, reservedUsernameHash, encryptedUsername);
}),
a -> accounts.confirmUsernameHash(a, reservedUsernameHash, encryptedUsername),
() -> accounts.getByAccountIdentifierAsync(account.getUuid()).thenApply(Optional::orElseThrow),
AccountChangeValidator.USERNAME_CHANGE_VALIDATOR,
MAX_UPDATE_ATTEMPTS
Expand Down
59 changes: 45 additions & 14 deletions service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
import org.whispersystems.textsecuregcm.storage.AccountsManager.UsernameReservation;
import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
import org.whispersystems.textsecuregcm.tests.util.DevicesHelper;
Expand All @@ -86,6 +87,7 @@
import org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil;
import org.whispersystems.textsecuregcm.util.Pair;
import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;

@Timeout(value = 10, threadMode = Timeout.ThreadMode.SEPARATE_THREAD)
class AccountsManagerTest {
Expand Down Expand Up @@ -193,7 +195,6 @@ void setup() throws InterruptedException {

enrollmentManager = mock(ExperimentEnrollmentManager.class);
when(enrollmentManager.isEnrolled(any(UUID.class), eq(AccountsManager.USERNAME_EXPERIMENT_NAME))).thenReturn(true);
when(accounts.usernameHashAvailable(any())).thenReturn(CompletableFuture.completedFuture(true));

final AccountLockManager accountLockManager = mock(AccountLockManager.class);

Expand Down Expand Up @@ -1587,18 +1588,36 @@ void testPniPqUpdate_incompleteKeys() {
@Test
void testReserveUsernameHash() throws UsernameHashNotAvailableException {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
final List<byte[]> usernameHashes = List.of(new byte[32], new byte[32]);
when(accounts.usernameHashAvailable(any())).thenReturn(CompletableFuture.completedFuture(true));
when(accounts.getByAccountIdentifierAsync(account.getUuid())).thenReturn(CompletableFuture.completedFuture(Optional.of(account)));

final List<byte[]> usernameHashes = List.of(TestRandomUtil.nextBytes(32), TestRandomUtil.nextBytes(32));
when(accounts.reserveUsernameHash(any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
accountsManager.reserveUsernameHash(account, usernameHashes);
verify(accounts).reserveUsernameHash(eq(account), eq(new byte[32]), eq(Duration.ofMinutes(5)));

UsernameReservation result = accountsManager.reserveUsernameHash(account, usernameHashes).join();
assertArrayEquals(usernameHashes.get(0), result.reservedUsernameHash());
verify(accounts, times(1)).reserveUsernameHash(eq(account), any(), eq(Duration.ofMinutes(5)));
}

@Test
void testReserveUsernameHashNotAvailable() {
void testReserveUsernameOptimisticLockingFailure() throws UsernameHashNotAvailableException {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
when(accounts.usernameHashAvailable(any())).thenReturn(CompletableFuture.completedFuture(false));
when(accounts.getByAccountIdentifierAsync(account.getUuid())).thenReturn(CompletableFuture.completedFuture(Optional.of(account)));

final List<byte[]> usernameHashes = List.of(TestRandomUtil.nextBytes(32), TestRandomUtil.nextBytes(32));
when(accounts.reserveUsernameHash(any(), any(), any()))
.thenReturn(CompletableFuture.failedFuture(new ContestedOptimisticLockException()))
.thenReturn(CompletableFuture.completedFuture(null));

UsernameReservation result = accountsManager.reserveUsernameHash(account, usernameHashes).join();
assertArrayEquals(usernameHashes.get(0), result.reservedUsernameHash());
verify(accounts, times(2)).reserveUsernameHash(eq(account), any(), eq(Duration.ofMinutes(5)));
}

@Test
void testReserveUsernameHashNotAvailable() {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
when(accounts.getByAccountIdentifierAsync(account.getUuid())).thenReturn(CompletableFuture.completedFuture(Optional.of(account)));
when(accounts.reserveUsernameHash(any(), any(), any())).thenReturn(CompletableFuture.failedFuture(new UsernameHashNotAvailableException()));
CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class,
accountsManager.reserveUsernameHash(account, List.of(USERNAME_HASH_1, USERNAME_HASH_2)));
}
Expand All @@ -1615,8 +1634,6 @@ void testReserveUsernameDisabled() {
void testConfirmReservedUsernameHash() throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
setReservationHash(account, USERNAME_HASH_1);
when(accounts.usernameHashAvailable(Optional.of(account.getUuid()), USERNAME_HASH_1))
.thenReturn(CompletableFuture.completedFuture(true));

when(accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1))
.thenReturn(CompletableFuture.completedFuture(null));
Expand All @@ -1625,12 +1642,26 @@ void testConfirmReservedUsernameHash() throws UsernameHashNotAvailableException,
verify(accounts).confirmUsernameHash(eq(account), eq(USERNAME_HASH_1), eq(ENCRYPTED_USERNAME_1));
}

@Test
void testConfirmReservedUsernameHashOptimisticLockingFailure() throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
setReservationHash(account, USERNAME_HASH_1);
when(accounts.getByAccountIdentifierAsync(account.getUuid())).thenReturn(CompletableFuture.completedFuture(Optional.of(account)));

when(accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1))
.thenReturn(CompletableFuture.failedFuture(new ContestedOptimisticLockException()))
.thenReturn(CompletableFuture.completedFuture(null));

accountsManager.confirmReservedUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
verify(accounts, times(2)).confirmUsernameHash(eq(account), eq(USERNAME_HASH_1), eq(ENCRYPTED_USERNAME_1));
}

@Test
void testConfirmReservedHashNameMismatch() {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
setReservationHash(account, USERNAME_HASH_1);
when(accounts.usernameHashAvailable(eq(Optional.of(account.getUuid())), eq(USERNAME_HASH_1)))
.thenReturn(CompletableFuture.completedFuture(true));
when(accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1))
.thenReturn(CompletableFuture.completedFuture(null));
CompletableFutureTestUtil.assertFailsWithCause(UsernameReservationNotFoundException.class,
accountsManager.confirmReservedUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_2));
}
Expand All @@ -1640,11 +1671,11 @@ void testConfirmReservedLapsed() {
final Account account = AccountsHelper.generateTestAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
// hash was reserved, but the reservation lapsed and another account took it
setReservationHash(account, USERNAME_HASH_1);
when(accounts.usernameHashAvailable(eq(Optional.of(account.getUuid())), eq(USERNAME_HASH_1)))
.thenReturn(CompletableFuture.completedFuture(false));
when(accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1))
.thenReturn(CompletableFuture.failedFuture(new UsernameHashNotAvailableException()));
CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class,
accountsManager.confirmReservedUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1));
verify(accounts, never()).confirmUsernameHash(any(), any(), any());
assertTrue(account.getUsernameHash().isEmpty());
}

@Test
Expand Down
13 changes: 2 additions & 11 deletions ...java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ void testNoUsernames() throws InterruptedException {
}

@Test
void testReserveUsernameSnatched() throws InterruptedException, UsernameHashNotAvailableException {
void testReserveUsernameGetFirstAvailableChoice() throws InterruptedException, UsernameHashNotAvailableException {
final Account account = AccountsHelper.createAccount(accountsManager, "+18005551111");

ArrayList<byte[]> usernameHashes = new ArrayList<>(Arrays.asList(USERNAME_HASH_1, USERNAME_HASH_2));
Expand All @@ -198,23 +198,14 @@ void testReserveUsernameSnatched() throws InterruptedException, UsernameHashNotA

byte[] availableHash = TestRandomUtil.nextBytes(32);
usernameHashes.add(availableHash);
usernameHashes.add(TestRandomUtil.nextBytes(32));

// first time this is called lie and say the username is available
// this simulates seeing an available username and then it being taken
// by someone before the write
doReturn(CompletableFuture.completedFuture(true))
.doCallRealMethod()
.when(accounts).usernameHashAvailable(any());
final byte[] username = accountsManager
.reserveUsernameHash(account, usernameHashes)
.join()
.reservedUsernameHash();

assertArrayEquals(username, availableHash);

// 1 attempt on first try (returns true),
// 5 more attempts until "availableHash" returns true
verify(accounts, times(4)).usernameHashAvailable(any());
}

@Test
Expand Down
Loading

0 comments on commit 184cdc0

Please sign in to comment.