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); messages.clear(account.getUuid(), deviceId);
account = accounts.update(account, a -> a.removeDevice(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 // ensure any messages that came in after the first clear() are also removed
messages.clear(account.getUuid(), deviceId); messages.clear(account.getUuid(), deviceId);
} }

View File

@ -223,7 +223,7 @@ public class AccountsManager {
maybeExistingAccount.ifPresent(definitelyExistingAccount -> { maybeExistingAccount.ifPresent(definitelyExistingAccount -> {
messagesManager.clear(definitelyExistingAccount.getUuid()); messagesManager.clear(definitelyExistingAccount.getUuid());
keysDynamoDb.delete(definitelyExistingAccount); keysDynamoDb.delete(definitelyExistingAccount.getUuid());
}); });
pendingAccounts.remove(number); pendingAccounts.remove(number);
@ -393,7 +393,7 @@ public class AccountsManager {
usernamesManager.delete(account.getUuid()); usernamesManager.delete(account.getUuid());
directoryQueue.deleteAccount(account); directoryQueue.deleteAccount(account);
profilesManager.deleteAll(account.getUuid()); profilesManager.deleteAll(account.getUuid());
keysDynamoDb.delete(account); keysDynamoDb.delete(account.getUuid());
messagesManager.clear(account.getUuid()); messagesManager.clear(account.getUuid());
deleteStorageServiceDataFuture.join(); 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) { public void store(final Account account, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> { STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId); delete(account.getUuid(), deviceId);
writeInBatches(keys, batch -> { writeInBatches(keys, batch -> {
List<WriteRequest> items = new ArrayList<>(); 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(() -> { DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder() final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName) .tableName(tableName)
.keyConditionExpression("#uuid = :uuid") .keyConditionExpression("#uuid = :uuid")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID)) .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
.expressionAttributeValues(Map.of( .expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()))) ":uuid", getPartitionKey(accountUuid)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID) .projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true) .consistentRead(true)
.build(); .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(() -> { DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder() final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName) .tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)") .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID)) .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of( .expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()), ":uuid", getPartitionKey(accountUuid),
":sortprefix", getSortKeyPrefix(deviceId))) ":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID) .projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true) .consistentRead(true)
.build(); .build();
deleteItemsForAccountMatchingQuery(account, queryRequest); deleteItemsForAccountMatchingQuery(accountUuid, queryRequest);
}); });
} }
private void deleteItemsForAccountMatchingQuery(final Account account, final QueryRequest querySpec) { private void deleteItemsForAccountMatchingQuery(final UUID accountUuid, final QueryRequest querySpec) {
final AttributeValue partitionKey = getPartitionKey(account.getUuid()); final AttributeValue partitionKey = getPartitionKey(accountUuid);
writeInBatches(db().query(querySpec).items(), batch -> { writeInBatches(db().query(querySpec).items(), batch -> {
List<WriteRequest> deletes = new ArrayList<>(); List<WriteRequest> deletes = new ArrayList<>();

View File

@ -116,7 +116,7 @@ public class KeysDynamoDbTest {
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID)); assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); 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));
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1)); assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1));
@ -130,7 +130,7 @@ public class KeysDynamoDbTest {
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID)); assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); 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(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); 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.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.eq; import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
@ -382,9 +381,7 @@ class DeviceControllerTest {
verify(messagesManager, times(2)).clear(AuthHelper.VALID_UUID, deviceId); verify(messagesManager, times(2)).clear(AuthHelper.VALID_UUID, deviceId);
verify(accountsManager, times(1)).update(eq(AuthHelper.VALID_ACCOUNT), any()); verify(accountsManager, times(1)).update(eq(AuthHelper.VALID_ACCOUNT), any());
verify(AuthHelper.VALID_ACCOUNT).removeDevice(deviceId); verify(AuthHelper.VALID_ACCOUNT).removeDevice(deviceId);
verify(keys).delete(AuthHelper.VALID_UUID, 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));
} }
} }