Remove row-based one-time PQ key store

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

@@ -46,7 +46,6 @@ import org.whispersystems.textsecuregcm.entities.ApnRegistrationId;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
-import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClient;
 import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
@@ -68,7 +67,6 @@ public class AccountCreationDeletionIntegrationTest {
       DynamoDbExtensionSchema.Tables.PNI_ASSIGNMENTS,
       DynamoDbExtensionSchema.Tables.USERNAMES,
       DynamoDbExtensionSchema.Tables.EC_KEYS,
-      DynamoDbExtensionSchema.Tables.PQ_KEYS,
       DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS,
       DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS,
       DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
@@ -101,7 +99,6 @@ public class AccountCreationDeletionIntegrationTest {
     final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
     keysManager = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.EC_KEYS.tableName()),
-        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName()),
         new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
             S3_EXTENSION.getS3Client(),
             DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),

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

@@ -60,7 +60,6 @@ class AccountsManagerChangeNumberIntegrationTest {
       Tables.PNI_ASSIGNMENTS,
       Tables.USERNAMES,
       Tables.EC_KEYS,
-      Tables.PQ_KEYS,
       Tables.PAGED_PQ_KEYS,
       Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS,
       Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
@@ -90,7 +89,6 @@ class AccountsManagerChangeNumberIntegrationTest {
       final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
       keysManager = new KeysManager(
           new SingleUseECPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.EC_KEYS.tableName()),
-          new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName()),
           new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
               S3_EXTENSION.getS3Client(),
               DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),

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

@@ -71,7 +71,6 @@ class AccountsManagerConcurrentModificationIntegrationTest {
       Tables.PNI_ASSIGNMENTS,
       Tables.DELETED_ACCOUNTS,
       Tables.EC_KEYS,
-      Tables.PQ_KEYS,
       Tables.PAGED_PQ_KEYS,
       Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS,
       Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);

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

@@ -73,7 +73,6 @@ class AccountsManagerUsernameIntegrationTest {
       Tables.PNI,
       Tables.PNI_ASSIGNMENTS,
       Tables.EC_KEYS,
-      Tables.PQ_KEYS,
       Tables.PAGED_PQ_KEYS,
       Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS,
       Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
@@ -103,7 +102,6 @@ class AccountsManagerUsernameIntegrationTest {
     final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
     final KeysManager keysManager = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.EC_KEYS.tableName()),
-        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName()),
         new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
             S3_EXTENSION.getS3Client(),
             DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),

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

@@ -62,7 +62,6 @@ public class AddRemoveDeviceIntegrationTest {
       DynamoDbExtensionSchema.Tables.PNI_ASSIGNMENTS,
       DynamoDbExtensionSchema.Tables.USERNAMES,
       DynamoDbExtensionSchema.Tables.EC_KEYS,
-      DynamoDbExtensionSchema.Tables.PQ_KEYS,
       DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS,
       DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS,
       DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
@@ -98,7 +97,6 @@ public class AddRemoveDeviceIntegrationTest {
     final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
     keysManager = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.EC_KEYS.tableName()),
-        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName()),
         new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
             S3_EXTENSION.getS3Client(),
             DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),

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

@@ -116,29 +116,15 @@ public final class DynamoDbExtensionSchema {
         List.of(), List.of()),
 
     EC_KEYS("keys_test",
-        SingleUsePreKeyStore.KEY_ACCOUNT_UUID,
-        SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID,
+        SingleUseECPreKeyStore.KEY_ACCOUNT_UUID,
+        SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID,
         List.of(
             AttributeDefinition.builder()
-                .attributeName(SingleUsePreKeyStore.KEY_ACCOUNT_UUID)
+                .attributeName(SingleUseECPreKeyStore.KEY_ACCOUNT_UUID)
                 .attributeType(ScalarAttributeType.B)
                 .build(),
             AttributeDefinition.builder()
-                .attributeName(SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID)
-                .attributeType(ScalarAttributeType.B)
-                .build()),
-        List.of(), List.of()),
-
-    PQ_KEYS("pq_keys_test",
-        SingleUsePreKeyStore.KEY_ACCOUNT_UUID,
-        SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID,
-        List.of(
-            AttributeDefinition.builder()
-                .attributeName(SingleUsePreKeyStore.KEY_ACCOUNT_UUID)
-                .attributeType(ScalarAttributeType.B)
-                .build(),
-            AttributeDefinition.builder()
-                .attributeName(SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID)
+                .attributeName(SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID)
                 .attributeType(ScalarAttributeType.B)
                 .build()),
         List.of(), List.of()),

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

@@ -30,12 +30,11 @@ class KeysManagerTest {
 
   private KeysManager keysManager;
 
-  private SingleUseKEMPreKeyStore singleUseKEMPreKeyStore;
   private PagedSingleUseKEMPreKeyStore pagedSingleUseKEMPreKeyStore;
 
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
-      Tables.EC_KEYS, Tables.PQ_KEYS, Tables.PAGED_PQ_KEYS,
+      Tables.EC_KEYS, Tables.PAGED_PQ_KEYS,
       Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
 
   @RegisterExtension
@@ -50,7 +49,6 @@ class KeysManagerTest {
   @BeforeEach
   void setup() {
     final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
-    singleUseKEMPreKeyStore = new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, Tables.PQ_KEYS.tableName());
     pagedSingleUseKEMPreKeyStore = new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
         S3_EXTENSION.getS3Client(),
         DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),
@@ -58,7 +56,6 @@ class KeysManagerTest {
 
     keysManager = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, Tables.EC_KEYS.tableName()),
-        singleUseKEMPreKeyStore,
         pagedSingleUseKEMPreKeyStore,
         new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient, Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
         new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));
@@ -77,24 +74,6 @@ class KeysManagerTest {
         "Repeatedly storing same key should have no effect");
   }
 
-  @Test
-  void storeKemOneTimePreKeysClearsOld() {
-    final List<KEMSignedPreKey> oldPreKeys = List.of(generateTestKEMSignedPreKey(1));
-
-    // Leave a key in the 'old' key store
-    singleUseKEMPreKeyStore.store(ACCOUNT_UUID, DEVICE_ID, oldPreKeys).join();
-
-    final List<KEMSignedPreKey> newPreKeys = List.of(generateTestKEMSignedPreKey(2));
-    keysManager.storeKemOneTimePreKeys(ACCOUNT_UUID, DEVICE_ID, newPreKeys).join();
-
-    assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());
-    assertEquals(1, pagedSingleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
-    assertEquals(0, singleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
-
-    final KEMSignedPreKey key = keysManager.takePQ(ACCOUNT_UUID, DEVICE_ID).join().orElseThrow();
-    assertEquals(2, key.keyId());
-  }
-
   @Test
   void storeKemOneTimePreKeys() {
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join(),
@@ -103,12 +82,10 @@ class KeysManagerTest {
     keysManager.storeKemOneTimePreKeys(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestKEMSignedPreKey(1))).join();
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());
     assertEquals(1, pagedSingleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
-    assertEquals(0, singleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
 
     keysManager.storeKemOneTimePreKeys(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestKEMSignedPreKey(1))).join();
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());
     assertEquals(1, pagedSingleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
-    assertEquals(0, singleUseKEMPreKeyStore.getCount(ACCOUNT_UUID, DEVICE_ID).join());
   }
 
 

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

@@ -5,7 +5,20 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.signal.libsignal.protocol.ecc.ECKeyPair;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
@@ -14,10 +27,10 @@ import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
 import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 import software.amazon.awssdk.services.dynamodb.paginators.ScanIterable;
-import java.util.Map;
 
-class SingleUseECPreKeyStoreTest extends SingleUsePreKeyStoreTest<ECPreKey> {
+class SingleUseECPreKeyStoreTest {
 
+  private static final int KEY_COUNT = 100;
   private SingleUseECPreKeyStore preKeyStore;
 
   @RegisterExtension
@@ -29,18 +42,11 @@ class SingleUseECPreKeyStoreTest extends SingleUsePreKeyStoreTest<ECPreKey> {
         DynamoDbExtensionSchema.Tables.EC_KEYS.tableName());
   }
 
-  @Override
-  protected SingleUsePreKeyStore<ECPreKey> getPreKeyStore() {
-    return preKeyStore;
-  }
-
-  @Override
-  protected ECPreKey generatePreKey(final long keyId) {
+  private ECPreKey generatePreKey(final long keyId) {
     return new ECPreKey(keyId, ECKeyPair.generate().getPublicKey());
   }
 
-  @Override
-  protected void clearKeyCountAttributes() {
+  private void clearKeyCountAttributes() {
     final ScanIterable scanIterable = DYNAMO_DB_EXTENSION.getDynamoDbClient().scanPaginator(ScanRequest.builder()
         .tableName(DynamoDbExtensionSchema.Tables.EC_KEYS.tableName())
         .build());
@@ -51,11 +57,113 @@ class SingleUseECPreKeyStoreTest extends SingleUsePreKeyStoreTest<ECPreKey> {
         DYNAMO_DB_EXTENSION.getDynamoDbClient().updateItem(UpdateItemRequest.builder()
             .tableName(DynamoDbExtensionSchema.Tables.EC_KEYS.tableName())
             .key(Map.of(
-                SingleUsePreKeyStore.KEY_ACCOUNT_UUID, item.get(SingleUsePreKeyStore.KEY_ACCOUNT_UUID),
-                SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID, item.get(SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID)))
-            .updateExpression("REMOVE " + SingleUsePreKeyStore.ATTR_REMAINING_KEYS)
+                SingleUseECPreKeyStore.KEY_ACCOUNT_UUID, item.get(SingleUseECPreKeyStore.KEY_ACCOUNT_UUID),
+                SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID, item.get(SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID)))
+            .updateExpression("REMOVE " + SingleUseECPreKeyStore.ATTR_REMAINING_KEYS)
             .build());
       }
     }
   }
+
+  @Test
+  void storeTake() {
+    final SingleUseECPreKeyStore preKeyStore = this.preKeyStore;
+
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(Optional.empty(), preKeyStore.take(accountIdentifier, deviceId).join());
+
+    final List<ECPreKey> sortedPreKeys;
+    {
+      final List<ECPreKey> preKeys = generateRandomPreKeys();
+      assertDoesNotThrow(() -> preKeyStore.store(accountIdentifier, deviceId, preKeys).join());
+
+      sortedPreKeys = new ArrayList<>(preKeys);
+      sortedPreKeys.sort(Comparator.comparing(preKey -> preKey.keyId()));
+    }
+
+    assertEquals(Optional.of(sortedPreKeys.get(0)), preKeyStore.take(accountIdentifier, deviceId).join());
+    assertEquals(Optional.of(sortedPreKeys.get(1)), preKeyStore.take(accountIdentifier, deviceId).join());
+  }
+
+  @Test
+  void getCount() {
+    final SingleUseECPreKeyStore preKeyStore = this.preKeyStore;
+
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+
+    final List<ECPreKey> preKeys = generateRandomPreKeys();
+
+    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
+
+    assertEquals(KEY_COUNT, preKeyStore.getCount(accountIdentifier, deviceId).join());
+
+    for (int i = 0; i < KEY_COUNT; i++) {
+      preKeyStore.take(accountIdentifier, deviceId).join();
+      assertEquals(KEY_COUNT - (i + 1), preKeyStore.getCount(accountIdentifier, deviceId).join());
+    }
+
+    preKeyStore.store(accountIdentifier, deviceId, List.of(generatePreKey(KEY_COUNT + 1))).join();
+    clearKeyCountAttributes();
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+  }
+
+  @Test
+  void deleteSingleDevice() {
+    final SingleUseECPreKeyStore preKeyStore = this.preKeyStore;
+
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier, deviceId).join());
+
+    final List<ECPreKey> preKeys = generateRandomPreKeys();
+
+    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
+    preKeyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
+
+    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier, deviceId).join());
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+    assertEquals(KEY_COUNT, preKeyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
+  }
+
+  @Test
+  void deleteAllDevices() {
+    final SingleUseECPreKeyStore preKeyStore = this.preKeyStore;
+
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier).join());
+
+    final List<ECPreKey> preKeys = generateRandomPreKeys();
+
+    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
+    preKeyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
+
+    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier).join());
+
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
+    assertEquals(0, preKeyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
+  }
+
+  private List<ECPreKey> generateRandomPreKeys() {
+    final Set<Integer> keyIds = new HashSet<>(KEY_COUNT);
+
+    while (keyIds.size() < KEY_COUNT) {
+      keyIds.add(Math.abs(ThreadLocalRandom.current().nextInt()));
+    }
+
+    return keyIds.stream()
+        .map(this::generatePreKey)
+        .toList();
+  }
 }

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

@@ -1,64 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.extension.RegisterExtension;
-import org.signal.libsignal.protocol.ecc.ECKeyPair;
-import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
-import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
-import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
-import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
-import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
-import software.amazon.awssdk.services.dynamodb.paginators.ScanIterable;
-import java.util.Map;
-
-class SingleUseKEMPreKeyStoreTest extends SingleUsePreKeyStoreTest<KEMSignedPreKey> {
-
-  private SingleUseKEMPreKeyStore preKeyStore;
-
-  private static final ECKeyPair IDENTITY_KEY_PAIR = ECKeyPair.generate();
-
-  @RegisterExtension
-  static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(DynamoDbExtensionSchema.Tables.PQ_KEYS);
-
-  @BeforeEach
-  void setUp() {
-    preKeyStore = new SingleUseKEMPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName());
-  }
-
-  @Override
-  protected SingleUsePreKeyStore<KEMSignedPreKey> getPreKeyStore() {
-    return preKeyStore;
-  }
-
-  @Override
-  protected KEMSignedPreKey generatePreKey(final long keyId) {
-    return KeysHelper.signedKEMPreKey(keyId, IDENTITY_KEY_PAIR);
-  }
-
-  @Override
-  protected void clearKeyCountAttributes() {
-    final ScanIterable scanIterable = DYNAMO_DB_EXTENSION.getDynamoDbClient().scanPaginator(ScanRequest.builder()
-        .tableName(DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName())
-        .build());
-
-    for (final ScanResponse response : scanIterable) {
-      for (final Map<String, AttributeValue> item : response.items()) {
-
-        DYNAMO_DB_EXTENSION.getDynamoDbClient().updateItem(UpdateItemRequest.builder()
-            .tableName(DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName())
-            .key(Map.of(
-                SingleUsePreKeyStore.KEY_ACCOUNT_UUID, item.get(SingleUsePreKeyStore.KEY_ACCOUNT_UUID),
-                SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID, item.get(SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID)))
-            .updateExpression("REMOVE " + SingleUsePreKeyStore.ATTR_REMAINING_KEYS)
-            .build());
-      }
-    }
-  }
-}

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

@@ -1,133 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
-import java.util.UUID;
-import java.util.concurrent.ThreadLocalRandom;
-import org.junit.jupiter.api.Test;
-import org.whispersystems.textsecuregcm.entities.PreKey;
-
-abstract class SingleUsePreKeyStoreTest<K extends PreKey<?>> {
-
-  private static final int KEY_COUNT = 100;
-
-  protected abstract SingleUsePreKeyStore<K> getPreKeyStore();
-
-  protected abstract K generatePreKey(final long keyId);
-
-  protected abstract void clearKeyCountAttributes();
-
-  @Test
-  void storeTake() {
-    final SingleUsePreKeyStore<K> preKeyStore = getPreKeyStore();
-
-    final UUID accountIdentifier = UUID.randomUUID();
-    final byte deviceId = 1;
-
-    assertEquals(Optional.empty(), preKeyStore.take(accountIdentifier, deviceId).join());
-
-    final List<K> sortedPreKeys;
-    {
-      final List<K> preKeys = generateRandomPreKeys();
-      assertDoesNotThrow(() -> preKeyStore.store(accountIdentifier, deviceId, preKeys).join());
-
-      sortedPreKeys = new ArrayList<>(preKeys);
-      sortedPreKeys.sort(Comparator.comparing(preKey -> preKey.keyId()));
-    }
-
-    assertEquals(Optional.of(sortedPreKeys.get(0)), preKeyStore.take(accountIdentifier, deviceId).join());
-    assertEquals(Optional.of(sortedPreKeys.get(1)), preKeyStore.take(accountIdentifier, deviceId).join());
-  }
-
-  @Test
-  void getCount() {
-    final SingleUsePreKeyStore<K> preKeyStore = getPreKeyStore();
-
-    final UUID accountIdentifier = UUID.randomUUID();
-    final byte deviceId = 1;
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-
-    final List<K> preKeys = generateRandomPreKeys();
-
-    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
-
-    assertEquals(KEY_COUNT, preKeyStore.getCount(accountIdentifier, deviceId).join());
-
-    for (int i = 0; i < KEY_COUNT; i++) {
-      preKeyStore.take(accountIdentifier, deviceId).join();
-      assertEquals(KEY_COUNT - (i + 1), preKeyStore.getCount(accountIdentifier, deviceId).join());
-    }
-
-    preKeyStore.store(accountIdentifier, deviceId, List.of(generatePreKey(KEY_COUNT + 1))).join();
-    clearKeyCountAttributes();
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-  }
-
-  @Test
-  void deleteSingleDevice() {
-    final SingleUsePreKeyStore<K> preKeyStore = getPreKeyStore();
-
-    final UUID accountIdentifier = UUID.randomUUID();
-    final byte deviceId = 1;
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier, deviceId).join());
-
-    final List<K> preKeys = generateRandomPreKeys();
-
-    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
-    preKeyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
-
-    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier, deviceId).join());
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-    assertEquals(KEY_COUNT, preKeyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
-  }
-
-  @Test
-  void deleteAllDevices() {
-    final SingleUsePreKeyStore<K> preKeyStore = getPreKeyStore();
-
-    final UUID accountIdentifier = UUID.randomUUID();
-    final byte deviceId = 1;
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier).join());
-
-    final List<K> preKeys = generateRandomPreKeys();
-
-    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();
-    preKeyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
-
-    assertDoesNotThrow(() -> preKeyStore.delete(accountIdentifier).join());
-
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
-    assertEquals(0, preKeyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
-  }
-
-  private List<K> generateRandomPreKeys() {
-    final Set<Integer> keyIds = new HashSet<>(KEY_COUNT);
-
-    while (keyIds.size() < KEY_COUNT) {
-      keyIds.add(Math.abs(ThreadLocalRandom.current().nextInt()));
-    }
-
-    return keyIds.stream()
-        .map(this::generatePreKey)
-        .toList();
-  }
-}