diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index 626aa1702..435a4cdae 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -182,7 +182,7 @@ import org.whispersystems.textsecuregcm.storage.DirectoryReconciliationClient; import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.MessagePersister; import org.whispersystems.textsecuregcm.storage.MessagesCache; import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb; @@ -379,7 +379,7 @@ public class WhisperServerService extends Application commonControllers = Lists.newArrayList( new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket()), @@ -657,7 +657,7 @@ public class WhisperServerService extends Application maxDeviceConfiguration; public DeviceController(StoredVerificationCodeManager pendingDevices, AccountsManager accounts, MessagesManager messages, - KeysDynamoDb keys, + Keys keys, RateLimiters rateLimiters, Map maxDeviceConfiguration) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java index 965cdcaee..df8c2923e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java @@ -44,7 +44,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.Device; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.util.Util; @SuppressWarnings("OptionalUsedAsFieldOrParameterType") @@ -52,7 +52,7 @@ import org.whispersystems.textsecuregcm.util.Util; public class KeysController { private final RateLimiters rateLimiters; - private final KeysDynamoDb keysDynamoDb; + private final Keys keys; private final AccountsManager accounts; private final PreKeyRateLimiter preKeyRateLimiter; @@ -64,11 +64,11 @@ public class KeysController { private static final String SOURCE_COUNTRY_TAG_NAME = "sourceCountry"; private static final String INTERNATIONAL_TAG_NAME = "international"; - public KeysController(RateLimiters rateLimiters, KeysDynamoDb keysDynamoDb, AccountsManager accounts, + public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts, PreKeyRateLimiter preKeyRateLimiter, RateLimitChallengeManager rateLimitChallengeManager) { this.rateLimiters = rateLimiters; - this.keysDynamoDb = keysDynamoDb; + this.keys = keys; this.accounts = accounts; this.preKeyRateLimiter = preKeyRateLimiter; this.rateLimitChallengeManager = rateLimitChallengeManager; @@ -77,7 +77,7 @@ public class KeysController { @GET @Produces(MediaType.APPLICATION_JSON) public PreKeyCount getStatus(@Auth AuthenticatedAccount auth) { - int count = keysDynamoDb.getCount(auth.getAccount(), auth.getAuthenticatedDevice().getId()); + int count = keys.getCount(auth.getAccount(), auth.getAuthenticatedDevice().getId()); if (count > 0) { count = count - 1; @@ -109,7 +109,7 @@ public class KeysController { }); } - keysDynamoDb.store(account, device.getId(), preKeys.getPreKeys()); + keys.store(account, device.getId(), preKeys.getPreKeys()); } @Timed @@ -213,12 +213,12 @@ public class KeysController { private Map getLocalKeys(Account destination, String deviceIdSelector) { try { if (deviceIdSelector.equals("*")) { - return keysDynamoDb.take(destination); + return keys.take(destination); } long deviceId = Long.parseLong(deviceIdSelector); - return keysDynamoDb.take(destination, deviceId) + return keys.take(destination, deviceId) .map(preKey -> Map.of(deviceId, preKey)) .orElse(Collections.emptyMap()); } catch (NumberFormatException e) { 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 db7cc42ed..d6c17de3a 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -68,7 +68,7 @@ public class AccountsManager { private final FaultTolerantRedisCluster cacheCluster; private final DeletedAccountsManager deletedAccountsManager; private final DirectoryQueue directoryQueue; - private final KeysDynamoDb keysDynamoDb; + private final Keys keys; private final MessagesManager messagesManager; private final UsernamesManager usernamesManager; private final ProfilesManager profilesManager; @@ -96,7 +96,7 @@ public class AccountsManager { final FaultTolerantRedisCluster cacheCluster, final DeletedAccountsManager deletedAccountsManager, final DirectoryQueue directoryQueue, - final KeysDynamoDb keysDynamoDb, + final Keys keys, final MessagesManager messagesManager, final UsernamesManager usernamesManager, final ProfilesManager profilesManager, @@ -110,7 +110,7 @@ public class AccountsManager { this.cacheCluster = cacheCluster; this.deletedAccountsManager = deletedAccountsManager; this.directoryQueue = directoryQueue; - this.keysDynamoDb = keysDynamoDb; + this.keys = keys; this.messagesManager = messagesManager; this.usernamesManager = usernamesManager; this.profilesManager = profilesManager; @@ -178,7 +178,7 @@ public class AccountsManager { // account and need to clear out messages and keys that may have been stored for the old account. if (!originalUuid.equals(actualUuid)) { messagesManager.clear(actualUuid); - keysDynamoDb.delete(actualUuid); + keys.delete(actualUuid); profilesManager.deleteAll(actualUuid); } @@ -437,11 +437,11 @@ public class AccountsManager { usernamesManager.delete(account.getUuid()); profilesManager.deleteAll(account.getUuid()); - keysDynamoDb.delete(account.getUuid()); + keys.delete(account.getUuid()); messagesManager.clear(account.getUuid()); account.getPhoneNumberIdentifier().ifPresent(pni -> { - keysDynamoDb.delete(pni); + keys.delete(pni); messagesManager.clear(pni); }); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java similarity index 93% rename from service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java rename to service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java index 975667443..f5d94d0f1 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java @@ -32,7 +32,7 @@ import software.amazon.awssdk.services.dynamodb.model.ReturnValue; import software.amazon.awssdk.services.dynamodb.model.Select; import software.amazon.awssdk.services.dynamodb.model.WriteRequest; -public class KeysDynamoDb extends AbstractDynamoDbStore { +public class Keys extends AbstractDynamoDbStore { private final String tableName; @@ -40,16 +40,16 @@ public class KeysDynamoDb extends AbstractDynamoDbStore { static final String KEY_DEVICE_ID_KEY_ID = "DK"; static final String KEY_PUBLIC_KEY = "P"; - private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(KeysDynamoDb.class, "storeKeys")); - private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice")); - private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForAccount")); - private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "getKeyCount")); - private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForDevice")); - private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForAccount")); - private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "contestedKeys")); - private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "keyCount")); + private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(Keys.class, "storeKeys")); + private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "takeKeyForDevice")); + private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "takeKeyForAccount")); + private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(Keys.class, "getKeyCount")); + private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForDevice")); + private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForAccount")); + private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(Keys.class, "contestedKeys")); + private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(Keys.class, "keyCount")); - public KeysDynamoDb(final DynamoDbClient dynamoDB, final String tableName) { + public Keys(final DynamoDbClient dynamoDB, final String tableName) { super(dynamoDB); this.tableName = tableName; } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java index b74275607..8ad9a4a80 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java @@ -44,7 +44,7 @@ import org.whispersystems.textsecuregcm.storage.DeletedAccounts; import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager; import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.MessagesCache; import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb; import org.whispersystems.textsecuregcm.storage.MessagesManager; @@ -167,7 +167,7 @@ public class DeleteUserCommand extends EnvironmentCommand maybeAccount; diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java index ee1373fa5..fdfaf9129 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java @@ -189,7 +189,7 @@ class AccountsManagerChangeNumberIntegrationTest { CACHE_CLUSTER_EXTENSION.getRedisCluster(), deletedAccountsManager, mock(DirectoryQueue.class), - mock(KeysDynamoDb.class), + mock(Keys.class), mock(MessagesManager.class), mock(UsernamesManager.class), mock(ProfilesManager.class), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java index 509bd5eeb..366bbe852 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java @@ -36,7 +36,6 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.ArgumentCaptor; -import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials; import org.whispersystems.textsecuregcm.entities.AccountAttributes; @@ -146,7 +145,7 @@ class AccountsManagerConcurrentModificationIntegrationTest { RedisClusterHelper.buildMockRedisCluster(commands), deletedAccountsManager, mock(DirectoryQueue.class), - mock(KeysDynamoDb.class), + mock(Keys.class), mock(MessagesManager.class), mock(UsernamesManager.class), mock(ProfilesManager.class), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbRule.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbRule.java index 8e5d5e464..297daf432 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbRule.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbRule.java @@ -22,15 +22,15 @@ public class KeysDynamoDbRule extends LocalDynamoDbRule { getDynamoDbClient().createTable(CreateTableRequest.builder() .tableName(TABLE_NAME) .keySchema( - KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(), - KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE) + KeySchemaElement.builder().attributeName(Keys.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(), + KeySchemaElement.builder().attributeName(Keys.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE) .build()) .attributeDefinitions(AttributeDefinition.builder() - .attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID) + .attributeName(Keys.KEY_ACCOUNT_UUID) .attributeType(ScalarAttributeType.B) .build(), AttributeDefinition.builder() - .attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID) + .attributeName(Keys.KEY_DEVICE_ID_KEY_ID) .attributeType(ScalarAttributeType.B) .build()) .provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build()) diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java deleted file mode 100644 index 973650752..000000000 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java +++ /dev/null @@ -1,144 +0,0 @@ -/* - * Copyright 2021 Signal Messenger, LLC - * SPDX-License-Identifier: AGPL-3.0-only - */ - -package org.whispersystems.textsecuregcm.storage; - -import org.junit.Before; -import org.junit.ClassRule; -import org.junit.Test; -import org.whispersystems.textsecuregcm.entities.PreKey; -import software.amazon.awssdk.services.dynamodb.model.AttributeValue; - -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.Set; -import java.util.UUID; - -import static org.junit.Assert.assertArrayEquals; -import static org.junit.Assert.assertEquals; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class KeysDynamoDbTest { - - private Account account; - private KeysDynamoDb keysDynamoDb; - - @ClassRule - public static KeysDynamoDbRule dynamoDbRule = new KeysDynamoDbRule(); - - private static final String ACCOUNT_NUMBER = "+18005551234"; - private static final long DEVICE_ID = 1L; - - @Before - public void setup() { - keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME); - - account = mock(Account.class); - when(account.getNumber()).thenReturn(ACCOUNT_NUMBER); - when(account.getUuid()).thenReturn(UUID.randomUUID()); - } - - @Test - public void testStore() { - assertEquals("Initial pre-key count for an account should be zero", - 0, keysDynamoDb.getCount(account, DEVICE_ID)); - - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID)); - - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); - assertEquals("Repeatedly storing same key should have no effect", - 1, keysDynamoDb.getCount(account, DEVICE_ID)); - - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(2, "different-public-key"))); - assertEquals("Inserting a new key should overwrite all prior keys for the given account/device", - 1, keysDynamoDb.getCount(account, DEVICE_ID)); - - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(3, "third-public-key"), new PreKey(4, "fourth-public-key"))); - assertEquals("Inserting multiple new keys should overwrite all prior keys for the given account/device", - 2, keysDynamoDb.getCount(account, DEVICE_ID)); - } - - @Test - public void testTakeAccount() { - final Device firstDevice = mock(Device.class); - final Device secondDevice = mock(Device.class); - - when(firstDevice.getId()).thenReturn(DEVICE_ID); - when(secondDevice.getId()).thenReturn(DEVICE_ID + 1); - when(account.getDevices()).thenReturn(Set.of(firstDevice, secondDevice)); - - assertEquals(Collections.emptyMap(), keysDynamoDb.take(account)); - - final PreKey firstDevicePreKey = new PreKey(1, "public-key"); - final PreKey secondDevicePreKey = new PreKey(2, "second-key"); - - keysDynamoDb.store(account, DEVICE_ID, List.of(firstDevicePreKey)); - keysDynamoDb.store(account, DEVICE_ID + 1, List.of(secondDevicePreKey)); - - final Map expectedKeys = Map.of(DEVICE_ID, firstDevicePreKey, - DEVICE_ID + 1, secondDevicePreKey); - - assertEquals(expectedKeys, keysDynamoDb.take(account)); - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID)); - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1)); - } - - @Test - public void testTakeAccountAndDeviceId() { - assertEquals(Optional.empty(), keysDynamoDb.take(account, DEVICE_ID)); - - final PreKey preKey = new PreKey(1, "public-key"); - - keysDynamoDb.store(account, DEVICE_ID, List.of(preKey, new PreKey(2, "different-pre-key"))); - assertEquals(Optional.of(preKey), keysDynamoDb.take(account, DEVICE_ID)); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID)); - } - - @Test - public void testGetCount() { - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID)); - - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID)); - } - - @Test - public void testDeleteByAccount() { - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key"))); - keysDynamoDb.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device"))); - - assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID)); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); - - keysDynamoDb.delete(account.getUuid()); - - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID)); - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1)); - } - - @Test - public void testDeleteByAccountAndDevice() { - keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key"))); - keysDynamoDb.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device"))); - - assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID)); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); - - keysDynamoDb.delete(account.getUuid(), DEVICE_ID); - - assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID)); - assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); - } - - @Test - public void testSortKeyPrefix() { - AttributeValue got = KeysDynamoDb.getSortKeyPrefix(123); - assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, got.b().asByteArray()); - } -} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysTest.java new file mode 100644 index 000000000..8b2ac3319 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysTest.java @@ -0,0 +1,144 @@ +/* + * Copyright 2021 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.storage; + +import org.junit.Before; +import org.junit.ClassRule; +import org.junit.Test; +import org.whispersystems.textsecuregcm.entities.PreKey; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class KeysTest { + + private Account account; + private Keys keys; + + @ClassRule + public static KeysDynamoDbRule dynamoDbRule = new KeysDynamoDbRule(); + + private static final String ACCOUNT_NUMBER = "+18005551234"; + private static final long DEVICE_ID = 1L; + + @Before + public void setup() { + keys = new Keys(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME); + + account = mock(Account.class); + when(account.getNumber()).thenReturn(ACCOUNT_NUMBER); + when(account.getUuid()).thenReturn(UUID.randomUUID()); + } + + @Test + public void testStore() { + assertEquals("Initial pre-key count for an account should be zero", + 0, keys.getCount(account, DEVICE_ID)); + + keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); + assertEquals(1, keys.getCount(account, DEVICE_ID)); + + keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); + assertEquals("Repeatedly storing same key should have no effect", + 1, keys.getCount(account, DEVICE_ID)); + + keys.store(account, DEVICE_ID, List.of(new PreKey(2, "different-public-key"))); + assertEquals("Inserting a new key should overwrite all prior keys for the given account/device", + 1, keys.getCount(account, DEVICE_ID)); + + keys.store(account, DEVICE_ID, List.of(new PreKey(3, "third-public-key"), new PreKey(4, "fourth-public-key"))); + assertEquals("Inserting multiple new keys should overwrite all prior keys for the given account/device", + 2, keys.getCount(account, DEVICE_ID)); + } + + @Test + public void testTakeAccount() { + final Device firstDevice = mock(Device.class); + final Device secondDevice = mock(Device.class); + + when(firstDevice.getId()).thenReturn(DEVICE_ID); + when(secondDevice.getId()).thenReturn(DEVICE_ID + 1); + when(account.getDevices()).thenReturn(Set.of(firstDevice, secondDevice)); + + assertEquals(Collections.emptyMap(), keys.take(account)); + + final PreKey firstDevicePreKey = new PreKey(1, "public-key"); + final PreKey secondDevicePreKey = new PreKey(2, "second-key"); + + keys.store(account, DEVICE_ID, List.of(firstDevicePreKey)); + keys.store(account, DEVICE_ID + 1, List.of(secondDevicePreKey)); + + final Map expectedKeys = Map.of(DEVICE_ID, firstDevicePreKey, + DEVICE_ID + 1, secondDevicePreKey); + + assertEquals(expectedKeys, keys.take(account)); + assertEquals(0, keys.getCount(account, DEVICE_ID)); + assertEquals(0, keys.getCount(account, DEVICE_ID + 1)); + } + + @Test + public void testTakeAccountAndDeviceId() { + assertEquals(Optional.empty(), keys.take(account, DEVICE_ID)); + + final PreKey preKey = new PreKey(1, "public-key"); + + keys.store(account, DEVICE_ID, List.of(preKey, new PreKey(2, "different-pre-key"))); + assertEquals(Optional.of(preKey), keys.take(account, DEVICE_ID)); + assertEquals(1, keys.getCount(account, DEVICE_ID)); + } + + @Test + public void testGetCount() { + assertEquals(0, keys.getCount(account, DEVICE_ID)); + + keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"))); + assertEquals(1, keys.getCount(account, DEVICE_ID)); + } + + @Test + public void testDeleteByAccount() { + keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key"))); + keys.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device"))); + + assertEquals(2, keys.getCount(account, DEVICE_ID)); + assertEquals(1, keys.getCount(account, DEVICE_ID + 1)); + + keys.delete(account.getUuid()); + + assertEquals(0, keys.getCount(account, DEVICE_ID)); + assertEquals(0, keys.getCount(account, DEVICE_ID + 1)); + } + + @Test + public void testDeleteByAccountAndDevice() { + keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key"))); + keys.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device"))); + + assertEquals(2, keys.getCount(account, DEVICE_ID)); + assertEquals(1, keys.getCount(account, DEVICE_ID + 1)); + + keys.delete(account.getUuid(), DEVICE_ID); + + assertEquals(0, keys.getCount(account, DEVICE_ID)); + assertEquals(1, keys.getCount(account, DEVICE_ID + 1)); + } + + @Test + public void testSortKeyPrefix() { + AttributeValue got = Keys.getSortKeyPrefix(123); + assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, got.b().asByteArray()); + } +} 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 e352646c2..484760381 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 @@ -52,7 +52,7 @@ import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.MessagesManager; import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager; import org.whispersystems.textsecuregcm.tests.util.AccountsHelper; @@ -66,7 +66,7 @@ class DeviceControllerTest { public DumbVerificationDeviceController(StoredVerificationCodeManager pendingDevices, AccountsManager accounts, MessagesManager messages, - KeysDynamoDb keys, + Keys keys, RateLimiters rateLimiters, Map deviceConfiguration) { @@ -82,7 +82,7 @@ class DeviceControllerTest { private static StoredVerificationCodeManager pendingDevicesManager = mock(StoredVerificationCodeManager.class); private static AccountsManager accountsManager = mock(AccountsManager.class ); private static MessagesManager messagesManager = mock(MessagesManager.class); - private static KeysDynamoDb keys = mock(KeysDynamoDb.class); + private static Keys keys = mock(Keys.class); private static RateLimiters rateLimiters = mock(RateLimiters.class ); private static RateLimiter rateLimiter = mock(RateLimiter.class ); private static Account account = mock(Account.class ); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java index ee90018c6..bf6d2fae6 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java @@ -61,7 +61,7 @@ import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.Device; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.tests.util.AccountsHelper; import org.whispersystems.textsecuregcm.tests.util.AuthHelper; @@ -89,7 +89,7 @@ class KeysControllerTest { private final SignedPreKey SAMPLE_SIGNED_KEY3 = new SignedPreKey( 3333, "barfoo", "sig33" ); private final SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid"); - private final static KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class ); + private final static Keys KEYS = mock(Keys.class ); private final static AccountsManager accounts = mock(AccountsManager.class ); private final static PreKeyRateLimiter preKeyRateLimiter = mock(PreKeyRateLimiter.class ); private final static RateLimitChallengeManager rateLimitChallengeManager = mock(RateLimitChallengeManager.class ); @@ -106,7 +106,7 @@ class KeysControllerTest { .addResource(new RateLimitChallengeExceptionMapper(rateLimitChallengeManager)) .addResource(new ServerRejectedExceptionMapper()) .addResource( - new KeysController(rateLimiters, keysDynamoDb, accounts, preKeyRateLimiter, rateLimitChallengeManager)) + new KeysController(rateLimiters, KEYS, accounts, preKeyRateLimiter, rateLimitChallengeManager)) .build(); @BeforeEach @@ -161,14 +161,14 @@ class KeysControllerTest { when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter); - when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(Optional.of(SAMPLE_KEY)); + when(KEYS.take(eq(existsAccount), eq(1L))).thenReturn(Optional.of(SAMPLE_KEY)); - when(keysDynamoDb.take(existsAccount)).thenReturn(Map.of(1L, SAMPLE_KEY, + when(KEYS.take(existsAccount)).thenReturn(Map.of(1L, SAMPLE_KEY, 2L, SAMPLE_KEY2, 3L, SAMPLE_KEY3, 4L, SAMPLE_KEY4)); - when(keysDynamoDb.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5); + when(KEYS.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5); when(AuthHelper.VALID_DEVICE.getSignedPreKey()).thenReturn(VALID_DEVICE_SIGNED_KEY); when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null); @@ -177,7 +177,7 @@ class KeysControllerTest { @AfterEach void teardown() { reset( - keysDynamoDb, + KEYS, accounts, preKeyRateLimiter, existsAccount, @@ -198,7 +198,7 @@ class KeysControllerTest { assertThat(result.getCount()).isEqualTo(4); - verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); + verify(KEYS).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); } @@ -257,8 +257,8 @@ class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keysDynamoDb); + verify(KEYS).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(KEYS); } @Test @@ -275,8 +275,8 @@ class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keysDynamoDb); + verify(KEYS).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(KEYS); } @Test @@ -303,7 +303,7 @@ class KeysControllerTest { .get(); assertThat(response.getStatus()).isEqualTo(401); - verifyNoMoreInteractions(keysDynamoDb); + verifyNoMoreInteractions(KEYS); } @Test @@ -315,7 +315,7 @@ class KeysControllerTest { .get(); assertThat(response.getStatus()).isEqualTo(401); - verifyNoMoreInteractions(keysDynamoDb); + verifyNoMoreInteractions(KEYS); } @@ -365,8 +365,8 @@ class KeysControllerTest { assertThat(signedPreKey).isNull(); assertThat(deviceId).isEqualTo(4); - verify(keysDynamoDb).take(eq(existsAccount)); - verifyNoMoreInteractions(keysDynamoDb); + verify(KEYS).take(eq(existsAccount)); + verifyNoMoreInteractions(KEYS); } @@ -434,7 +434,7 @@ class KeysControllerTest { assertThat(response.getStatus()).isEqualTo(204); ArgumentCaptor listCaptor = ArgumentCaptor.forClass(List.class); - verify(keysDynamoDb).store(eqUuid(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture()); + verify(KEYS).store(eqUuid(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture()); List capturedList = listCaptor.getValue(); assertThat(capturedList.size()).isEqualTo(1); @@ -468,7 +468,7 @@ class KeysControllerTest { assertThat(response.getStatus()).isEqualTo(204); ArgumentCaptor listCaptor = ArgumentCaptor.forClass(List.class); - verify(keysDynamoDb).store(eqUuid(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture()); + verify(KEYS).store(eqUuid(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture()); List capturedList = listCaptor.getValue(); assertThat(capturedList.size()).isEqualTo(1); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java index 7a9bca8f3..e7d1a373e 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java @@ -56,7 +56,7 @@ import org.whispersystems.textsecuregcm.storage.ContestedOptimisticLockException import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager; import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities; -import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; +import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.MessagesManager; import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers; import org.whispersystems.textsecuregcm.storage.ProfilesManager; @@ -69,7 +69,7 @@ class AccountsManagerTest { private Accounts accounts; private DeletedAccountsManager deletedAccountsManager; private DirectoryQueue directoryQueue; - private KeysDynamoDb keys; + private Keys keys; private MessagesManager messagesManager; private ProfilesManager profilesManager; @@ -89,7 +89,7 @@ class AccountsManagerTest { accounts = mock(Accounts.class); deletedAccountsManager = mock(DeletedAccountsManager.class); directoryQueue = mock(DirectoryQueue.class); - keys = mock(KeysDynamoDb.class); + keys = mock(Keys.class); messagesManager = mock(MessagesManager.class); profilesManager = mock(ProfilesManager.class);