s/KeysDynamoDb/Keys/

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<WhisperServerConfiguration
     Usernames usernames = new Usernames(accountDatabase);
     ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
     Profiles profiles = new Profiles(accountDatabase);
-    KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName());
+    Keys keys = new Keys(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName());
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
         config.getMessageDynamoDbConfiguration().getTableName(),
         config.getMessageDynamoDbConfiguration().getTimeToLive());
@@ -467,7 +467,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
         deletedAccountsLockDynamoDbClient, config.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
+        deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
         pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
     RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
     DeadLetterHandler          deadLetterHandler          = new DeadLetterHandler(accountsManager, messagesManager);
@@ -649,7 +649,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             smsSender, dynamicConfigurationManager, turnTokenGenerator, config.getTestDevices(),
             transitionalRecaptchaClient, gcmSender, apnSender, backupCredentialsGenerator,
             verifyExperimentEnrollmentManager));
-    environment.jersey().register(new KeysController(rateLimiters, keysDynamoDb, accountsManager, preKeyRateLimiter, rateLimitChallengeManager));
+    environment.jersey().register(new KeysController(rateLimiters, keys, accountsManager, preKeyRateLimiter, rateLimitChallengeManager));
 
     final List<Object> commonControllers = Lists.newArrayList(
         new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket()),
@@ -657,7 +657,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new AttachmentControllerV3(rateLimiters, config.getGcpAttachmentsConfiguration().getDomain(), config.getGcpAttachmentsConfiguration().getEmail(), config.getGcpAttachmentsConfiguration().getMaxSizeInBytes(), config.getGcpAttachmentsConfiguration().getPathPrefix(), config.getGcpAttachmentsConfiguration().getRsaSigningKey()),
         new CertificateController(new CertificateGenerator(config.getDeliveryCertificate().getCertificate(), config.getDeliveryCertificate().getPrivateKey(), config.getDeliveryCertificate().getExpiresDays()), zkAuthOperations),
         new ChallengeController(rateLimitChallengeManager),
-        new DeviceController(pendingDevicesManager, accountsManager, messagesManager, keysDynamoDb, rateLimiters, config.getMaxDevices()),
+        new DeviceController(pendingDevicesManager, accountsManager, messagesManager, keys, rateLimiters, config.getMaxDevices()),
         new DirectoryController(directoryCredentialsGenerator),
         new DirectoryV2Controller(directoryV2CredentialsGenerator),
         new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),

service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java

@@ -41,7 +41,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.util.Util;
@@ -57,14 +57,14 @@ public class DeviceController {
   private final StoredVerificationCodeManager pendingDevices;
   private final AccountsManager       accounts;
   private final MessagesManager       messages;
-  private final KeysDynamoDb          keys;
+  private final Keys keys;
   private final RateLimiters          rateLimiters;
   private final Map<String, Integer>  maxDeviceConfiguration;
 
   public DeviceController(StoredVerificationCodeManager pendingDevices,
                           AccountsManager accounts,
                           MessagesManager messages,
-                          KeysDynamoDb keys,
+                          Keys keys,
                           RateLimiters rateLimiters,
                           Map<String, Integer> maxDeviceConfiguration)
   {

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<Long, PreKey> 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) {

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);
     });
 

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               STORE_KEYS_TIMER              = Metrics.timer(name(Keys.class, "storeKeys"));
-  private static final Timer               TAKE_KEY_FOR_DEVICE_TIMER     = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice"));
+  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(KeysDynamoDb.class, "takeKeyForAccount"));
+  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(KeysDynamoDb.class, "getKeyCount"));
+  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(KeysDynamoDb.class, "deleteKeysForDevice"));
+  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(KeysDynamoDb.class, "deleteKeysForAccount"));
+  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(KeysDynamoDb.class, "contestedKeys"));
+  private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION    = Metrics.summary(name(Keys.class, "contestedKeys"));
-  private static final DistributionSummary KEY_COUNT_DISTRIBUTION        = Metrics.summary(name(KeysDynamoDb.class, "keyCount"));
+  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;
   }

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<WhisperServerConfigura
       Usernames usernames = new Usernames(accountDatabase);
       Profiles profiles = new Profiles(accountDatabase);
       ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
-      KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeysDynamoDb,
+      Keys keys = new Keys(preKeysDynamoDb,
           configuration.getKeysDynamoDbConfiguration().getTableName());
       MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
           configuration.getMessageDynamoDbConfiguration().getTableName(),
@@ -207,7 +207,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
           configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
       StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
       AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-          deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
+          deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
           pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
 
       for (String user : users) {

service/src/main/java/org/whispersystems/textsecuregcm/workers/SetUserDiscoverabilityCommand.java

@@ -42,7 +42,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;
@@ -172,7 +172,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
       Usernames usernames = new Usernames(accountDatabase);
       Profiles profiles = new Profiles(accountDatabase);
       ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
-      KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeysDynamoDb,
+      Keys keys = new Keys(preKeysDynamoDb,
           configuration.getKeysDynamoDbConfiguration().getTableName());
       MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
           configuration.getMessageDynamoDbConfiguration().getTableName(),
@@ -210,7 +210,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
           configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
       StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
       AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-          deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
+          deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
           pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
 
       Optional<Account> maybeAccount;

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),

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),

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(Keys.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(),
-              KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE)
+              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())

service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java

@@ -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<Long, PreKey> 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());
-    }
-}

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<Long, PreKey> 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());
+    }
+}

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<String, Integer> 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               );

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));
+    verify(KEYS).take(eq(existsAccount), eq(1L));
-    verifyNoMoreInteractions(keysDynamoDb);
+    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));
+    verify(KEYS).take(eq(existsAccount), eq(1L));
-    verifyNoMoreInteractions(keysDynamoDb);
+    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));
+    verify(KEYS).take(eq(existsAccount));
-    verifyNoMoreInteractions(keysDynamoDb);
+    verifyNoMoreInteractions(KEYS);
   }
 
 
@@ -434,7 +434,7 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     ArgumentCaptor<List> 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<PreKey> capturedList = listCaptor.getValue();
     assertThat(capturedList.size()).isEqualTo(1);
@@ -468,7 +468,7 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     ArgumentCaptor<List> 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<PreKey> capturedList = listCaptor.getValue();
     assertThat(capturedList.size()).isEqualTo(1);

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);