Write "recently deleted account" rows exclusively by PNI

This commit is contained in:
Jon Chambers 2024-11-27 10:59:01 -05:00 committed by Jon Chambers
parent f5d3d1e65d
commit d08bc4c413
2 changed files with 33 additions and 49 deletions

View File

@ -47,8 +47,6 @@ import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.Util; import org.whispersystems.textsecuregcm.util.Util;
import reactor.core.publisher.Flux; import reactor.core.publisher.Flux;
import reactor.core.scheduler.Scheduler; import reactor.core.scheduler.Scheduler;
import reactor.util.function.Tuple3;
import reactor.util.function.Tuples;
import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@ -59,7 +57,6 @@ import software.amazon.awssdk.services.dynamodb.model.Delete;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.Put; import software.amazon.awssdk.services.dynamodb.model.Put;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest; import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryResponse; import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure; import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
@ -245,11 +242,10 @@ public class Accounts extends AbstractDynamoDbStore {
// Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
// the newly-created account. // the newly-created account.
final TransactWriteItem deletedAccountDelete = buildRemoveDeletedAccount(account.getNumber()); final TransactWriteItem deletedAccountDelete = buildRemoveDeletedAccount(account.getPhoneNumberIdentifier());
final TransactWriteItem deletedAccountDeletePNI = buildRemoveDeletedAccount(account.getPhoneNumberIdentifier());
final Collection<TransactWriteItem> writeItems = new ArrayList<>( final Collection<TransactWriteItem> writeItems = new ArrayList<>(
List.of(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut, deletedAccountDeletePNI, deletedAccountDelete)); List.of(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut, deletedAccountDelete));
writeItems.addAll(additionalWriteItems); writeItems.addAll(additionalWriteItems);
@ -443,12 +439,9 @@ public class Accounts extends AbstractDynamoDbStore {
writeItems.add(buildConstraintTablePut(phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr)); writeItems.add(buildConstraintTablePut(phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr));
writeItems.add(buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, originalPni)); writeItems.add(buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, originalPni));
writeItems.add(buildConstraintTablePut(phoneNumberIdentifierConstraintTableName, uuidAttr, ATTR_PNI_UUID, pniAttr)); writeItems.add(buildConstraintTablePut(phoneNumberIdentifierConstraintTableName, uuidAttr, ATTR_PNI_UUID, pniAttr));
writeItems.add(buildRemoveDeletedAccount(number));
writeItems.add(buildRemoveDeletedAccount(phoneNumberIdentifier)); writeItems.add(buildRemoveDeletedAccount(phoneNumberIdentifier));
maybeDisplacedAccountIdentifier.ifPresent(displacedAccountIdentifier -> { maybeDisplacedAccountIdentifier.ifPresent(displacedAccountIdentifier ->
writeItems.add(buildPutDeletedAccount(displacedAccountIdentifier, originalNumber)); writeItems.add(buildPutDeletedAccount(displacedAccountIdentifier, originalPni)));
writeItems.add(buildPutDeletedAccount(displacedAccountIdentifier, originalPni));
});
// The `catch (TransactionCanceledException) block needs to check whether the cancellation reason is the account // The `catch (TransactionCanceledException) block needs to check whether the cancellation reason is the account
// update write item // update write item
@ -618,7 +611,7 @@ public class Accounts extends AbstractDynamoDbStore {
if (conditionalCheckFailed(e.cancellationReasons().get(0))) { if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
// The constraint table update failed the condition check. It could be because the username was taken, // 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 // or because we need to retry with a longer TTL
final Map<String, AttributeValue> item = e.cancellationReasons().get(0).item(); final Map<String, AttributeValue> item = e.cancellationReasons().getFirst().item();
final UUID existingOwner = AttributeValues.getUUID(item, UsernameTable.ATTR_ACCOUNT_UUID, null); final UUID existingOwner = AttributeValues.getUUID(item, UsernameTable.ATTR_ACCOUNT_UUID, null);
final boolean confirmed = AttributeValues.getBool(item, UsernameTable.ATTR_CONFIRMED, false); final boolean confirmed = AttributeValues.getBool(item, UsernameTable.ATTR_CONFIRMED, false);
final long existingTtl = AttributeValues.getLong(item, UsernameTable.ATTR_TTL, 0L); final long existingTtl = AttributeValues.getLong(item, UsernameTable.ATTR_TTL, 0L);
@ -1076,7 +1069,7 @@ public class Accounts extends AbstractDynamoDbStore {
final Throwable unwrapped = ExceptionUtils.unwrap(throwable); final Throwable unwrapped = ExceptionUtils.unwrap(throwable);
if (unwrapped instanceof TransactionCanceledException transactionCanceledException) { if (unwrapped instanceof TransactionCanceledException transactionCanceledException) {
if (CONDITIONAL_CHECK_FAILED.equals(transactionCanceledException.cancellationReasons().get(0).code())) { if (CONDITIONAL_CHECK_FAILED.equals(transactionCanceledException.cancellationReasons().getFirst().code())) {
throw new ContestedOptimisticLockException(); throw new ContestedOptimisticLockException();
} }
@ -1164,18 +1157,6 @@ public class Accounts extends AbstractDynamoDbStore {
.map(Accounts::fromItem))); .map(Accounts::fromItem)));
} }
private TransactWriteItem buildPutDeletedAccount(final UUID uuid, final String e164) {
return TransactWriteItem.builder()
.put(Put.builder()
.tableName(deletedAccountsTableName)
.item(Map.of(
DELETED_ACCOUNTS_KEY_ACCOUNT_PNI, AttributeValues.fromString(e164),
DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
DELETED_ACCOUNTS_ATTR_EXPIRES, AttributeValues.fromLong(clock.instant().plus(DELETED_ACCOUNTS_TIME_TO_LIVE).getEpochSecond())))
.build())
.build();
}
private TransactWriteItem buildPutDeletedAccount(final UUID aci, final UUID pni) { private TransactWriteItem buildPutDeletedAccount(final UUID aci, final UUID pni) {
return TransactWriteItem.builder() return TransactWriteItem.builder()
.put(Put.builder() .put(Put.builder()
@ -1188,15 +1169,6 @@ public class Accounts extends AbstractDynamoDbStore {
.build(); .build();
} }
private TransactWriteItem buildRemoveDeletedAccount(final String e164) {
return TransactWriteItem.builder()
.delete(Delete.builder()
.tableName(deletedAccountsTableName)
.key(Map.of(DELETED_ACCOUNTS_KEY_ACCOUNT_PNI, AttributeValues.fromString(e164)))
.build())
.build();
}
private TransactWriteItem buildRemoveDeletedAccount(final UUID pni) { private TransactWriteItem buildRemoveDeletedAccount(final UUID pni) {
return TransactWriteItem.builder() return TransactWriteItem.builder()
.delete(Delete.builder() .delete(Delete.builder()
@ -1253,7 +1225,6 @@ public class Accounts extends AbstractDynamoDbStore {
buildDelete(phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, account.getNumber()), buildDelete(phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, account.getNumber()),
buildDelete(accountsTableName, KEY_ACCOUNT_UUID, uuid), buildDelete(accountsTableName, KEY_ACCOUNT_UUID, uuid),
buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, account.getPhoneNumberIdentifier()), buildDelete(phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, account.getPhoneNumberIdentifier()),
buildPutDeletedAccount(uuid, account.getNumber()),
buildPutDeletedAccount(uuid, account.getPhoneNumberIdentifier()) buildPutDeletedAccount(uuid, account.getPhoneNumberIdentifier())
)); ));
@ -1290,7 +1261,7 @@ public class Accounts extends AbstractDynamoDbStore {
.sequential(); .sequential();
} }
public Flux<Tuple3<String, UUID, Long>> getE164KeyedDeletedAccounts(final int segments, final Scheduler scheduler) { public Flux<String> getE164sForRecentlyDeletedAccounts(final int segments, final Scheduler scheduler) {
if (segments < 1) { if (segments < 1) {
throw new IllegalArgumentException("Total number of segments must be positive"); throw new IllegalArgumentException("Total number of segments must be positive");
} }
@ -1303,14 +1274,12 @@ public class Accounts extends AbstractDynamoDbStore {
.consistentRead(true) .consistentRead(true)
.segment(segment) .segment(segment)
.totalSegments(segments) .totalSegments(segments)
.filterExpression("begins_with(#key, :e164Prefix)")
.expressionAttributeNames(Map.of("#key", DELETED_ACCOUNTS_KEY_ACCOUNT_PNI))
.expressionAttributeValues(Map.of(":e164Prefix", AttributeValue.fromS("+")))
.build()) .build())
.items()) .items())
.map(item -> .map(item -> item.get(DELETED_ACCOUNTS_KEY_ACCOUNT_PNI).s())
Tuples.of(
item.get(DELETED_ACCOUNTS_KEY_ACCOUNT_PNI).s(),
AttributeValues.getUUID(item, DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID, null),
AttributeValues.getLong(item, DELETED_ACCOUNTS_ATTR_EXPIRES, 0)))
.filter(item -> item.getT1().startsWith("+"))
.sequential(); .sequential();
} }
@ -1437,7 +1406,7 @@ public class Accounts extends AbstractDynamoDbStore {
.formatted(indexName, keyName, keyValue))); .formatted(indexName, keyName, keyValue)));
} }
final AttributeValue primaryKeyValue = response.items().get(0).get(KEY_ACCOUNT_UUID); final AttributeValue primaryKeyValue = response.items().getFirst().get(KEY_ACCOUNT_UUID);
return itemByKeyAsync(table, KEY_ACCOUNT_UUID, primaryKeyValue); return itemByKeyAsync(table, KEY_ACCOUNT_UUID, primaryKeyValue);
}); });
} }

View File

@ -22,6 +22,7 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.i18n.phonenumbers.PhoneNumberUtil;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.time.Duration; import java.time.Duration;
import java.time.Instant; import java.time.Instant;
@ -65,7 +66,6 @@ import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.TestClock; import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil; import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import reactor.core.scheduler.Schedulers; import reactor.core.scheduler.Schedulers;
import reactor.util.function.Tuples;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@ -75,7 +75,6 @@ import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.Put; import software.amazon.awssdk.services.dynamodb.model.Put;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure; import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
@ -1700,13 +1699,29 @@ class AccountsTest {
} }
@Test @Test
public void getE164KeyedDeletedAccounts() { public void getE164sForRecentlyDeletedAccounts() {
final Account deletedAccount = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID()); final UUID accountIdentifier = UUID.randomUUID();
final UUID phoneNumberIdentifier = UUID.randomUUID();
final String phoneNumber = PhoneNumberUtil.getInstance().format(
PhoneNumberUtil.getInstance().getExampleNumber("US"),
PhoneNumberUtil.PhoneNumberFormat.E164);
final Account deletedAccount = generateAccount(phoneNumber, accountIdentifier, phoneNumberIdentifier);
createAccount(deletedAccount); createAccount(deletedAccount);
accounts.delete(deletedAccount.getUuid(), List.of()).join(); accounts.delete(deletedAccount.getUuid(), List.of()).join();
// Artificially insert this row to simulate legacy data
DYNAMO_DB_EXTENSION.getDynamoDbClient().putItem(PutItemRequest.builder()
.tableName(Tables.DELETED_ACCOUNTS.tableName())
.item(Map.of(
Accounts.DELETED_ACCOUNTS_KEY_ACCOUNT_PNI, AttributeValues.fromString(phoneNumber),
Accounts.DELETED_ACCOUNTS_ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier),
Accounts.DELETED_ACCOUNTS_ATTR_EXPIRES, AttributeValues.fromLong(clock.instant().plus(Accounts.DELETED_ACCOUNTS_TIME_TO_LIVE).getEpochSecond())))
.build());
assertEquals( assertEquals(
List.of(Tuples.of("+18005551234", deletedAccount.getUuid(), clock.instant().plus(Accounts.DELETED_ACCOUNTS_TIME_TO_LIVE).toEpochMilli() / 1000)), List.of(phoneNumber),
accounts.getE164KeyedDeletedAccounts(1, Schedulers.immediate()).collectList().block()); accounts.getE164sForRecentlyDeletedAccounts(1, Schedulers.immediate()).collectList().block());
} }
private static Device generateDevice(byte id) { private static Device generateDevice(byte id) {