MirrorMirror
/
Signal-Server
mirror of https://github.com/signalapp/Signal-Server.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Simplify username operations in `Accounts` 

- Group username table constants together
- Rethrow JsonProcessingException earlier
- Use UpdateAccountSpec.forAccount in username operations
- Inline confirm/clear transaction helpers
This commit is contained in:
Ravi Khadiwala 2024-02-28 22:23:20 -06:00 committed by ravi-signal
parent 8f100a792e
commit 47b24b5dff
3 changed files with 223 additions and 305 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

484
service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
View File

@ -119,14 +119,9 @@ public class Accounts extends AbstractDynamoDbStore {
static final String ATTR_CANONICALLY_DISCOVERABLE = "C"; static final String ATTR_CANONICALLY_DISCOVERABLE = "C";
// username hash; byte[] or null // username hash; byte[] or null
static final String ATTR_USERNAME_HASH = "N"; static final String ATTR_USERNAME_HASH = "N";
// confirmed; bool
static final String ATTR_CONFIRMED = "F";
// reclaimable; bool. Indicates that on confirmation the username link should be preserved
static final String ATTR_RECLAIMABLE = "R";
// unidentified access key; byte[] or null // unidentified access key; byte[] or null
static final String ATTR_UAK = "UAK"; static final String ATTR_UAK = "UAK";
// time to live; number
static final String ATTR_TTL = "TTL";
static final String DELETED_ACCOUNTS_KEY_ACCOUNT_E164 = "P"; static final String DELETED_ACCOUNTS_KEY_ACCOUNT_E164 = "P";
static final String DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID = "U"; static final String DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID = "U";
@ -182,6 +177,19 @@ public class Accounts extends AbstractDynamoDbStore {
deletedAccountsTableName); deletedAccountsTableName);
} }
static class UsernameTable {
// usernameHash; bytes.
static final String KEY_USERNAME_HASH = Accounts.ATTR_USERNAME_HASH;
// uuid, bytes. The owner of the username or reservation
static final String ATTR_ACCOUNT_UUID = Accounts.KEY_ACCOUNT_UUID;
// confirmed; bool
static final String ATTR_CONFIRMED = "F";
// reclaimable; bool. Indicates that on confirmation the username link should be preserved
static final String ATTR_RECLAIMABLE = "R";
// time to live; number
static final String ATTR_TTL = "TTL";
}
boolean create(final Account account, final List<TransactWriteItem> additionalWriteItems) boolean create(final Account account, final List<TransactWriteItem> additionalWriteItems)
throws AccountAlreadyExistsException { throws AccountAlreadyExistsException {
@ -255,8 +263,6 @@ public class Accounts extends AbstractDynamoDbStore {
// this shouldn't happen // this shouldn't happen
throw new RuntimeException("could not create account: " + extractCancellationReasonCodes(e)); throw new RuntimeException("could not create account: " + extractCancellationReasonCodes(e));
} }
} catch (final JsonProcessingException e) {
throw new IllegalArgumentException(e);
} finally { } finally {
sample.stop(CREATE_TIMER); sample.stop(CREATE_TIMER);
} }
@ -313,16 +319,16 @@ public class Accounts extends AbstractDynamoDbStore {
.put(Put.builder() .put(Put.builder()
.tableName(usernamesConstraintTableName) .tableName(usernamesConstraintTableName)
.item(Map.of( .item(Map.of(
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountToCreate.getUuid()), UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash),
ATTR_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash), UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(accountToCreate.getUuid()),
ATTR_TTL, AttributeValues.fromLong(expirationTime), UsernameTable.ATTR_TTL, AttributeValues.fromLong(expirationTime),
ATTR_CONFIRMED, AttributeValues.fromBool(false), UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(false),
ATTR_RECLAIMABLE, AttributeValues.fromBool(true))) UsernameTable.ATTR_RECLAIMABLE, AttributeValues.fromBool(true)))
.conditionExpression("attribute_not_exists(#username_hash) OR (#ttl < :now) OR #uuid = :uuid") .conditionExpression("attribute_not_exists(#username_hash) OR (#ttl < :now) OR #uuid = :uuid")
.expressionAttributeNames(Map.of( .expressionAttributeNames(Map.of(
"#username_hash", ATTR_USERNAME_HASH, "#username_hash", UsernameTable.KEY_USERNAME_HASH,
"#ttl", ATTR_TTL, "#ttl", UsernameTable.ATTR_TTL,
"#uuid", KEY_ACCOUNT_UUID)) "#uuid", UsernameTable.ATTR_ACCOUNT_UUID))
.expressionAttributeValues(Map.of( .expressionAttributeValues(Map.of(
":now", AttributeValues.fromLong(clock.instant().getEpochSecond()), ":now", AttributeValues.fromLong(clock.instant().getEpochSecond()),
":uuid", AttributeValues.fromUUID(accountToCreate.getUuid()))) ":uuid", AttributeValues.fromUUID(accountToCreate.getUuid())))
@ -431,8 +437,6 @@ public class Accounts extends AbstractDynamoDbStore {
account.setVersion(account.getVersion() + 1); account.setVersion(account.getVersion() + 1);
succeeded = true; succeeded = true;
} catch (final JsonProcessingException e) {
throw new IllegalArgumentException(e);
} catch (final TransactionCanceledException e) { } catch (final TransactionCanceledException e) {
if (e.hasCancellationReasons()) { if (e.hasCancellationReasons()) {
if (CONDITIONAL_CHECK_FAILED.equals(e.cancellationReasons().get(accountUpdateIndex).code())) { if (CONDITIONAL_CHECK_FAILED.equals(e.cancellationReasons().get(accountUpdateIndex).code())) {
@ -474,13 +478,6 @@ public class Accounts extends AbstractDynamoDbStore {
// Use account UUID as a "reservation token" - by providing this, the client proves ownership of the hash // 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 = account.getUuid();
final byte[] accountJsonBytes;
try {
accountJsonBytes = ACCOUNT_DDB_JSON_WRITER.writeValueAsBytes(account);
} catch (final JsonProcessingException e) {
throw new IllegalArgumentException(e);
}
final List<TransactWriteItem> writeItems = new ArrayList<>(); final List<TransactWriteItem> writeItems = new ArrayList<>();
@ -488,13 +485,17 @@ public class Accounts extends AbstractDynamoDbStore {
.put(Put.builder() .put(Put.builder()
.tableName(usernamesConstraintTableName) .tableName(usernamesConstraintTableName)
.item(Map.of( .item(Map.of(
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid), UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
ATTR_USERNAME_HASH, AttributeValues.fromByteArray(reservedUsernameHash), UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(reservedUsernameHash),
ATTR_TTL, AttributeValues.fromLong(expirationTime), UsernameTable.ATTR_TTL, AttributeValues.fromLong(expirationTime),
ATTR_CONFIRMED, AttributeValues.fromBool(false), UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(false),
ATTR_RECLAIMABLE, AttributeValues.fromBool(false))) UsernameTable.ATTR_RECLAIMABLE, AttributeValues.fromBool(false)))
.conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :confirmed)") .conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :confirmed)")
.expressionAttributeNames(Map.of("#username_hash", ATTR_USERNAME_HASH, "#ttl", ATTR_TTL, "#aci", KEY_ACCOUNT_UUID, "#confirmed", ATTR_CONFIRMED)) .expressionAttributeNames(Map.of(
"#username_hash", UsernameTable.KEY_USERNAME_HASH,
"#ttl", UsernameTable.ATTR_TTL,
"#aci", UsernameTable.ATTR_ACCOUNT_UUID,
"#confirmed", UsernameTable.ATTR_CONFIRMED))
.expressionAttributeValues(Map.of( .expressionAttributeValues(Map.of(
":now", AttributeValues.fromLong(clock.instant().getEpochSecond()), ":now", AttributeValues.fromLong(clock.instant().getEpochSecond()),
":aci", AttributeValues.fromUUID(uuid), ":aci", AttributeValues.fromUUID(uuid),
@ -503,39 +504,22 @@ public class Accounts extends AbstractDynamoDbStore {
.build()) .build())
.build()); .build());
writeItems.add( writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, account).transactItem());
TransactWriteItem.builder()
.update(Update.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.updateExpression("SET #data = :data ADD #version :version_increment")
.conditionExpression("#version = :version")
.expressionAttributeNames(Map.of("#data", ATTR_ACCOUNT_DATA, "#version", ATTR_VERSION))
.expressionAttributeValues(Map.of(
":data", AttributeValues.fromByteArray(accountJsonBytes),
":version", AttributeValues.fromInt(account.getVersion()),
":version_increment", AttributeValues.fromInt(1)))
.build())
.build());
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder() return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
.transactItems(writeItems) .exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> {
.build()) // If the constraint table update failed the condition check, the username's taken and we should stop
.exceptionally(throwable -> { // trying. However, if the accounts table fails the conditional check or
if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { // either table was concurrently updated, it's an optimistic locking failure and we should try again.
// If the constraint table update failed the condition check, the username's taken and we should stop if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
// trying. However, if the accounts table fails the conditional check or throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
// either table was concurrently updated, it's an optimistic locking failure and we should try again. } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) ||
if (conditionalCheckFailed(e.cancellationReasons().get(0))) { e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
throw ExceptionUtils.wrap(new UsernameHashNotAvailableException()); throw new ContestedOptimisticLockException();
} else if (conditionalCheckFailed(e.cancellationReasons().get(1)) || } else {
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { throw ExceptionUtils.wrap(e);
throw new ContestedOptimisticLockException();
}
} }
}))
throw ExceptionUtils.wrap(throwable);
})
.whenComplete((response, throwable) -> { .whenComplete((response, throwable) -> {
sample.stop(RESERVE_USERNAME_TIMER); sample.stop(RESERVE_USERNAME_TIMER);
@ -561,22 +545,52 @@ public class Accounts extends AbstractDynamoDbStore {
*/ */
public CompletableFuture<Void> confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername) { public CompletableFuture<Void> confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername) {
final Timer.Sample sample = Timer.start(); final Timer.Sample sample = Timer.start();
if (usernameHash == null) {
throw new IllegalArgumentException("Cannot confirm a null usernameHash");
}
return pickLinkHandle(account, usernameHash) return pickLinkHandle(account, usernameHash)
.thenCompose(linkHandle -> { .thenCompose(linkHandle -> {
final TransactWriteItemsRequest request; final Account updatedAccount = AccountUtil.cloneAccountAsNotStale(account);
try { updatedAccount.setUsernameHash(usernameHash);
final Account updatedAccount = AccountUtil.cloneAccountAsNotStale(account); updatedAccount.setReservedUsernameHash(null);
updatedAccount.setUsernameHash(usernameHash); updatedAccount.setUsernameLinkDetails(encryptedUsername == null ? null : linkHandle, encryptedUsername);
updatedAccount.setReservedUsernameHash(null);
updatedAccount.setUsernameLinkDetails(encryptedUsername == null ? null : linkHandle, encryptedUsername);
request = buildConfirmUsernameHashRequest(updatedAccount, account.getUsernameHash()); final Optional<byte[]> maybeOriginalUsernameHash = account.getUsernameHash();
} catch (final JsonProcessingException e) {
throw new IllegalArgumentException(e);
}
return asyncClient.transactWriteItems(request).thenApply(ignored -> linkHandle); final List<TransactWriteItem> writeItems = new ArrayList<>();
// 0: 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()
.tableName(usernamesConstraintTableName)
.item(Map.of(
UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(updatedAccount.getUuid()),
UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash),
UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(true)))
// it's not in the constraint table OR it's expired OR it was reserved by us
.conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :confirmed)")
.expressionAttributeNames(Map.of(
"#username_hash", UsernameTable.KEY_USERNAME_HASH,
"#ttl", UsernameTable.ATTR_TTL,
"#aci", UsernameTable.ATTR_ACCOUNT_UUID,
"#confirmed", UsernameTable.ATTR_CONFIRMED))
.expressionAttributeValues(Map.of(
":now", AttributeValues.fromLong(clock.instant().getEpochSecond()),
":aci", AttributeValues.fromUUID(updatedAccount.getUuid()),
":confirmed", AttributeValues.fromBool(false)))
.returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
.build())
.build());
// 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)));
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build())
.thenApply(ignored -> linkHandle);
}) })
.thenApply(linkHandle -> { .thenApply(linkHandle -> {
account.setUsernameHash(usernameHash); account.setUsernameHash(usernameHash);
@ -586,23 +600,19 @@ public class Accounts extends AbstractDynamoDbStore {
account.setVersion(account.getVersion() + 1); account.setVersion(account.getVersion() + 1);
return (Void) null; return (Void) null;
}) })
.exceptionally(throwable -> { .exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> {
if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { // If the constraint table update failed the condition check, the username's taken and we should stop
// 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
// trying. However, if the accounts table fails the conditional check or // updated, it's an optimistic locking failure and we should try again.
// either table was concurrently updated, it's an optimistic locking failure and we should try again. if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
// NOTE: the fixed indices here must be kept in sync with the creation of the TransactWriteItems in throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
// buildConfirmUsernameHashRequest! } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) ||
if (conditionalCheckFailed(e.cancellationReasons().get(0))) { e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
throw ExceptionUtils.wrap(new UsernameHashNotAvailableException()); throw new ContestedOptimisticLockException();
} else if (conditionalCheckFailed(e.cancellationReasons().get(1)) || } else {
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { throw ExceptionUtils.wrap(e);
throw new ContestedOptimisticLockException();
}
} }
}))
throw ExceptionUtils.wrap(throwable);
})
.whenComplete((ignored, throwable) -> sample.stop(SET_USERNAME_TIMER)); .whenComplete((ignored, throwable) -> sample.stop(SET_USERNAME_TIMER));
} }
@ -616,10 +626,10 @@ public class Accounts extends AbstractDynamoDbStore {
// preserve the link handle. // preserve the link handle.
return asyncClient.getItem(GetItemRequest.builder() return asyncClient.getItem(GetItemRequest.builder()
.tableName(usernamesConstraintTableName) .tableName(usernamesConstraintTableName)
.key(Map.of(ATTR_USERNAME_HASH, AttributeValues.b(usernameHash))) .key(Map.of(UsernameTable.KEY_USERNAME_HASH, AttributeValues.b(usernameHash)))
.projectionExpression(ATTR_RECLAIMABLE).build()) .projectionExpression(UsernameTable.ATTR_RECLAIMABLE).build())
.thenApply(response -> { .thenApply(response -> {
if (response.hasItem() && AttributeValues.getBool(response.item(), ATTR_RECLAIMABLE, false)) { if (response.hasItem() && AttributeValues.getBool(response.item(), UsernameTable.ATTR_RECLAIMABLE, false)) {
// this username reservation indicates it's a username waiting to be "reclaimed" // this username reservation indicates it's a username waiting to be "reclaimed"
return account.getUsernameLinkHandle(); return account.getUsernameLinkHandle();
} }
@ -629,72 +639,6 @@ public class Accounts extends AbstractDynamoDbStore {
}); });
} }
private TransactWriteItemsRequest buildConfirmUsernameHashRequest(final Account updatedAccount,
final Optional<byte[]> maybeOriginalUsernameHash)
throws JsonProcessingException {
final List<TransactWriteItem> writeItems = new ArrayList<>();
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()
.tableName(usernamesConstraintTableName)
.item(Map.of(
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(updatedAccount.getUuid()),
ATTR_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash),
ATTR_CONFIRMED, AttributeValues.fromBool(true)))
// it's not in the constraint table OR it's expired OR it was reserved by us
.conditionExpression("attribute_not_exists(#username_hash) OR #ttl < :now OR (#aci = :aci AND #confirmed = :confirmed)")
.expressionAttributeNames(Map.of("#username_hash", ATTR_USERNAME_HASH, "#ttl", ATTR_TTL, "#aci", KEY_ACCOUNT_UUID, "#confirmed", ATTR_CONFIRMED))
.expressionAttributeValues(Map.of(
":now", AttributeValues.fromLong(clock.instant().getEpochSecond()),
":aci", AttributeValues.fromUUID(updatedAccount.getUuid()),
":confirmed", AttributeValues.fromBool(false)))
.returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
.build())
.build());
final StringBuilder updateExpr = new StringBuilder("SET #data = :data, #username_hash = :username_hash");
final Map<String, AttributeValue> expressionAttributeValues = new HashMap<>(Map.of(
":data", accountDataAttributeValue(updatedAccount),
":username_hash", AttributeValues.fromByteArray(usernameHash),
":version", AttributeValues.fromInt(updatedAccount.getVersion()),
":version_increment", AttributeValues.fromInt(1)));
if (updatedAccount.getUsernameLinkHandle() != null) {
updateExpr.append(", #ul = :ul");
expressionAttributeValues.put(":ul", AttributeValues.fromUUID(updatedAccount.getUsernameLinkHandle()));
} else {
updateExpr.append(" REMOVE #ul");
}
updateExpr.append(" ADD #version :version_increment");
writeItems.add(
TransactWriteItem.builder()
.update(Update.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(updatedAccount.getUuid())))
.updateExpression(updateExpr.toString())
.conditionExpression("#version = :version")
.expressionAttributeNames(Map.of("#data", ATTR_ACCOUNT_DATA,
"#username_hash", ATTR_USERNAME_HASH,
"#ul", ATTR_USERNAME_LINK_UUID,
"#version", ATTR_VERSION))
.expressionAttributeValues(expressionAttributeValues)
.build())
.build());
maybeOriginalUsernameHash.ifPresent(originalUsernameHash -> writeItems.add(
buildDelete(usernamesConstraintTableName, ATTR_USERNAME_HASH, originalUsernameHash)));
return TransactWriteItemsRequest.builder()
.transactItems(writeItems)
.build();
}
/** /**
* Clear the username hash and link from the given account * Clear the username hash and link from the given account
* *
@ -704,73 +648,40 @@ public class Accounts extends AbstractDynamoDbStore {
* to the account or username constraint records. * to the account or username constraint records.
*/ */
public CompletableFuture<Void> clearUsernameHash(final Account account) { public CompletableFuture<Void> clearUsernameHash(final Account account) {
return account.getUsernameHash().map(usernameHash -> { if (account.getUsernameHash().isEmpty()) {
final Timer.Sample sample = Timer.start(); // no username to clear
return CompletableFuture.completedFuture(null);
}
final byte[] usernameHash = account.getUsernameHash().get();
final Timer.Sample sample = Timer.start();
final TransactWriteItemsRequest request; final Account updatedAccount = AccountUtil.cloneAccountAsNotStale(account);
updatedAccount.setUsernameHash(null);
updatedAccount.setUsernameLinkDetails(null, null);
try { return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
final Account updatedAccount = AccountUtil.cloneAccountAsNotStale(account); .transactItems(List.of(
updatedAccount.setUsernameHash(null); // 0: remove the username from the account object, conditioned on account version
updatedAccount.setUsernameLinkDetails(null, null); UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem(),
// 1: remote the username from the constraint table
request = buildClearUsernameHashRequest(updatedAccount, usernameHash); buildDelete(usernamesConstraintTableName, UsernameTable.KEY_USERNAME_HASH, usernameHash)))
} catch (final JsonProcessingException e) { .build())
throw new IllegalArgumentException(e); .thenAccept(ignored -> {
} account.setUsernameHash(null);
account.setUsernameLinkDetails(null, null);
return asyncClient.transactWriteItems(request) account.setVersion(account.getVersion() + 1);
.thenAccept(ignored -> { })
account.setUsernameHash(null); .exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> {
account.setUsernameLinkDetails(null, null); if (conditionalCheckFailed(e.cancellationReasons().get(0)) ||
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
account.setVersion(account.getVersion() + 1); throw new ContestedOptimisticLockException();
}) } else {
.exceptionally(throwable -> { throw ExceptionUtils.wrap(e);
if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) { }
if (conditionalCheckFailed(e.cancellationReasons().get(0)) || }))
e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) { .whenComplete((ignored, throwable) -> sample.stop(CLEAR_USERNAME_HASH_TIMER));
throw new ContestedOptimisticLockException();
}
}
throw ExceptionUtils.wrap(throwable);
})
.whenComplete((ignored, throwable) -> sample.stop(CLEAR_USERNAME_HASH_TIMER));
}).orElseGet(() -> CompletableFuture.completedFuture(null));
} }
private TransactWriteItemsRequest buildClearUsernameHashRequest(final Account updatedAccount, final byte[] originalUsernameHash)
throws JsonProcessingException {
final List<TransactWriteItem> writeItems = new ArrayList<>();
writeItems.add(
TransactWriteItem.builder()
.update(Update.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(updatedAccount.getUuid())))
.updateExpression("SET #data = :data REMOVE #username_hash, #username_link ADD #version :version_increment")
.conditionExpression("#version = :version")
.expressionAttributeNames(Map.of("#data", ATTR_ACCOUNT_DATA,
"#username_hash", ATTR_USERNAME_HASH,
"#username_link", ATTR_USERNAME_LINK_UUID,
"#version", ATTR_VERSION))
.expressionAttributeValues(Map.of(
":data", accountDataAttributeValue(updatedAccount),
":version", AttributeValues.fromInt(updatedAccount.getVersion()),
":version_increment", AttributeValues.fromInt(1)))
.build())
.build());
writeItems.add(buildDelete(usernamesConstraintTableName, ATTR_USERNAME_HASH, originalUsernameHash));
return TransactWriteItemsRequest.builder()
.transactItems(writeItems)
.build();
}
/** /**
* A ddb update that can be used as part of a transaction or single-item update statement. * A ddb update that can be used as part of a transaction or single-item update statement.
*/ */
@ -806,65 +717,68 @@ public class Accounts extends AbstractDynamoDbStore {
static UpdateAccountSpec forAccount( static UpdateAccountSpec forAccount(
final String accountTableName, final String accountTableName,
final Account account) { final Account account) {
try { // username, e164, and pni cannot be modified through this method
// username, e164, and pni cannot be modified through this method final Map<String, String> attrNames = new HashMap<>(Map.of(
final Map<String, String> attrNames = new HashMap<>(Map.of( "#number", ATTR_ACCOUNT_E164,
"#number", ATTR_ACCOUNT_E164, "#data", ATTR_ACCOUNT_DATA,
"#data", ATTR_ACCOUNT_DATA, "#cds", ATTR_CANONICALLY_DISCOVERABLE,
"#cds", ATTR_CANONICALLY_DISCOVERABLE, "#version", ATTR_VERSION));
"#version", ATTR_VERSION));
final Map<String, AttributeValue> attrValues = new HashMap<>(Map.of( final Map<String, AttributeValue> attrValues = new HashMap<>(Map.of(
":data", accountDataAttributeValue(account), ":data", accountDataAttributeValue(account),
":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()), ":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
":version", AttributeValues.fromInt(account.getVersion()), ":version", AttributeValues.fromInt(account.getVersion()),
":version_increment", AttributeValues.fromInt(1))); ":version_increment", AttributeValues.fromInt(1)));
final StringBuilder updateExpressionBuilder = new StringBuilder("SET #data = :data, #cds = :cds"); final StringBuilder updateExpressionBuilder = new StringBuilder("SET #data = :data, #cds = :cds");
if (account.getUnidentifiedAccessKey().isPresent()) { if (account.getUnidentifiedAccessKey().isPresent()) {
// if it's present in the account, also set the uak // if it's present in the account, also set the uak
attrNames.put("#uak", ATTR_UAK); attrNames.put("#uak", ATTR_UAK);
attrValues.put(":uak", AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get())); attrValues.put(":uak", AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get()));
updateExpressionBuilder.append(", #uak = :uak"); updateExpressionBuilder.append(", #uak = :uak");
}
// If the account has a username/handle pair, we should add it to the top level attributes.
// When we remove an encryptedUsername but preserve the link (re-registration), it's possible that the account
// has a usernameLinkHandle but not an encrypted username. In this case there should already be a top-level
// usernameLink attribute.
if (account.getEncryptedUsername().isPresent() && account.getUsernameLinkHandle() != null) {
attrNames.put("#ul", ATTR_USERNAME_LINK_UUID);
attrValues.put(":ul", AttributeValues.fromUUID(account.getUsernameLinkHandle()));
updateExpressionBuilder.append(", #ul = :ul");
}
// Some operations may remove the usernameLink or the usernameHash (re-registration, clear username link, and
// clear username hash). Since these also have top-level ddb attributes, we need to make sure to remove those
// as well.
final List<String> removes = new ArrayList<>();
if (account.getUsernameLinkHandle() == null) {
attrNames.put("#ul", ATTR_USERNAME_LINK_UUID);
removes.add("#ul");
}
if (account.getUsernameHash().isEmpty()) {
attrNames.put("#username_hash", ATTR_USERNAME_HASH);
removes.add("#username_hash");
}
if (!removes.isEmpty()) {
updateExpressionBuilder.append(" REMOVE %s".formatted(String.join(",", removes)));
}
updateExpressionBuilder.append(" ADD #version :version_increment");
return new UpdateAccountSpec(
accountTableName,
Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())),
attrNames,
attrValues,
updateExpressionBuilder.toString(),
"attribute_exists(#number) AND #version = :version");
} catch (final JsonProcessingException e) {
throw new IllegalArgumentException(e);
} }
if (account.getUsernameHash().isPresent()) {
// if it's present in the account, also set the username hash
attrNames.put("#usernameHash", ATTR_USERNAME_HASH);
attrValues.put(":usernameHash", AttributeValues.fromByteArray(account.getUsernameHash().get()));
updateExpressionBuilder.append(", #usernameHash = :usernameHash");
}
// If the account has a username/handle pair, we should add it to the top level attributes.
// When we remove an encryptedUsername but preserve the link (re-registration), it's possible that the account
// has a usernameLinkHandle but not an encrypted username. In this case there should already be a top-level
// usernameLink attribute.
if (account.getEncryptedUsername().isPresent() && account.getUsernameLinkHandle() != null) {
attrNames.put("#ul", ATTR_USERNAME_LINK_UUID);
attrValues.put(":ul", AttributeValues.fromUUID(account.getUsernameLinkHandle()));
updateExpressionBuilder.append(", #ul = :ul");
}
// Some operations may remove the usernameLink or the usernameHash (re-registration, clear username link, and
// clear username hash). Since these also have top-level ddb attributes, we need to make sure to remove those
// as well.
final List<String> removes = new ArrayList<>();
if (account.getUsernameLinkHandle() == null) {
attrNames.put("#ul", ATTR_USERNAME_LINK_UUID);
removes.add("#ul");
}
if (account.getUsernameHash().isEmpty()) {
attrNames.put("#username_hash", ATTR_USERNAME_HASH);
removes.add("#username_hash");
}
if (!removes.isEmpty()) {
updateExpressionBuilder.append(" REMOVE %s".formatted(String.join(",", removes)));
}
updateExpressionBuilder.append(" ADD #version :version_increment");
return new UpdateAccountSpec(
accountTableName,
Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())),
attrNames,
attrValues,
updateExpressionBuilder.toString(),
"attribute_exists(#number) AND #version = :version");
} }
} }
@ -980,9 +894,9 @@ public class Accounts extends AbstractDynamoDbStore {
public CompletableFuture<Optional<Account>> getByUsernameHash(final byte[] usernameHash) { public CompletableFuture<Optional<Account>> getByUsernameHash(final byte[] usernameHash) {
return getByIndirectLookupAsync(GET_BY_USERNAME_HASH_TIMER, return getByIndirectLookupAsync(GET_BY_USERNAME_HASH_TIMER,
usernamesConstraintTableName, usernamesConstraintTableName,
ATTR_USERNAME_HASH, UsernameTable.KEY_USERNAME_HASH,
AttributeValues.fromByteArray(usernameHash), AttributeValues.fromByteArray(usernameHash),
item -> AttributeValues.getBool(item, ATTR_CONFIRMED, false) // ignore items that are reservations (not confirmed) item -> AttributeValues.getBool(item, UsernameTable.ATTR_CONFIRMED, false) // ignore items that are reservations (not confirmed)
); );
} }
@ -1074,7 +988,7 @@ public class Accounts extends AbstractDynamoDbStore {
)); ));
account.getUsernameHash().ifPresent(usernameHash -> transactWriteItems.add( account.getUsernameHash().ifPresent(usernameHash -> transactWriteItems.add(
buildDelete(usernamesConstraintTableName, ATTR_USERNAME_HASH, usernameHash))); buildDelete(usernamesConstraintTableName, UsernameTable.KEY_USERNAME_HASH, usernameHash)));
transactWriteItems.addAll(additionalWriteItems); transactWriteItems.addAll(additionalWriteItems);
@ -1239,7 +1153,7 @@ public class Accounts extends AbstractDynamoDbStore {
final Account account, final Account account,
final AttributeValue uuidAttr, final AttributeValue uuidAttr,
final AttributeValue numberAttr, final AttributeValue numberAttr,
final AttributeValue pniUuidAttr) throws JsonProcessingException { final AttributeValue pniUuidAttr) {
final Map<String, AttributeValue> item = new HashMap<>(Map.of( final Map<String, AttributeValue> item = new HashMap<>(Map.of(
KEY_ACCOUNT_UUID, uuidAttr, KEY_ACCOUNT_UUID, uuidAttr,
@ -1378,8 +1292,12 @@ public class Accounts extends AbstractDynamoDbStore {
} }
} }
private static AttributeValue accountDataAttributeValue(final Account account) throws JsonProcessingException { private static AttributeValue accountDataAttributeValue(final Account account) {
return AttributeValues.fromByteArray(ACCOUNT_DDB_JSON_WRITER.writeValueAsBytes(account)); try {
return AttributeValues.fromByteArray(ACCOUNT_DDB_JSON_WRITER.writeValueAsBytes(account));
} catch (JsonProcessingException e) {
throw new IllegalArgumentException(e);
}
} }
private static boolean conditionalCheckFailed(final CancellationReason reason) { private static boolean conditionalCheckFailed(final CancellationReason reason) {

14
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java
View File

@ -163,11 +163,11 @@ class AccountsManagerUsernameIntegrationTest {
int i = 0; int i = 0;
for (byte[] hash : usernameHashes) { for (byte[] hash : usernameHashes) {
final Map<String, AttributeValue> item = new HashMap<>(Map.of( final Map<String, AttributeValue> item = new HashMap<>(Map.of(
Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(UUID.randomUUID()), Accounts.UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(UUID.randomUUID()),
Accounts.ATTR_USERNAME_HASH, AttributeValues.fromByteArray(hash))); Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(hash)));
// half of these are taken usernames, half are only reservations (have a TTL) // half of these are taken usernames, half are only reservations (have a TTL)
if (i % 2 == 0) { if (i % 2 == 0) {
item.put(Accounts.ATTR_TTL, item.put(Accounts.UsernameTable.ATTR_TTL,
AttributeValues.fromLong(Instant.now().plus(Duration.ofMinutes(1)).getEpochSecond())); AttributeValues.fromLong(Instant.now().plus(Duration.ofMinutes(1)).getEpochSecond()));
} }
i++; i++;
@ -192,8 +192,8 @@ class AccountsManagerUsernameIntegrationTest {
DYNAMO_DB_EXTENSION.getDynamoDbClient().putItem(PutItemRequest.builder() DYNAMO_DB_EXTENSION.getDynamoDbClient().putItem(PutItemRequest.builder()
.tableName(Tables.USERNAMES.tableName()) .tableName(Tables.USERNAMES.tableName())
.item(Map.of( .item(Map.of(
Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(UUID.randomUUID()), Accounts.UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(UUID.randomUUID()),
Accounts.ATTR_USERNAME_HASH, AttributeValues.fromByteArray(hash))) Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(hash)))
.build()); .build());
} }
@ -250,9 +250,9 @@ class AccountsManagerUsernameIntegrationTest {
// force expiration // force expiration
DYNAMO_DB_EXTENSION.getDynamoDbClient().updateItem(UpdateItemRequest.builder() DYNAMO_DB_EXTENSION.getDynamoDbClient().updateItem(UpdateItemRequest.builder()
.tableName(Tables.USERNAMES.tableName()) .tableName(Tables.USERNAMES.tableName())
.key(Map.of(Accounts.ATTR_USERNAME_HASH, AttributeValues.fromByteArray(USERNAME_HASH_1))) .key(Map.of(Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(USERNAME_HASH_1)))
.updateExpression("SET #ttl = :ttl") .updateExpression("SET #ttl = :ttl")
.expressionAttributeNames(Map.of("#ttl", Accounts.ATTR_TTL)) .expressionAttributeNames(Map.of("#ttl", Accounts.UsernameTable.ATTR_TTL))
.expressionAttributeValues(Map.of(":ttl", AttributeValues.fromLong(past))) .expressionAttributeValues(Map.of(":ttl", AttributeValues.fromLong(past)))
.build()); .build());

30
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
View File

@ -1212,8 +1212,8 @@ class AccountsTest {
final Map<String, AttributeValue> usernameConstraintRecord = getUsernameConstraintTableItem(USERNAME_HASH_1); final Map<String, AttributeValue> usernameConstraintRecord = getUsernameConstraintTableItem(USERNAME_HASH_1);
assertThat(usernameConstraintRecord).containsKey(Accounts.ATTR_USERNAME_HASH); assertThat(usernameConstraintRecord).containsKey(Accounts.UsernameTable.KEY_USERNAME_HASH);
assertThat(usernameConstraintRecord).doesNotContainKey(Accounts.ATTR_TTL); assertThat(usernameConstraintRecord).doesNotContainKey(Accounts.UsernameTable.ATTR_TTL);
} }
@Test @Test
@ -1231,10 +1231,10 @@ class AccountsTest {
final Map<String, AttributeValue> usernameConstraintRecord1 = getUsernameConstraintTableItem(USERNAME_HASH_1); final Map<String, AttributeValue> usernameConstraintRecord1 = getUsernameConstraintTableItem(USERNAME_HASH_1);
final Map<String, AttributeValue> usernameConstraintRecord2 = getUsernameConstraintTableItem(USERNAME_HASH_2); final Map<String, AttributeValue> usernameConstraintRecord2 = getUsernameConstraintTableItem(USERNAME_HASH_2);
assertThat(usernameConstraintRecord1).containsKey(Accounts.ATTR_USERNAME_HASH); assertThat(usernameConstraintRecord1).containsKey(Accounts.UsernameTable.KEY_USERNAME_HASH);
assertThat(usernameConstraintRecord2).containsKey(Accounts.ATTR_USERNAME_HASH); assertThat(usernameConstraintRecord2).containsKey(Accounts.UsernameTable.KEY_USERNAME_HASH);
assertThat(usernameConstraintRecord1).containsKey(Accounts.ATTR_TTL); assertThat(usernameConstraintRecord1).containsKey(Accounts.UsernameTable.ATTR_TTL);
assertThat(usernameConstraintRecord2).containsKey(Accounts.ATTR_TTL); assertThat(usernameConstraintRecord2).containsKey(Accounts.UsernameTable.ATTR_TTL);
clock.pin(Instant.EPOCH.plus(Duration.ofMinutes(1))); clock.pin(Instant.EPOCH.plus(Duration.ofMinutes(1)));
@ -1243,10 +1243,10 @@ class AccountsTest {
assertThat(account.getUsernameHash()).isEmpty(); assertThat(account.getUsernameHash()).isEmpty();
final Map<String, AttributeValue> newUsernameConstraintRecord1 = getUsernameConstraintTableItem(USERNAME_HASH_1); final Map<String, AttributeValue> newUsernameConstraintRecord1 = getUsernameConstraintTableItem(USERNAME_HASH_1);
assertThat(newUsernameConstraintRecord1).containsKey(Accounts.ATTR_USERNAME_HASH); assertThat(newUsernameConstraintRecord1).containsKey(Accounts.UsernameTable.KEY_USERNAME_HASH);
assertThat(newUsernameConstraintRecord1).containsKey(Accounts.ATTR_TTL); assertThat(newUsernameConstraintRecord1).containsKey(Accounts.UsernameTable.ATTR_TTL);
assertThat(usernameConstraintRecord1.get(Accounts.ATTR_TTL)) assertThat(usernameConstraintRecord1.get(Accounts.UsernameTable.ATTR_TTL))
.isNotEqualTo(newUsernameConstraintRecord1.get(Accounts.ATTR_TTL)); .isNotEqualTo(newUsernameConstraintRecord1.get(Accounts.UsernameTable.ATTR_TTL));
} }
@Test @Test
@ -1257,20 +1257,20 @@ class AccountsTest {
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join(); accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
assertArrayEquals(account.getReservedUsernameHash().orElseThrow(), USERNAME_HASH_1); assertArrayEquals(account.getReservedUsernameHash().orElseThrow(), USERNAME_HASH_1);
assertThat(account.getUsernameHash()).isEmpty(); assertThat(account.getUsernameHash()).isEmpty();
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).containsKey(Accounts.ATTR_TTL); assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).containsKey(Accounts.UsernameTable.ATTR_TTL);
accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join(); accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
assertThat(account.getReservedUsernameHash()).isEmpty(); assertThat(account.getReservedUsernameHash()).isEmpty();
assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1); assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1);
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).doesNotContainKey(Accounts.ATTR_TTL); assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).doesNotContainKey(Accounts.UsernameTable.ATTR_TTL);
CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class, CompletableFutureTestUtil.assertFailsWithCause(UsernameHashNotAvailableException.class,
accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1))); accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)));
assertThat(account.getReservedUsernameHash()).isEmpty(); assertThat(account.getReservedUsernameHash()).isEmpty();
assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1); assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1);
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).containsKey(Accounts.ATTR_USERNAME_HASH); assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).containsKey(Accounts.UsernameTable.KEY_USERNAME_HASH);
assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).doesNotContainKey(Accounts.ATTR_TTL); assertThat(getUsernameConstraintTableItem(USERNAME_HASH_1)).doesNotContainKey(Accounts.UsernameTable.ATTR_TTL);
} }
@Test @Test
@ -1492,7 +1492,7 @@ class AccountsTest {
return DYNAMO_DB_EXTENSION.getDynamoDbClient() return DYNAMO_DB_EXTENSION.getDynamoDbClient()
.getItem(GetItemRequest.builder() .getItem(GetItemRequest.builder()
.tableName(Tables.USERNAMES.tableName()) .tableName(Tables.USERNAMES.tableName())
.key(Map.of(Accounts.ATTR_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash))) .key(Map.of(Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash)))
.build()) .build())
.item(); .item();
} }