Identify accounts for which to delete keys by UUID.

This commit is contained in:
Jon Chambers 2021-07-26 09:57:20 -04:00 committed by Jon Chambers
parent d09dcc90fe
commit be20c04cd8
5 changed files with 15 additions and 18 deletions

View File

@ -102,7 +102,7 @@ public class DeviceController {
messages.clear(account.getUuid(), deviceId);
account = accounts.update(account, a -> a.removeDevice(deviceId));
keys.delete(account, deviceId);
keys.delete(account.getUuid(), deviceId);
// ensure any messages that came in after the first clear() are also removed
messages.clear(account.getUuid(), deviceId);
}

View File

@ -223,7 +223,7 @@ public class AccountsManager {
maybeExistingAccount.ifPresent(definitelyExistingAccount -> {
messagesManager.clear(definitelyExistingAccount.getUuid());
keysDynamoDb.delete(definitelyExistingAccount);
keysDynamoDb.delete(definitelyExistingAccount.getUuid());
});
pendingAccounts.remove(number);
@ -393,7 +393,7 @@ public class AccountsManager {
usernamesManager.delete(account.getUuid());
directoryQueue.deleteAccount(account);
profilesManager.deleteAll(account.getUuid());
keysDynamoDb.delete(account);
keysDynamoDb.delete(account.getUuid());
messagesManager.clear(account.getUuid());
deleteStorageServiceDataFuture.join();

View File

@ -56,7 +56,7 @@ public class KeysDynamoDb extends AbstractDynamoDbStore {
public void store(final Account account, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId);
delete(account.getUuid(), deviceId);
writeInBatches(keys, batch -> {
List<WriteRequest> items = new ArrayList<>();
@ -152,41 +152,41 @@ public class KeysDynamoDb extends AbstractDynamoDbStore {
});
}
public void delete(final Account account) {
public void delete(final UUID accountUuid) {
DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid())))
":uuid", getPartitionKey(accountUuid)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
deleteItemsForAccountMatchingQuery(account, queryRequest);
deleteItemsForAccountMatchingQuery(accountUuid, queryRequest);
});
}
public void delete(final Account account, final long deviceId) {
public void delete(final UUID accountUuid, final long deviceId) {
DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":uuid", getPartitionKey(accountUuid),
":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
deleteItemsForAccountMatchingQuery(account, queryRequest);
deleteItemsForAccountMatchingQuery(accountUuid, queryRequest);
});
}
private void deleteItemsForAccountMatchingQuery(final Account account, final QueryRequest querySpec) {
final AttributeValue partitionKey = getPartitionKey(account.getUuid());
private void deleteItemsForAccountMatchingQuery(final UUID accountUuid, final QueryRequest querySpec) {
final AttributeValue partitionKey = getPartitionKey(accountUuid);
writeInBatches(db().query(querySpec).items(), batch -> {
List<WriteRequest> deletes = new ArrayList<>();

View File

@ -116,7 +116,7 @@ public class KeysDynamoDbTest {
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
keysDynamoDb.delete(account);
keysDynamoDb.delete(account.getUuid());
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1));
@ -130,7 +130,7 @@ public class KeysDynamoDbTest {
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
keysDynamoDb.delete(account, DEVICE_ID);
keysDynamoDb.delete(account.getUuid(), DEVICE_ID);
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));

View File

@ -7,7 +7,6 @@ package org.whispersystems.textsecuregcm.tests.controllers;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
@ -382,9 +381,7 @@ class DeviceControllerTest {
verify(messagesManager, times(2)).clear(AuthHelper.VALID_UUID, deviceId);
verify(accountsManager, times(1)).update(eq(AuthHelper.VALID_ACCOUNT), any());
verify(AuthHelper.VALID_ACCOUNT).removeDevice(deviceId);
// The account instance may have changed as part of a call to `AccountManager#update`
verify(keys).delete(argThat(account -> account.getUuid().equals(AuthHelper.VALID_UUID)), eq(deviceId));
verify(keys).delete(AuthHelper.VALID_UUID, deviceId);
}
}