Temporarily hold a username after an account releases it

This commit is contained in:
Ravi Khadiwala 2024-03-01 17:22:11 -06:00 committed by ravi-signal
parent 47b24b5dff
commit 3cc740cda3
7 changed files with 744 additions and 40 deletions

View File

@ -238,6 +238,7 @@ import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand;
import org.whispersystems.textsecuregcm.workers.ProcessPushNotificationFeedbackCommand;
import org.whispersystems.textsecuregcm.workers.RemoveExpiredAccountsCommand;
import org.whispersystems.textsecuregcm.workers.RemoveExpiredLinkedDevicesCommand;
import org.whispersystems.textsecuregcm.workers.RemoveExpiredUsernameHoldsCommand;
import org.whispersystems.textsecuregcm.workers.ScheduledApnPushNotificationSenderServiceCommand;
import org.whispersystems.textsecuregcm.workers.ServerVersionCommand;
import org.whispersystems.textsecuregcm.workers.SetRequestLoggingEnabledTask;
@ -297,6 +298,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
bootstrap.addCommand(new ScheduledApnPushNotificationSenderServiceCommand());
bootstrap.addCommand(new MessagePersisterServiceCommand());
bootstrap.addCommand(new RemoveExpiredAccountsCommand(Clock.systemUTC()));
bootstrap.addCommand(new RemoveExpiredUsernameHoldsCommand(Clock.systemUTC()));
bootstrap.addCommand(new ProcessPushNotificationFeedbackCommand(Clock.systemUTC()));
bootstrap.addCommand(new RemoveExpiredLinkedDevicesCommand());
}

View File

@ -13,6 +13,7 @@ import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import java.time.Clock;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
@ -106,9 +107,14 @@ public class Account {
@JsonProperty
private int version;
@JsonProperty("holds")
private List<UsernameHold> usernameHolds = Collections.emptyList();
@JsonIgnore
private boolean stale;
public record UsernameHold(@JsonProperty("uh") byte[] usernameHash, @JsonProperty("e") long expirationSecs) {}
public UUID getIdentifier(final IdentityType identityType) {
return switch (identityType) {
case ACI -> getUuid();
@ -525,6 +531,15 @@ public class Account {
devices.forEach(Device::lockAuthTokenHash);
}
public List<UsernameHold> getUsernameHolds() {
return Collections.unmodifiableList(usernameHolds);
}
public void setUsernameHolds(final List<UsernameHold> usernameHolds) {
this.requireNotStale();
this.usernameHolds = usernameHolds;
}
boolean isStale() {
return stale;
}

View File

@ -18,6 +18,7 @@ import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
@ -98,6 +99,7 @@ public class Accounts extends AbstractDynamoDbStore {
private static final Timer GET_BY_PNI_TIMER = Metrics.timer(name(Accounts.class, "getByPni"));
private static final Timer GET_BY_UUID_TIMER = Metrics.timer(name(Accounts.class, "getByUuid"));
private static final Timer DELETE_TIMER = Metrics.timer(name(Accounts.class, "delete"));
private static final String USERNAME_HOLD_ADDED_COUNTER_NAME = name(Accounts.class, "usernameHoldAdded");
private static final String CONDITIONAL_CHECK_FAILED = "ConditionalCheckFailed";
@ -132,6 +134,18 @@ public class Accounts extends AbstractDynamoDbStore {
static final Duration DELETED_ACCOUNTS_TIME_TO_LIVE = Duration.ofDays(30);
/**
* Maximum number of temporary username holds an account can have on recently used usernames
*/
@VisibleForTesting
static final int MAX_USERNAME_HOLDS = 3;
/**
* How long an old username is held for an account after the account initially clears/switches the username
*/
@VisibleForTesting
static final Duration USERNAME_HOLD_DURATION = Duration.ofDays(7);
private final Clock clock;
private final DynamoDbAsyncClient asyncClient;
@ -456,6 +470,7 @@ 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
@ -470,14 +485,61 @@ public class Accounts extends AbstractDynamoDbStore {
final Timer.Sample sample = Timer.start();
// if there is an existing old reservation it will be cleaned up via ttl
// if there is an existing old reservation it will be cleaned up via ttl. Save it so we can restore it to the local
// account if the update fails though.
final Optional<byte[]> maybeOriginalReservation = account.getReservedUsernameHash();
account.setReservedUsernameHash(reservedUsernameHash);
// Normally when a username is reserved for the first time we reserve it for the provided TTL. But if the
// reservation is for a username that we already have a reservation for (for example, if it's reclaimable, or there
// is a hold) we might own that reservation for longer anyways, so we should preserve the original TTL in that case.
// What we'd really like to do is set expirationTime = max(oldExpirationTime, now + ttl), but dynamodb doesn't
// support that. Instead, we'll set expiration if it's greater than the existing expiration, otherwise retry
final long expirationTime = clock.instant().plus(ttl).getEpochSecond();
return tryReserveUsernameHash(account, reservedUsernameHash, expirationTime)
.exceptionallyCompose(ExceptionUtils.exceptionallyHandler(TtlConflictException.class, ttlConflict ->
// retry (once) with the returned expiration time
tryReserveUsernameHash(account, reservedUsernameHash, ttlConflict.getExistingExpirationSeconds())))
.whenComplete((response, throwable) -> {
sample.stop(RESERVE_USERNAME_TIMER);
if (throwable == null) {
account.setVersion(account.getVersion() + 1);
} else {
account.setReservedUsernameHash(maybeOriginalReservation.orElse(null));
}
});
}
private static class TtlConflictException extends ContestedOptimisticLockException {
private final long existingExpirationSeconds;
TtlConflictException(final long existingExpirationSeconds) {
super();
this.existingExpirationSeconds = existingExpirationSeconds;
}
long getExistingExpirationSeconds() {
return existingExpirationSeconds;
}
}
/**
* Try to reserve the provided usernameHash
*
* @param updatedAccount The account, already updated to reserve the provided usernameHash
* @param reservedUsernameHash The usernameHash to reserve
* @param expirationTimeSeconds When the reservation should expire
* @return A future that completes successfully if the usernameHash was reserved
* @throws TtlConflictException if the usernameHash was already reserved but with a longer TTL. The operation should
* be retried with the returned {@link TtlConflictException#getExistingExpirationSeconds()}
*/
private CompletableFuture<Void> tryReserveUsernameHash(
final Account updatedAccount,
final byte[] reservedUsernameHash,
final long expirationTimeSeconds) {
// Use account UUID as a "reservation token" - by providing this, the client proves ownership of the hash
final UUID uuid = account.getUuid();
final UUID uuid = updatedAccount.getUuid();
final List<TransactWriteItem> writeItems = new ArrayList<>();
@ -487,10 +549,13 @@ public class Accounts extends AbstractDynamoDbStore {
.item(Map.of(
UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(reservedUsernameHash),
UsernameTable.ATTR_TTL, AttributeValues.fromLong(expirationTime),
UsernameTable.ATTR_TTL, AttributeValues.fromLong(expirationTimeSeconds),
UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(false),
UsernameTable.ATTR_RECLAIMABLE, AttributeValues.fromBool(false)))
.conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :confirmed)")
// we can make a reservation if no reservation exists for the name, or that reservation is expired, or there
// is a reservation but it's ours and we haven't confirmed it yet and we're not accidentally reducing our
// reservation's TTL. Note that confirmed=false => a TTL exists
.conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :false AND #ttl <= :expirationTime)")
.expressionAttributeNames(Map.of(
"#username_hash", UsernameTable.KEY_USERNAME_HASH,
"#ttl", UsernameTable.ATTR_TTL,
@ -499,37 +564,134 @@ public class Accounts extends AbstractDynamoDbStore {
.expressionAttributeValues(Map.of(
":now", AttributeValues.fromLong(clock.instant().getEpochSecond()),
":aci", AttributeValues.fromUUID(uuid),
":confirmed", AttributeValues.fromBool(false)))
":false", AttributeValues.fromBool(false),
":expirationTime", AttributeValues.fromLong(expirationTimeSeconds)))
.returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
.build())
.build());
writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, account).transactItem());
writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem());
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
return asyncClient
.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
.thenRun(Util.NOOP)
.exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> {
// If the constraint table update failed the condition check, the username's taken and we should stop
// 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.
// trying. However,
if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
// The constraint table update failed the condition check. It could be because the username was taken,
// or because we need to retry with a longer TTL
final Map<String, AttributeValue> item = e.cancellationReasons().get(0).item();
final UUID existingOwner = AttributeValues.getUUID(item, UsernameTable.ATTR_ACCOUNT_UUID, null);
final boolean confirmed = AttributeValues.getBool(item, UsernameTable.ATTR_CONFIRMED, false);
final long existingTtl = AttributeValues.getLong(item, UsernameTable.ATTR_TTL, 0L);
if (uuid.equals(existingOwner) && !confirmed && existingTtl > expirationTimeSeconds) {
// We failed because we provided a shorter TTL than the one that exists on the reservation. The caller
// can retry with updated expiration time.
throw new TtlConflictException(existingTtl);
}
throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
} else if (conditionalCheckFailed(e.cancellationReasons().get(1)) ||
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
// The accounts table fails the conditional check or either table was concurrently updated, it's an
// optimistic locking failure and we should try again.
throw new ContestedOptimisticLockException();
} else {
throw ExceptionUtils.wrap(e);
}
}))
.whenComplete((response, throwable) -> {
sample.stop(RESERVE_USERNAME_TIMER);
}));
}
if (throwable == null) {
account.setVersion(account.getVersion() + 1);
} else {
account.setReservedUsernameHash(maybeOriginalReservation.orElse(null));
}
})
.thenRun(() -> {});
/**
* Add a held usernameHash to the account object.
* <p>
* An account may only have up to MAX_USERNAME_HOLDS held usernames. If adding this hold pushes the account over this
* limit, a usernameHash is returned that the caller must release their hold on.
* <p>
* This only tracks the holds associated with the account, ensuring that no other account can take a held username is
* done via the username constraint table, and should be done transactionally with writing the updated account.
*
* @param accountToUpdate The account to update (in-place)
* @param newHold A username hash to add to the account's holds
* @param now The current time
* @return If present, an old hold that the caller should remove from the username constraint table
*/
private Optional<byte[]> addToHolds(final Account accountToUpdate, final byte[] newHold, final Instant now) {
List<Account.UsernameHold> holds = new ArrayList<>(accountToUpdate.getUsernameHolds());
final Account.UsernameHold holdToAdd = new Account.UsernameHold(newHold,
now.plus(USERNAME_HOLD_DURATION).getEpochSecond());
// Remove any holds that are
// - expired
// - match what we're trying to add (we'll re-add it at the end of the list to refresh the ttl)
// - match our current username
holds.removeIf(hold -> hold.expirationSecs() < now.getEpochSecond()
|| Arrays.equals(newHold, hold.usernameHash())
|| accountToUpdate.getUsernameHash().map(curr -> Arrays.equals(curr, hold.usernameHash())).orElse(false));
// add the new hold
holds.add(holdToAdd);
if (holds.size() <= MAX_USERNAME_HOLDS) {
accountToUpdate.setUsernameHolds(holds);
Metrics.counter(USERNAME_HOLD_ADDED_COUNTER_NAME, "max", String.valueOf(false)).increment();
return Optional.empty();
} else {
accountToUpdate.setUsernameHolds(holds.subList(1, holds.size()));
Metrics.counter(USERNAME_HOLD_ADDED_COUNTER_NAME, "max", String.valueOf(true)).increment();
// Newer holds are always added to the end of the holds list, so the first hold is always the oldest hold. Note
// that if a duplicate hold is added, we remove it from the list and re-add it at the end, this preserves hold
// ordering
return Optional.of(holds.getFirst().usernameHash());
}
}
/**
* Transaction item to update the usernameConstraintTable to "hold" a usernameHash for an account
*
* @param holder The account with the hold.
* @param usernameHash The hash to reserve for the account
* @param now The current time
* @return A transaction item that will update the usernameConstraintTable.
*/
private TransactWriteItem holdUsernameTransactItem(final UUID holder, final byte[] usernameHash, final Instant now) {
return TransactWriteItem.builder().put(Put.builder()
.tableName(usernamesConstraintTableName)
.item(Map.of(
UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash),
UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(holder),
UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(false),
UsernameTable.ATTR_TTL,
AttributeValues.fromLong(now.plus(USERNAME_HOLD_DURATION).getEpochSecond())))
.build()).build();
}
/**
* Transaction item to release a hold on the usernameConstraintTable
*
* @param holder The account with the hold.
* @param usernameHashToRelease The hash to release for the account
* @param now The current time
* @return A transaction item that will update the usernameConstraintTable. The transaction will fail with a condition
* exception if someone else has a reservation for usernameHashToRelease
*/
private TransactWriteItem releaseHoldIfAllowedTransactItem(
final UUID holder, final byte[] usernameHashToRelease, final Instant now) {
return TransactWriteItem.builder().delete(Delete.builder()
.tableName(usernamesConstraintTableName)
.key(Map.of(UsernameTable.KEY_USERNAME_HASH, AttributeValues.b(usernameHashToRelease)))
// we can release the hold if we own it (and it's not our confirmed username) or if no one owns it
.conditionExpression("(#aci = :aci AND #confirmed = :false) OR #ttl < :now OR attribute_not_exists(#usernameHash)")
.expressionAttributeNames(Map.of(
"#usernameHash", UsernameTable.KEY_USERNAME_HASH,
"#aci", UsernameTable.ATTR_ACCOUNT_UUID,
"#confirmed", UsernameTable.ATTR_CONFIRMED,
"#ttl", UsernameTable.ATTR_TTL))
.expressionAttributeValues(Map.of(
":aci", AttributeValues.b(holder),
":now", AttributeValues.n(now.getEpochSecond()),
":false", AttributeValues.fromBool(false)))
.build()).build();
}
/**
@ -551,12 +713,14 @@ public class Accounts extends AbstractDynamoDbStore {
return pickLinkHandle(account, usernameHash)
.thenCompose(linkHandle -> {
final Optional<byte[]> maybeOriginalUsernameHash = account.getUsernameHash();
final Account updatedAccount = AccountUtil.cloneAccountAsNotStale(account);
updatedAccount.setUsernameHash(usernameHash);
updatedAccount.setReservedUsernameHash(null);
updatedAccount.setUsernameLinkDetails(encryptedUsername == null ? null : linkHandle, encryptedUsername);
final Optional<byte[]> maybeOriginalUsernameHash = account.getUsernameHash();
final Instant now = clock.instant();
final Optional<byte[]> holdToRemove = maybeOriginalUsernameHash
.flatMap(hold -> addToHolds(updatedAccount, hold, now));
final List<TransactWriteItem> writeItems = new ArrayList<>();
@ -585,18 +749,22 @@ public class Accounts extends AbstractDynamoDbStore {
// 1: update the account object (conditioned on the version increment)
writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem());
// 2?: remove the old username hash (if it exists) from the username constraint table
maybeOriginalUsernameHash.ifPresent(originalUsernameHash -> writeItems.add(
buildDelete(usernamesConstraintTableName, UsernameTable.KEY_USERNAME_HASH, originalUsernameHash)));
// 2?: Add a temporary hold for the old username to stop others from claiming it
maybeOriginalUsernameHash.ifPresent(originalUsernameHash ->
writeItems.add(holdUsernameTransactItem(updatedAccount.getUuid(), originalUsernameHash, now)));
// 3?: Adding that hold may have caused our account to exceed our maximum holds. Release an old hold
holdToRemove.ifPresent(oldHold ->
writeItems.add(releaseHoldIfAllowedTransactItem(updatedAccount.getUuid(), oldHold, now)));
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
.thenApply(ignored -> linkHandle);
.thenApply(ignored -> updatedAccount);
})
.thenApply(linkHandle -> {
.thenApply(updatedAccount -> {
account.setUsernameHash(usernameHash);
account.setReservedUsernameHash(null);
account.setUsernameLinkDetails(encryptedUsername == null ? null : linkHandle, encryptedUsername);
account.setUsernameLinkDetails(updatedAccount.getUsernameLinkHandle(), updatedAccount.getEncryptedUsername().orElse(null));
account.setUsernameHolds(updatedAccount.getUsernameHolds());
account.setVersion(account.getVersion() + 1);
return (Void) null;
})
@ -606,8 +774,14 @@ public class Accounts extends AbstractDynamoDbStore {
// 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)) ||
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
} else if (conditionalCheckFailed(e.cancellationReasons().get(1)) // Account version conflict
// When we looked at the holds on our account, we thought we still held the corresponding username
// reservation. But it turned out that someone else has taken the reservation since. This means that the
// TTL on the hold must have just expired, so if we retry we should see that our hold is expired, and we
// won't try to remove it again.
|| (e.cancellationReasons().size() > 3 && conditionalCheckFailed(e.cancellationReasons().get(3)))
// concurrent update on any table
|| e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
throw new ContestedOptimisticLockException();
} else {
throw ExceptionUtils.wrap(e);
@ -659,21 +833,36 @@ public class Accounts extends AbstractDynamoDbStore {
updatedAccount.setUsernameHash(null);
updatedAccount.setUsernameLinkDetails(null, null);
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
.transactItems(List.of(
// 0: remove the username from the account object, conditioned on account version
UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem(),
// 1: remote the username from the constraint table
buildDelete(usernamesConstraintTableName, UsernameTable.KEY_USERNAME_HASH, usernameHash)))
.build())
final Instant now = clock.instant();
final Optional<byte[]> holdToRemove = addToHolds(updatedAccount, usernameHash, now);
final List<TransactWriteItem> items = new ArrayList<>();
// 0: remove the username from the account object, conditioned on account version
items.add(UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem());
// 1: Un-confirm our username, adding a temporary hold for the old username to stop others from claiming it
items.add(holdUsernameTransactItem(updatedAccount.getUuid(), usernameHash, now));
// 2?: Adding that hold may have caused our account to exceed our maximum holds. Release an old hold
holdToRemove.ifPresent(oldHold -> items.add(releaseHoldIfAllowedTransactItem(updatedAccount.getUuid(), oldHold, now)));
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(items).build())
.thenAccept(ignored -> {
account.setUsernameHash(null);
account.setUsernameLinkDetails(null, null);
account.setVersion(account.getVersion() + 1);
account.setUsernameHolds(updatedAccount.getUsernameHolds());
})
.exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> {
if (conditionalCheckFailed(e.cancellationReasons().get(0)) ||
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
if (conditionalCheckFailed(e.cancellationReasons().get(0)) // Account version conflict
// When we looked at the holds on our account, we thought we still held the corresponding username
// reservation. But it turned out that someone else has taken the reservation since. This means that the
// TTL on the hold must have just expired, so if we retry we should see that our hold is expired, and we
// won't try to remove it again.
|| (e.cancellationReasons().size() > 2 && conditionalCheckFailed(e.cancellationReasons().get(2)))
// concurrent update on any table
|| e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
throw new ContestedOptimisticLockException();
} else {
throw ExceptionUtils.wrap(e);

View File

@ -0,0 +1,114 @@
/*
* Copyright 2024 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.workers;
import com.google.common.annotations.VisibleForTesting;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Metrics;
import net.sourceforge.argparse4j.inf.Subparser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.util.Util;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
public class RemoveExpiredUsernameHoldsCommand extends AbstractSinglePassCrawlAccountsCommand {
private final Clock clock;
@VisibleForTesting
static final String DRY_RUN_ARGUMENT = "dry-run";
@VisibleForTesting
static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency";
private static final int DEFAULT_MAX_CONCURRENCY = 16;
private static final String DELETED_HOLDS_COUNTER_NAME =
name(RemoveExpiredUsernameHoldsCommand.class, "expiredHolds");
private static final String UPDATED_ACCOUNTS_COUNTER_NAME =
name(RemoveExpiredUsernameHoldsCommand.class, "accountsWithExpiredHolds");
private static final Logger log = LoggerFactory.getLogger(RemoveExpiredUsernameHoldsCommand.class);
public RemoveExpiredUsernameHoldsCommand(final Clock clock) {
super("remove-expired-username-holds", "Removes expired username holds from account records");
this.clock = clock;
}
@Override
public void configure(final Subparser subparser) {
super.configure(subparser);
subparser.addArgument("--dry-run")
.type(Boolean.class)
.dest(DRY_RUN_ARGUMENT)
.required(false)
.setDefault(true)
.help("If true, don't actually delete holds");
subparser.addArgument("--max-concurrency")
.type(Integer.class)
.dest(MAX_CONCURRENCY_ARGUMENT)
.setDefault(DEFAULT_MAX_CONCURRENCY)
.help("Max concurrency for DynamoDB operations");
}
@Override
protected void crawlAccounts(final Flux<Account> accounts) {
final boolean isDryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT);
final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT);
final Counter deletedHoldsCounter =
Metrics.counter(DELETED_HOLDS_COUNTER_NAME, "dryRun", String.valueOf(isDryRun));
final Counter updatedAccountsCounter =
Metrics.counter(UPDATED_ACCOUNTS_COUNTER_NAME, "dryRun", String.valueOf(isDryRun));
final AccountsManager accountManager = getCommandDependencies().accountsManager();
accounts.flatMap(account -> {
final List<Account.UsernameHold> holds = new ArrayList<>(account.getUsernameHolds());
int holdsToRemove = removeExpired(holds);
final Mono<Void> purgeMono = isDryRun || holdsToRemove == 0
? Mono.empty()
: Mono.fromFuture(() ->
accountManager.updateAsync(account, a -> a.setUsernameHolds(holds)).thenRun(Util.NOOP));
return purgeMono
.doOnSuccess(ignored -> {
deletedHoldsCounter.increment(holdsToRemove);
updatedAccountsCounter.increment();
})
.onErrorResume(throwable -> {
log.warn("Failed to purge {} expired holds on account {}", holdsToRemove, account.getUuid());
return Mono.empty();
});
}, maxConcurrency)
.then().block();
}
@VisibleForTesting
int removeExpired(final List<Account.UsernameHold> holds) {
final Instant now = this.clock.instant();
int holdsToRemove = 0;
for (Iterator<Account.UsernameHold> it = holds.iterator(); it.hasNext(); ) {
if (it.next().expirationSecs() < now.getEpochSecond()) {
holdsToRemove++;
it.remove();
}
}
return holdsToRemove;
}
}

View File

@ -239,6 +239,32 @@ class AccountsManagerUsernameIntegrationTest {
assertThat(accountsManager.getByAccountIdentifier(account.getUuid()).orElseThrow().getUsernameHash()).isEmpty();
}
@Test
public void testHold() throws InterruptedException {
Account account = AccountsHelper.createAccount(accountsManager, "+18005551111");
AccountsManager.UsernameReservation reservation =
accountsManager.reserveUsernameHash(account, List.of(USERNAME_HASH_1)).join();
// confirm
account = accountsManager.confirmReservedUsernameHash(
reservation.account(),
reservation.reservedUsernameHash(),
ENCRYPTED_USERNAME_1).join();
// clear
account = accountsManager.clearUsernameHash(account).join();
assertThat(accountsManager.getByUsernameHash(USERNAME_HASH_1).join()).isEmpty();
assertThat(accountsManager.getByAccountIdentifier(account.getUuid()).orElseThrow().getUsernameHash()).isEmpty();
assertThat(accountsManager.getByUsernameHash(reservation.reservedUsernameHash()).join()).isEmpty();
Account account2 = AccountsHelper.createAccount(accountsManager, "+18005552222");
CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class,
accountsManager.reserveUsernameHash(account2, List.of(USERNAME_HASH_1)),
"account2 should not be able to reserve a held hash");
}
@Test
public void testReservationLapsed() throws InterruptedException {
final Account account = AccountsHelper.createAccount(accountsManager, "+18005551111");

View File

@ -25,6 +25,7 @@ import com.fasterxml.jackson.core.JsonProcessingException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
@ -39,7 +40,9 @@ import java.util.concurrent.CompletionException;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@ -120,6 +123,7 @@ class AccountsTest {
when(mockDynamicConfigManager.getConfiguration())
.thenReturn(new DynamicConfiguration());
clock.pin(Instant.EPOCH);
accounts = new Accounts(
clock,
DYNAMO_DB_EXTENSION.getDynamoDbClient(),
@ -944,8 +948,14 @@ class AccountsTest {
accounts.confirmUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_2).join();
final UUID newHandle = account.getUsernameLinkHandle();
// switching usernames should put a hold on our original username
assertThat(accounts.getByUsernameHash(USERNAME_HASH_1).join()).isEmpty();
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).isEmpty();
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).containsExactlyInAnyOrderEntriesOf(Map.of(
Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.b(USERNAME_HASH_1),
Accounts.UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.b(account.getUuid()),
Accounts.UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(false),
Accounts.UsernameTable.ATTR_TTL,
AttributeValues.n(clock.instant().plus(Accounts.USERNAME_HOLD_DURATION).getEpochSecond())));
assertThat(accounts.getByUsernameLinkHandle(oldHandle).join()).isEmpty();
{
@ -1331,6 +1341,217 @@ class AccountsTest {
assertThat(account.getUsernameHash()).isEmpty();
}
@ParameterizedTest
@ValueSource(booleans = {false, true})
void testRemoveOldestHold(boolean clearUsername) {
Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
final List<byte[]> usernames = IntStream.range(0, 7).mapToObj(i -> TestRandomUtil.nextBytes(32)).toList();
final ArrayDeque<byte[]> expectedHolds = new ArrayDeque<>();
expectedHolds.add(USERNAME_HASH_1);
for (byte[] username : usernames) {
accounts.reserveUsernameHash(account, username, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, username, ENCRYPTED_USERNAME_1).join();
assertThat(accounts.getByUsernameHash(username).join()).isPresent();
final Account read = accounts.getByAccountIdentifier(account.getUuid()).orElseThrow();
assertThat(read.getUsernameHolds().stream().map(Account.UsernameHold::usernameHash).toList())
.containsExactlyElementsOf(expectedHolds);
expectedHolds.add(username);
if (expectedHolds.size() == Accounts.MAX_USERNAME_HOLDS + 1) {
expectedHolds.pop();
}
// clearing the username adds a hold, but the subsequent confirm in the next iteration should add the same hold
// (should be a noop) so we don't need to touch expectedHolds
if (clearUsername) {
accounts.clearUsernameHash(account).join();
}
}
final Account account2 = generateAccount("+18005554321", UUID.randomUUID(), UUID.randomUUID());
createAccount(account2);
// someone else should be able to get any of the usernames except the held usernames (MAX_HOLDS) +1 for the username
// currently held by the other account if we didn't clear it
final int numFree = usernames.size() - Accounts.MAX_USERNAME_HOLDS - (clearUsername ? 0 : 1);
final List<byte[]> freeUsernames = usernames.subList(0, numFree);
final List<byte[]> heldUsernames = usernames.subList(numFree, usernames.size());
for (byte[] username : freeUsernames) {
assertDoesNotThrow(() ->
accounts.reserveUsernameHash(account2, username, Duration.ofDays(2)).join());
}
for (byte[] username : heldUsernames) {
CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class,
accounts.reserveUsernameHash(account2, username, Duration.ofDays(2)));
}
}
@Test
void testHoldUsername() {
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
accounts.clearUsernameHash(account);
Account account2 = generateAccount("+18005554321", UUID.randomUUID(), UUID.randomUUID());
createAccount(account2);
CompletableFutureTestUtil.assertFailsWithCause(
UsernameHashNotAvailableException.class,
accounts.reserveUsernameHash(account2, USERNAME_HASH_1, Duration.ofDays(1)),
"account2 should not be able reserve username held by account");
// but we should be able to get it back
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
}
@Test
void testNoHoldsBarred() {
// should be able to reserve all MAX_HOLDS usernames
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
final List<byte[]> usernames = IntStream.range(0, Accounts.MAX_USERNAME_HOLDS + 1)
.mapToObj(i -> TestRandomUtil.nextBytes(32))
.toList();
for (byte[] username : usernames) {
accounts.reserveUsernameHash(account, username, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, username, ENCRYPTED_USERNAME_1).join();
}
// someone else shouldn't be able to get any of our holds
Account account2 = generateAccount("+18005554321", UUID.randomUUID(), UUID.randomUUID());
createAccount(account2);
for (byte[] username : usernames) {
CompletableFutureTestUtil.assertFailsWithCause(
UsernameHashNotAvailableException.class,
accounts.reserveUsernameHash(account2, username, Duration.ofDays(1)),
"account2 should not be able reserve username held by account");
}
// once the hold expires it's fine though
clock.pin(Instant.EPOCH.plus(Accounts.USERNAME_HOLD_DURATION).plus(Duration.ofSeconds(1)));
accounts.reserveUsernameHash(account2, usernames.get(0), Duration.ofDays(1)).join();
// if account1 modifies their username, we should also clear out the old holds, leaving only their newly added hold
accounts.clearUsernameHash(account).join();
assertThat(account.getUsernameHolds().stream().map(Account.UsernameHold::usernameHash))
.containsExactly(usernames.getLast());
}
@Test
public void testCannotRemoveHold() {
// Tests the case where we are trying to remove a hold we think we have, but it turns out we've already lost it.
// This means that the Account record an account has a hold on a particular username, but that hold is held by
// someone else in the username table. This can happen when the hold TTL expires while we are performing the update
// operation that attempts to remove the hold, and another user swoops in and takes the held username. In this
// case, a simple retry should let us check the clock again and notice that our hold in our account has expired.
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
accounts.reserveUsernameHash(account, USERNAME_HASH_2, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_1).join();
// Now we have a hold on username_hash_1. Simulate a race where the TTL on username_hash_1 expires, and someone
// else picks up the username by going forward and then back in time
Account account2 = generateAccount("+18005554321", UUID.randomUUID(), UUID.randomUUID());
createAccount(account2);
clock.pin(Instant.EPOCH.plus(Accounts.USERNAME_HOLD_DURATION).plus(Duration.ofSeconds(1)));
accounts.reserveUsernameHash(account2, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account2, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
clock.pin(Instant.EPOCH);
// already have 1 hold, should be able to get to MAX_HOLDS without a problem
for (int i = 1; i < Accounts.MAX_USERNAME_HOLDS; i++) {
accounts.reserveUsernameHash(account, TestRandomUtil.nextBytes(32), Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, TestRandomUtil.nextBytes(32), ENCRYPTED_USERNAME_1).join();
}
accounts.reserveUsernameHash(account, TestRandomUtil.nextBytes(32), Duration.ofDays(1)).join();
// Should fail, because we cannot remove our hold on USERNAME_HASH_1
CompletableFutureTestUtil.assertFailsWithCause(ContestedOptimisticLockException.class,
accounts.confirmUsernameHash(account, TestRandomUtil.nextBytes(32), ENCRYPTED_USERNAME_1));
// Should now pass once we realize our hold's TTL is over
clock.pin(Instant.EPOCH.plus(Accounts.USERNAME_HOLD_DURATION).plus(Duration.ofSeconds(1)));
accounts.confirmUsernameHash(account, TestRandomUtil.nextBytes(32), ENCRYPTED_USERNAME_1).join();
}
@Test
void testDeduplicateHoldsOnSwappedUsernames() {
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
final Consumer<byte[]> assertSingleHold = (byte[] usernameToCheck) -> {
// our account should have exactly one hold for the username
assertThat(account.getUsernameHolds().stream().map(Account.UsernameHold::usernameHash).toList())
.containsExactly(usernameToCheck);
// the username should be reserved for USERNAME_HOLD_DURATION (a re-reservation shouldn't reduce our expiration to
// the provided reservation TTL)
assertThat(
AttributeValues.getLong(getUsernameConstraintTableItem(usernameToCheck), Accounts.UsernameTable.ATTR_TTL, 0L))
.isEqualTo(Accounts.USERNAME_HOLD_DURATION.getSeconds());
};
// Swap back and forth between username 1 and 2. Username hashes shouldn't reappear in our holds if we already have
// a hold
for (int i = 0; i < 5; i++) {
accounts.reserveUsernameHash(account, USERNAME_HASH_2, Duration.ofSeconds(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_1).join();
assertSingleHold.accept(USERNAME_HASH_1);
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofSeconds(1)).join();
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
assertSingleHold.accept(USERNAME_HASH_2);
}
}
@Test
void testRemoveHoldAfterConfirm() {
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
createAccount(account);
final List<byte[]> usernames = IntStream.range(0, Accounts.MAX_USERNAME_HOLDS)
.mapToObj(i -> TestRandomUtil.nextBytes(32)).toList();
for (byte[] username : usernames) {
accounts.reserveUsernameHash(account, username, Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, username, ENCRYPTED_USERNAME_1).join();
}
int holdToRereserve = (Accounts.MAX_USERNAME_HOLDS / 2) - 1;
// should have MAX_HOLDS - 1 holds (everything in usernames except the last username, which is our current)
assertThat(account.getUsernameHolds().stream().map(Account.UsernameHold::usernameHash).toList())
.containsExactlyElementsOf(usernames.subList(0, usernames.size() - 1));
// if we confirm a username we already have held, it should just drop out of the holds list
accounts.reserveUsernameHash(account, usernames.get(holdToRereserve), Duration.ofDays(1)).join();
accounts.confirmUsernameHash(account, usernames.get(holdToRereserve), ENCRYPTED_USERNAME_1).join();
// should have a hold on every username but the one we just confirmed
assertThat(account.getUsernameHolds().stream().map(Account.UsernameHold::usernameHash).toList())
.containsExactlyElementsOf(Stream.concat(
usernames.subList(0, holdToRereserve).stream(),
usernames.subList(holdToRereserve + 1, usernames.size()).stream())
.toList());
}
@Test
public void testIgnoredFieldsNotAddedToDataAttribute() throws Exception {
final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
@ -1405,6 +1626,8 @@ class AccountsTest {
assertInstanceOf(DeviceIdDeserializer.DeviceIdDeserializationException.class, cause);
}
private static Device generateDevice(byte id) {
return DevicesHelper.createDevice(id);
}

View File

@ -0,0 +1,135 @@
/*
* Copyright 2024 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.workers;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.argThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Consumer;
import java.util.stream.IntStream;
import net.sourceforge.argparse4j.inf.Namespace;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.ArgumentCaptor;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import reactor.core.publisher.Flux;
class RemoveExpiredUsernameHoldsCommandTest {
private static class TestRemoveExpiredUsernameHoldsCommand extends RemoveExpiredUsernameHoldsCommand {
private final CommandDependencies commandDependencies;
private final Namespace namespace;
public TestRemoveExpiredUsernameHoldsCommand(final Clock clock, final AccountsManager accountsManager,
final boolean isDryRun) {
super(clock);
commandDependencies = mock(CommandDependencies.class);
when(commandDependencies.accountsManager()).thenReturn(accountsManager);
namespace = mock(Namespace.class);
when(namespace.getBoolean(RemoveExpiredUsernameHoldsCommand.DRY_RUN_ARGUMENT)).thenReturn(isDryRun);
when(namespace.getInt(RemoveExpiredUsernameHoldsCommand.MAX_CONCURRENCY_ARGUMENT)).thenReturn(16);
}
@Override
protected CommandDependencies getCommandDependencies() {
return commandDependencies;
}
@Override
protected Namespace getNamespace() {
return namespace;
}
}
@ParameterizedTest
@ValueSource(booleans = {true, false})
void crawlAccounts(final boolean isDryRun) {
final TestClock clock = TestClock.pinned(Instant.EPOCH.plus(Duration.ofSeconds(1)));
final AccountsManager accountsManager = mock(AccountsManager.class);
when(accountsManager.updateAsync(any(), any()))
.thenReturn(CompletableFuture.completedFuture(null));
final RemoveExpiredUsernameHoldsCommand removeExpiredUsernameHoldsCommand =
new TestRemoveExpiredUsernameHoldsCommand(clock, accountsManager, isDryRun);
final Account hasHolds = mock(Account.class);
final List<Account.UsernameHold> originalHolds = List.of(
// expired
new Account.UsernameHold(TestRandomUtil.nextBytes(32), Instant.EPOCH.getEpochSecond()),
// not expired
new Account.UsernameHold(TestRandomUtil.nextBytes(32),
Instant.EPOCH.plus(Duration.ofSeconds(5)).getEpochSecond()));
when(hasHolds.getUsernameHolds()).thenReturn(originalHolds);
final Account noHolds = mock(Account.class);
removeExpiredUsernameHoldsCommand.crawlAccounts(Flux.just(hasHolds, noHolds));
if (isDryRun) {
verifyNoInteractions(accountsManager);
} else {
ArgumentCaptor<Consumer<Account>> updaterCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(accountsManager, times(1)).updateAsync(eq(hasHolds), updaterCaptor.capture());
final Consumer<Account> consumer = updaterCaptor.getValue();
consumer.accept(hasHolds);
verify(hasHolds, times(1)).setUsernameHolds(argThat(holds ->
holds.equals(List.of(originalHolds.getLast()))));
verifyNoMoreInteractions(accountsManager);
}
}
@Test
public void removeHolds() {
final List<Account.UsernameHold> holds = IntStream.range(0, 100)
.mapToObj(i -> new Account.UsernameHold(TestRandomUtil.nextBytes(32), i)).toList();
final List<Account.UsernameHold> shuffled = new ArrayList<>(holds);
Collections.shuffle(shuffled);
final int currentTime = ThreadLocalRandom.current().nextInt(0, 100);
final Clock clock = TestClock.pinned(Instant.EPOCH.plus(Duration.ofSeconds(currentTime)));
final RemoveExpiredUsernameHoldsCommand removeExpiredUsernameHoldsCommand =
new TestRemoveExpiredUsernameHoldsCommand(clock, mock(AccountsManager.class), false);
final List<Account.UsernameHold> actual = new ArrayList<>(shuffled);
final int numRemoved = removeExpiredUsernameHoldsCommand.removeExpired(actual);
assertThat(numRemoved).isEqualTo(currentTime);
assertThat(actual).hasSize(100 - currentTime);
// should preserve order
final Iterator<Account.UsernameHold> expected = shuffled.iterator();
for (Account.UsernameHold hold : actual) {
while (!Arrays.equals(expected.next().usernameHash(), hold.usernameHash())) {
assertThat(expected).as("expected should be in order").hasNext();
}
}
}
}