From be20c04cd8395c5681fd192fec69c0123a810df2 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 26 Jul 2021 09:57:20 -0400 Subject: [PATCH] Identify accounts for which to delete keys by UUID. --- .../controllers/DeviceController.java | 2 +- .../textsecuregcm/storage/AccountsManager.java | 4 ++-- .../textsecuregcm/storage/KeysDynamoDb.java | 18 +++++++++--------- .../storage/KeysDynamoDbTest.java | 4 ++-- .../controllers/DeviceControllerTest.java | 5 +---- 5 files changed, 15 insertions(+), 18 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java index d1f22f988..0eeb16851 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java @@ -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); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index d27aa0228..05de23b7c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -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(); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java index fb0c924ab..975667443 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java @@ -56,7 +56,7 @@ public class KeysDynamoDb extends AbstractDynamoDbStore { public void store(final Account account, final long deviceId, final List keys) { STORE_KEYS_TIMER.record(() -> { - delete(account, deviceId); + delete(account.getUuid(), deviceId); writeInBatches(keys, batch -> { List 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 deletes = new ArrayList<>(); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java index 3dffbd7e0..973650752 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java @@ -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)); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java index e8bc69cc7..9d34464e2 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java @@ -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); } }