Remove row-based one-time PQ key store

service/config/sample.yml

@@ -108,8 +108,6 @@ dynamoDbTables:
     tableName: Example_Keys
   ecSignedPreKeys:
     tableName: Example_EC_Signed_Pre_Keys
-  pqKeys:
-    tableName: Example_PQ_Keys
   pagedPqKeys:
     tableName: Example_PQ_Paged_Keys
   pqLastResortKeys:

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -243,7 +243,6 @@ import org.whispersystems.textsecuregcm.storage.RepeatedUseKEMSignedPreKeyStore;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
 import org.whispersystems.textsecuregcm.storage.SingleUseECPreKeyStore;
-import org.whispersystems.textsecuregcm.storage.SingleUseKEMPreKeyStore;
 import org.whispersystems.textsecuregcm.storage.SubscriptionManager;
 import org.whispersystems.textsecuregcm.storage.Subscriptions;
 import org.whispersystems.textsecuregcm.storage.VerificationSessionManager;
@@ -469,7 +468,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         .build();
     KeysManager keysManager = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, config.getDynamoDbTables().getEcKeys().getTableName()),
-        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, config.getDynamoDbTables().getKemKeys().getTableName()),
         new PagedSingleUseKEMPreKeyStore(
             dynamoDbAsyncClient,
             asyncKeysS3Client,

service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbTables.java

@@ -58,7 +58,6 @@ public class DynamoDbTables {
   private final IssuedReceiptsTableConfiguration issuedReceipts;
   private final Table ecKeys;
   private final Table ecSignedPreKeys;
-  private final Table kemKeys;
   private final Table kemLastResortKeys;
   private final Table pagedKemKeys;
   private final TableWithExpiration messages;
@@ -87,7 +86,6 @@ public class DynamoDbTables {
       @JsonProperty("issuedReceipts") final IssuedReceiptsTableConfiguration issuedReceipts,
       @JsonProperty("ecKeys") final Table ecKeys,
       @JsonProperty("ecSignedPreKeys") final Table ecSignedPreKeys,
-      @JsonProperty("pqKeys") final Table kemKeys,
       @JsonProperty("pqLastResortKeys") final Table kemLastResortKeys,
       @JsonProperty("pagedPqKeys") final Table pagedKemKeys,
       @JsonProperty("messages") final TableWithExpiration messages,
@@ -115,7 +113,6 @@ public class DynamoDbTables {
     this.issuedReceipts = issuedReceipts;
     this.ecKeys = ecKeys;
     this.ecSignedPreKeys = ecSignedPreKeys;
-    this.kemKeys = kemKeys;
     this.pagedKemKeys = pagedKemKeys;
     this.kemLastResortKeys = kemLastResortKeys;
     this.messages = messages;
@@ -199,12 +196,6 @@ public class DynamoDbTables {
     return ecSignedPreKeys;
   }
 
-  @NotNull
-  @Valid
-  public Table getKemKeys() {
-    return kemKeys;
-  }
-
   @NotNull
   @Valid
   public Table getPagedKemKeys() {

service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java

@@ -7,12 +7,13 @@ package org.whispersystems.textsecuregcm.storage;
 
 import com.google.common.annotations.VisibleForTesting;
 import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tag;
+import io.micrometer.core.instrument.Tags;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
-import io.micrometer.core.instrument.Tag;
+import javax.annotation.Nullable;
-import io.micrometer.core.instrument.Tags;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
@@ -23,14 +24,12 @@ import org.whispersystems.textsecuregcm.util.Futures;
 import org.whispersystems.textsecuregcm.util.Optionals;
 import reactor.core.publisher.Flux;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
-import javax.annotation.Nullable;
 
 public class KeysManager {
   // KeysController for backwards compatibility
   private static final String GET_KEYS_COUNTER_NAME = MetricsUtil.name(KeysManager.class, "getKeys");
 
   private final SingleUseECPreKeyStore ecPreKeys;
-  private final SingleUseKEMPreKeyStore pqPreKeys;
   private final PagedSingleUseKEMPreKeyStore pagedPqPreKeys;
   private final RepeatedUseECSignedPreKeyStore ecSignedPreKeys;
   private final RepeatedUseKEMSignedPreKeyStore pqLastResortKeys;
@@ -39,12 +38,10 @@ public class KeysManager {
 
   public KeysManager(
       final SingleUseECPreKeyStore ecPreKeys,
-      final SingleUseKEMPreKeyStore pqPreKeys,
       final PagedSingleUseKEMPreKeyStore pagedPqPreKeys,
       final RepeatedUseECSignedPreKeyStore ecSignedPreKeys,
       final RepeatedUseKEMSignedPreKeyStore pqLastResortKeys) {
     this.ecPreKeys = ecPreKeys;
-    this.pqPreKeys = pqPreKeys;
     this.pagedPqPreKeys = pagedPqPreKeys;
     this.ecSignedPreKeys = ecSignedPreKeys;
     this.pqLastResortKeys = pqLastResortKeys;
@@ -109,9 +106,7 @@ public class KeysManager {
 
   public CompletableFuture<Void> storeKemOneTimePreKeys(final UUID identifier, final byte deviceId,
       final List<KEMSignedPreKey> preKeys) {
-    // Unconditionally delete keys in old format keystore, then write to the pagedPqPreKeys store
+    return pagedPqPreKeys.store(identifier, deviceId, preKeys);
-    return pqPreKeys.delete(identifier, deviceId)
-        .thenCompose(_ -> pagedPqPreKeys.store(identifier, deviceId, preKeys));
 
   }
 
@@ -124,16 +119,12 @@ public class KeysManager {
   CompletableFuture<Optional<KEMSignedPreKey>> takePQ(final UUID identifier, final byte deviceId) {
     return tagTakePQ(pagedPqPreKeys.take(identifier, deviceId), PQSource.PAGE)
         .thenCompose(maybeSingleUsePreKey -> maybeSingleUsePreKey
-              .map(ignored -> CompletableFuture.completedFuture(maybeSingleUsePreKey))
+            .map(_ -> CompletableFuture.completedFuture(maybeSingleUsePreKey))
-              .orElseGet(() -> tagTakePQ(pqPreKeys.take(identifier, deviceId), PQSource.ROW)))
-        .thenCompose(maybeSingleUsePreKey -> maybeSingleUsePreKey
-            .map(singleUsePreKey -> CompletableFuture.completedFuture(maybeSingleUsePreKey))
             .orElseGet(() -> tagTakePQ(pqLastResortKeys.find(identifier, deviceId), PQSource.LAST_RESORT)));
   }
 
   private enum PQSource {
     PAGE,
-    ROW,
     LAST_RESORT
   }
   private CompletableFuture<Optional<KEMSignedPreKey>> tagTakePQ(CompletableFuture<Optional<KEMSignedPreKey>> prekey, final PQSource source) {
@@ -163,15 +154,12 @@ public class KeysManager {
   }
 
   public CompletableFuture<Integer> getPqCount(final UUID identifier, final byte deviceId) {
-    // We only return the paged prekey count to encourage clients to upload more prekeys if they only have prekeys
-    // stored in the previous key store format
     return pagedPqPreKeys.getCount(identifier, deviceId);
   }
 
   public CompletableFuture<Void> deleteSingleUsePreKeys(final UUID identifier) {
     return CompletableFuture.allOf(
         ecPreKeys.delete(identifier),
-        pqPreKeys.delete(identifier),
         pagedPqPreKeys.delete(identifier)
     );
   }
@@ -179,7 +167,6 @@ public class KeysManager {
   public CompletableFuture<Void> deleteSingleUsePreKeys(final UUID accountUuid, final byte deviceId) {
     return CompletableFuture.allOf(
         ecPreKeys.delete(accountUuid, deviceId),
-        pqPreKeys.delete(accountUuid, deviceId),
         pagedPqPreKeys.delete(accountUuid, deviceId)
     );
   }

service/src/main/java/org/whispersystems/textsecuregcm/storage/PagedSingleUseKEMPreKeyStore.java

@@ -51,11 +51,11 @@ import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.S3Object;
 
 /**
- * @implNote This version of a {@link SingleUsePreKeyStore} store bundles prekeys into "pages", which are stored in on
+ * @implNote This an analog of {@link SingleUseECPreKeyStore} store bundles prekeys into "pages", which are stored in on
  * an object store and referenced via dynamodb. Each device may only have a single active page at a time. Crashes or
  * errors may leave orphaned pages which are no longer referenced by the database. A background process must
  * periodically check for orphaned pages and remove them.
- * @see SingleUsePreKeyStore
+ * @see SingleUseECPreKeyStore
  */
 public class PagedSingleUseKEMPreKeyStore {
 

service/src/main/java/org/whispersystems/textsecuregcm/storage/SingleUseECPreKeyStore.java

@@ -5,30 +5,305 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
+import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+import static org.whispersystems.textsecuregcm.storage.AbstractDynamoDbStore.DYNAMO_DB_MAX_BATCH_SIZE;
+
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.DistributionSummary;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Timer;
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.signal.libsignal.protocol.InvalidKeyException;
 import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
+import org.whispersystems.textsecuregcm.util.Util;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import java.util.Map;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
-import java.util.UUID;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
+import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
+import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
 
-import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+/**
+ * A single-use EC pre-key store stores single-use EC prekeys. Keys returned by a single-use pre-key
+ * store's {@link #take(UUID, byte)} method are guaranteed to be returned exactly once, and repeated calls will never
+ * yield the same key.
+ * <p/>
+ * Each {@link Account} may have one or more {@link Device devices}. Clients <em>should</em> regularly check their
+ * supply of single-use pre-keys (see {@link #getCount(UUID, byte)}) and upload new keys when their supply runs low. In
+ * the event that a party wants to begin a session with a device that has no single-use pre-keys remaining, that party
+ * may fall back to using the device's repeated-use ("last-resort") signed pre-key instead.
+ */
+public class SingleUseECPreKeyStore {
+
+  private final DynamoDbAsyncClient dynamoDbAsyncClient;
+  private final String tableName;
+
+  private final Timer getKeyCountTimer = Metrics.timer(name(getClass(), "getCount"));
+  private final Timer storeKeyTimer = Metrics.timer(name(getClass(), "storeKey"));
+  private final Timer storeKeyBatchTimer = Metrics.timer(name(getClass(), "storeKeyBatch"));
+  private final Timer deleteForDeviceTimer = Metrics.timer(name(getClass(), "deleteForDevice"));
+  private final Timer deleteForAccountTimer = Metrics.timer(name(getClass(), "deleteForAccount"));
+
+  private final Counter noKeyCountAvailableCounter = Metrics.counter(name(getClass(), "noKeyCountAvailable"));
+
+  final DistributionSummary keysConsideredForTakeDistributionSummary = DistributionSummary
+      .builder(name(getClass(), "keysConsideredForTake"))
+      .publishPercentiles(0.5, 0.75, 0.95, 0.99, 0.999)
+      .distributionStatisticExpiry(Duration.ofMinutes(10))
+      .register(Metrics.globalRegistry);
+
+  final DistributionSummary availableKeyCountDistributionSummary = DistributionSummary
+      .builder(name(getClass(), "availableKeyCount"))
+      .publishPercentiles(0.5, 0.75, 0.95, 0.99, 0.999)
+      .distributionStatisticExpiry(Duration.ofMinutes(10))
+      .register(Metrics.globalRegistry);
 
-public class SingleUseECPreKeyStore extends SingleUsePreKeyStore<ECPreKey> {
   private static final String PARSE_BYTE_ARRAY_COUNTER_NAME = name(SingleUseECPreKeyStore.class, "parseByteArray");
 
+  private final String takeKeyTimerName = name(getClass(), "takeKey");
+  private static final String KEY_PRESENT_TAG_NAME = "keyPresent";
+
+  static final String KEY_ACCOUNT_UUID = "U";
+  static final String KEY_DEVICE_ID_KEY_ID = "DK";
+  static final String ATTR_PUBLIC_KEY = "P";
+  static final String ATTR_REMAINING_KEYS = "R";
+
   public SingleUseECPreKeyStore(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) {
-    super(dynamoDbAsyncClient, tableName);
+    this.dynamoDbAsyncClient = dynamoDbAsyncClient;
+    this.tableName = tableName;
   }
 
-  @Override
+  /**
-  protected Map<String, AttributeValue> getItemFromPreKey(final UUID identifier,
+   * Stores a batch of single-use pre-keys for a specific device. All previously-stored keys for the device are cleared
+   * before storing new keys.
+   *
+   * @param identifier the identifier for the account/identity with which the target device is associated
+   * @param deviceId the identifier for the device within the given account/identity
+   * @param preKeys a collection of single-use pre-keys to store for the target device
+   *
+   * @return a future that completes when all previously-stored keys have been removed and the given collection of
+   * pre-keys has been stored in its place
+   */
+  public CompletableFuture<Void> store(final UUID identifier, final byte deviceId, final List<ECPreKey> preKeys) {
+    final Timer.Sample sample = Timer.start();
+
+    return Mono.fromFuture(() -> delete(identifier, deviceId))
+        .thenMany(
+            Flux.fromIterable(preKeys)
+                .sort(Comparator.comparing(preKey -> preKey.keyId()))
+                .zipWith(Flux.range(0, preKeys.size()).map(i -> preKeys.size() - i))
+                .flatMap(preKeyAndRemainingCount -> Mono.fromFuture(() ->
+                        store(identifier, deviceId, preKeyAndRemainingCount.getT1(), preKeyAndRemainingCount.getT2())),
+                    DYNAMO_DB_MAX_BATCH_SIZE))
+        .then()
+        .toFuture()
+        .thenRun(() -> sample.stop(storeKeyBatchTimer));
+  }
+
+  private CompletableFuture<Void> store(final UUID identifier, final byte deviceId, final ECPreKey preKey, final int remainingKeys) {
+    final Timer.Sample sample = Timer.start();
+
+    return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
+            .tableName(tableName)
+            .item(getItemFromPreKey(identifier, deviceId, preKey, remainingKeys))
+            .build())
+        .thenRun(() -> sample.stop(storeKeyTimer));
+  }
+
+  /**
+   * Attempts to retrieve a single-use pre-key for a specific device. Keys may only be returned by this method at most
+   * once; once the key is returned, it is removed from the key store and subsequent calls to this method will never
+   * return the same key.
+   *
+   * @param identifier the identifier for the account/identity with which the target device is associated
+   * @param deviceId the identifier for the device within the given account/identity
+   *
+   * @return a future that yields a single-use pre-key if one is available or empty if no single-use pre-keys are
+   * available for the target device
+   */
+  public CompletableFuture<Optional<ECPreKey>> take(final UUID identifier, final byte deviceId) {
+    final Timer.Sample sample = Timer.start();
+    final AttributeValue partitionKey = getPartitionKey(identifier);
+    final AtomicInteger keysConsidered = new AtomicInteger(0);
+
+    return Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
+                .tableName(tableName)
+                .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
+                .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
+                .expressionAttributeValues(Map.of(
+                    ":uuid", partitionKey,
+                    ":sortprefix", getSortKeyPrefix(deviceId)))
+                .projectionExpression(KEY_DEVICE_ID_KEY_ID)
+                .consistentRead(false)
+                .limit(1)
+                .build())
+            .items())
+        .map(item -> DeleteItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(
+                KEY_ACCOUNT_UUID, partitionKey,
+                KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)))
+            .returnValues(ReturnValue.ALL_OLD)
+            .build())
+        .flatMap(deleteItemRequest -> Mono.fromFuture(() -> dynamoDbAsyncClient.deleteItem(deleteItemRequest)), 1)
+        .doOnNext(deleteItemResponse -> keysConsidered.incrementAndGet())
+        .filter(DeleteItemResponse::hasAttributes)
+        .next()
+        .map(deleteItemResponse -> getPreKeyFromItem(deleteItemResponse.attributes()))
+        .toFuture()
+        .thenApply(Optional::ofNullable)
+        .whenComplete((maybeKey, throwable) -> {
+          sample.stop(Metrics.timer(takeKeyTimerName, KEY_PRESENT_TAG_NAME, String.valueOf(maybeKey != null && maybeKey.isPresent())));
+          keysConsideredForTakeDistributionSummary.record(keysConsidered.get());
+        });
+  }
+
+  /**
+   * Estimates the number of single-use pre-keys available for a given device.
+
+   * @param identifier the identifier for the account/identity with which the target device is associated
+   * @param deviceId the identifier for the device within the given account/identity
+
+   * @return a future that yields the approximate number of single-use pre-keys currently available for the target
+   * device
+   */
+  public CompletableFuture<Integer> getCount(final UUID identifier, final byte deviceId) {
+    final Timer.Sample sample = Timer.start();
+
+    return dynamoDbAsyncClient.query(QueryRequest.builder()
+            .tableName(tableName)
+            .consistentRead(false)
+            .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
+            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
+            .expressionAttributeValues(Map.of(
+                ":uuid", getPartitionKey(identifier),
+                ":sortprefix", getSortKeyPrefix(deviceId)))
+            .projectionExpression(ATTR_REMAINING_KEYS)
+            .limit(1)
+            .build())
+        .thenApply(response -> {
+          if (response.count() > 0) {
+            final Map<String, AttributeValue> item = response.items().getFirst();
+
+            if (item.containsKey(ATTR_REMAINING_KEYS)) {
+              return Integer.parseInt(item.get(ATTR_REMAINING_KEYS).n());
+            } else {
+              // Some legacy keys sets may not have pre-counted keys; in that case, we'll tell the owners of those key
+              // sets that they have none remaining, prompting an upload of a fresh set that we'll pre-count. This has
+              // no effect on consumers of keys, which will still be able to take keys if any are actually present.
+              noKeyCountAvailableCounter.increment();
+              return 0;
+            }
+          } else {
+            return 0;
+          }
+        })
+        .whenComplete((keyCount, throwable) -> {
+          sample.stop(getKeyCountTimer);
+
+          if (throwable == null && keyCount != null) {
+            availableKeyCountDistributionSummary.record(keyCount);
+          }
+        });
+  }
+
+  /**
+   * Removes all single-use pre-keys for all devices associated with the given account/identity.
+   *
+   * @param identifier the identifier for the account/identity for which to remove single-use pre-keys
+   *
+   * @return a future that completes when all single-use pre-keys have been removed for all devices associated with the
+   * given account/identity
+   */
+  public CompletableFuture<Void> delete(final UUID identifier) {
+    final Timer.Sample sample = Timer.start();
+
+    return deleteItems(getPartitionKey(identifier), Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
+            .tableName(tableName)
+            .keyConditionExpression("#uuid = :uuid")
+            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
+            .expressionAttributeValues(Map.of(":uuid", getPartitionKey(identifier)))
+            .projectionExpression(KEY_DEVICE_ID_KEY_ID)
+            .consistentRead(true)
+            .build())
+        .items()))
+        .thenRun(() -> sample.stop(deleteForAccountTimer));
+  }
+
+  /**
+   * Removes all single-use pre-keys for a specific device.
+   *
+   * @param identifier the identifier for the account/identity with which the target device is associated
+   * @param deviceId the identifier for the device within the given account/identity
+
+   * @return a future that completes when all single-use pre-keys have been removed for the target device
+   */
+  public CompletableFuture<Void> delete(final UUID identifier, final byte deviceId) {
+    final Timer.Sample sample = Timer.start();
+
+    return deleteItems(getPartitionKey(identifier), Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
+            .tableName(tableName)
+            .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
+            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
+            .expressionAttributeValues(Map.of(
+                ":uuid", getPartitionKey(identifier),
+                ":sortprefix", getSortKeyPrefix(deviceId)))
+            .projectionExpression(KEY_DEVICE_ID_KEY_ID)
+            .consistentRead(true)
+            .build())
+        .items()))
+        .thenRun(() -> sample.stop(deleteForDeviceTimer));
+  }
+
+  private CompletableFuture<Void> deleteItems(final AttributeValue partitionKey, final Flux<Map<String, AttributeValue>> items) {
+    return items
+        .map(item -> DeleteItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(
+                KEY_ACCOUNT_UUID, partitionKey,
+                KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)
+            ))
+            .build())
+        .flatMap(deleteItemRequest -> Mono.fromFuture(() -> dynamoDbAsyncClient.deleteItem(deleteItemRequest)), DYNAMO_DB_MAX_BATCH_SIZE)
+        .then()
+        .toFuture()
+        .thenRun(Util.NOOP);
+  }
+
+  protected static AttributeValue getPartitionKey(final UUID accountUuid) {
+    return AttributeValues.fromUUID(accountUuid);
+  }
+
+  protected static AttributeValue getSortKey(final byte deviceId, final long keyId) {
+    final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
+    byteBuffer.putLong(deviceId);
+    byteBuffer.putLong(keyId);
+    return AttributeValues.fromByteBuffer(byteBuffer.flip());
+  }
+
+  private static AttributeValue getSortKeyPrefix(final byte deviceId) {
+    final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
+    byteBuffer.putLong(deviceId);
+    return AttributeValues.fromByteBuffer(byteBuffer.flip());
+  }
+
+  private Map<String, AttributeValue> getItemFromPreKey(final UUID identifier,
       final byte deviceId,
       final ECPreKey preKey,
       final int remainingKeys) {
-
     return Map.of(
         KEY_ACCOUNT_UUID, getPartitionKey(identifier),
         KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, preKey.keyId()),
@@ -36,8 +311,7 @@ public class SingleUseECPreKeyStore extends SingleUsePreKeyStore<ECPreKey> {
         ATTR_REMAINING_KEYS, AttributeValues.fromInt(remainingKeys));
   }
 
-  @Override
+  private ECPreKey getPreKeyFromItem(final Map<String, AttributeValue> item) {
-  protected ECPreKey getPreKeyFromItem(final Map<String, AttributeValue> item) {
     final long keyId = item.get(KEY_DEVICE_ID_KEY_ID).b().asByteBuffer().getLong(8);
     final byte[] publicKey = AttributeValues.extractByteArray(item.get(ATTR_PUBLIC_KEY), PARSE_BYTE_ARRAY_COUNTER_NAME);
 

service/src/main/java/org/whispersystems/textsecuregcm/storage/SingleUseKEMPreKeyStore.java

@@ -1,50 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import org.signal.libsignal.protocol.InvalidKeyException;
-import org.signal.libsignal.protocol.kem.KEMPublicKey;
-import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
-import org.whispersystems.textsecuregcm.util.AttributeValues;
-import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
-import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import java.util.Map;
-import java.util.UUID;
-
-public class SingleUseKEMPreKeyStore extends SingleUsePreKeyStore<KEMSignedPreKey> {
-
-  public SingleUseKEMPreKeyStore(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) {
-    super(dynamoDbAsyncClient, tableName);
-  }
-
-  @Override
-  protected Map<String, AttributeValue> getItemFromPreKey(final UUID identifier,
-      final byte deviceId,
-      final KEMSignedPreKey signedPreKey,
-      final int remainingKeys) {
-
-    return Map.of(
-        KEY_ACCOUNT_UUID, getPartitionKey(identifier),
-        KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, signedPreKey.keyId()),
-        ATTR_PUBLIC_KEY, AttributeValues.fromByteArray(signedPreKey.serializedPublicKey()),
-        ATTR_SIGNATURE, AttributeValues.fromByteArray(signedPreKey.signature()),
-        ATTR_REMAINING_KEYS, AttributeValues.fromInt(remainingKeys));
-  }
-
-  @Override
-  protected KEMSignedPreKey getPreKeyFromItem(final Map<String, AttributeValue> item) {
-    final long keyId = item.get(KEY_DEVICE_ID_KEY_ID).b().asByteBuffer().getLong(8);
-    final byte[] publicKey = item.get(ATTR_PUBLIC_KEY).b().asByteArray();
-    final byte[] signature = item.get(ATTR_SIGNATURE).b().asByteArray();
-
-    try {
-      return new KEMSignedPreKey(keyId, new KEMPublicKey(publicKey), signature);
-    } catch (final InvalidKeyException e) {
-      // This should never happen since we're serializing keys directly from `KEMPublicKey` instances on the way in
-      throw new IllegalArgumentException(e);
-    }
-  }
-}

service/src/main/java/org/whispersystems/textsecuregcm/storage/SingleUsePreKeyStore.java

@@ -1,306 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
-import static org.whispersystems.textsecuregcm.storage.AbstractDynamoDbStore.DYNAMO_DB_MAX_BATCH_SIZE;
-
-import io.micrometer.core.instrument.Counter;
-import io.micrometer.core.instrument.DistributionSummary;
-import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Timer;
-import java.nio.ByteBuffer;
-import java.time.Duration;
-import java.util.Comparator;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.UUID;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicInteger;
-import org.whispersystems.textsecuregcm.entities.PreKey;
-import org.whispersystems.textsecuregcm.util.AttributeValues;
-import org.whispersystems.textsecuregcm.util.Util;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
-import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
-import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
-import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
-
-/**
- * A single-use pre-key store stores single-use pre-keys of a specific type. Keys returned by a single-use pre-key
- * store's {@link #take(UUID, byte)} method are guaranteed to be returned exactly once, and repeated calls will never
- * yield the same key.
- * <p/>
- * Each {@link Account} may have one or more {@link Device devices}. Clients <em>should</em> regularly check their
- * supply of single-use pre-keys (see {@link #getCount(UUID, byte)}) and upload new keys when their supply runs low. In
- * the event that a party wants to begin a session with a device that has no single-use pre-keys remaining, that party
- * may fall back to using the device's repeated-use ("last-resort") signed pre-key instead.
- */
-public abstract class SingleUsePreKeyStore<K extends PreKey<?>> {
-
-  private final DynamoDbAsyncClient dynamoDbAsyncClient;
-  private final String tableName;
-
-  private final Timer getKeyCountTimer = Metrics.timer(name(getClass(), "getCount"));
-  private final Timer storeKeyTimer = Metrics.timer(name(getClass(), "storeKey"));
-  private final Timer storeKeyBatchTimer = Metrics.timer(name(getClass(), "storeKeyBatch"));
-  private final Timer deleteForDeviceTimer = Metrics.timer(name(getClass(), "deleteForDevice"));
-  private final Timer deleteForAccountTimer = Metrics.timer(name(getClass(), "deleteForAccount"));
-
-  private final Counter noKeyCountAvailableCounter = Metrics.counter(name(getClass(), "noKeyCountAvailable"));
-
-  final DistributionSummary keysConsideredForTakeDistributionSummary = DistributionSummary
-      .builder(name(getClass(), "keysConsideredForTake"))
-      .publishPercentiles(0.5, 0.75, 0.95, 0.99, 0.999)
-      .distributionStatisticExpiry(Duration.ofMinutes(10))
-      .register(Metrics.globalRegistry);
-
-  final DistributionSummary availableKeyCountDistributionSummary = DistributionSummary
-      .builder(name(getClass(), "availableKeyCount"))
-      .publishPercentiles(0.5, 0.75, 0.95, 0.99, 0.999)
-      .distributionStatisticExpiry(Duration.ofMinutes(10))
-      .register(Metrics.globalRegistry);
-
-  private final String takeKeyTimerName = name(getClass(), "takeKey");
-  private static final String KEY_PRESENT_TAG_NAME = "keyPresent";
-
-  static final String KEY_ACCOUNT_UUID = "U";
-  static final String KEY_DEVICE_ID_KEY_ID = "DK";
-  static final String ATTR_PUBLIC_KEY = "P";
-  static final String ATTR_SIGNATURE = "S";
-  static final String ATTR_REMAINING_KEYS = "R";
-
-  protected SingleUsePreKeyStore(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) {
-    this.dynamoDbAsyncClient = dynamoDbAsyncClient;
-    this.tableName = tableName;
-  }
-
-  /**
-   * Stores a batch of single-use pre-keys for a specific device. All previously-stored keys for the device are cleared
-   * before storing new keys.
-   *
-   * @param identifier the identifier for the account/identity with which the target device is associated
-   * @param deviceId the identifier for the device within the given account/identity
-   * @param preKeys a collection of single-use pre-keys to store for the target device
-   *
-   * @return a future that completes when all previously-stored keys have been removed and the given collection of
-   * pre-keys has been stored in its place
-   */
-  public CompletableFuture<Void> store(final UUID identifier, final byte deviceId, final List<K> preKeys) {
-    final Timer.Sample sample = Timer.start();
-
-    return Mono.fromFuture(() -> delete(identifier, deviceId))
-        .thenMany(
-            Flux.fromIterable(preKeys)
-                .sort(Comparator.comparing(preKey -> preKey.keyId()))
-                .zipWith(Flux.range(0, preKeys.size()).map(i -> preKeys.size() - i))
-                .flatMap(preKeyAndRemainingCount -> Mono.fromFuture(() ->
-                        store(identifier, deviceId, preKeyAndRemainingCount.getT1(), preKeyAndRemainingCount.getT2())),
-                    DYNAMO_DB_MAX_BATCH_SIZE))
-        .then()
-        .toFuture()
-        .thenRun(() -> sample.stop(storeKeyBatchTimer));
-  }
-
-  private CompletableFuture<Void> store(final UUID identifier, final byte deviceId, final K preKey, final int remainingKeys) {
-    final Timer.Sample sample = Timer.start();
-
-    return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
-            .tableName(tableName)
-            .item(getItemFromPreKey(identifier, deviceId, preKey, remainingKeys))
-            .build())
-        .thenRun(() -> sample.stop(storeKeyTimer));
-  }
-
-  /**
-   * Attempts to retrieve a single-use pre-key for a specific device. Keys may only be returned by this method at most
-   * once; once the key is returned, it is removed from the key store and subsequent calls to this method will never
-   * return the same key.
-   *
-   * @param identifier the identifier for the account/identity with which the target device is associated
-   * @param deviceId the identifier for the device within the given account/identity
-   *
-   * @return a future that yields a single-use pre-key if one is available or empty if no single-use pre-keys are
-   * available for the target device
-   */
-  public CompletableFuture<Optional<K>> take(final UUID identifier, final byte deviceId) {
-    final Timer.Sample sample = Timer.start();
-    final AttributeValue partitionKey = getPartitionKey(identifier);
-    final AtomicInteger keysConsidered = new AtomicInteger(0);
-
-    return Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
-                .tableName(tableName)
-                .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
-                .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
-                .expressionAttributeValues(Map.of(
-                    ":uuid", partitionKey,
-                    ":sortprefix", getSortKeyPrefix(deviceId)))
-                .projectionExpression(KEY_DEVICE_ID_KEY_ID)
-                .consistentRead(false)
-                .limit(1)
-                .build())
-            .items())
-        .map(item -> DeleteItemRequest.builder()
-            .tableName(tableName)
-            .key(Map.of(
-                KEY_ACCOUNT_UUID, partitionKey,
-                KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)))
-            .returnValues(ReturnValue.ALL_OLD)
-            .build())
-        .flatMap(deleteItemRequest -> Mono.fromFuture(() -> dynamoDbAsyncClient.deleteItem(deleteItemRequest)), 1)
-        .doOnNext(deleteItemResponse -> keysConsidered.incrementAndGet())
-        .filter(DeleteItemResponse::hasAttributes)
-        .next()
-        .map(deleteItemResponse -> getPreKeyFromItem(deleteItemResponse.attributes()))
-        .toFuture()
-        .thenApply(Optional::ofNullable)
-        .whenComplete((maybeKey, throwable) -> {
-          sample.stop(Metrics.timer(takeKeyTimerName, KEY_PRESENT_TAG_NAME, String.valueOf(maybeKey != null && maybeKey.isPresent())));
-          keysConsideredForTakeDistributionSummary.record(keysConsidered.get());
-        });
-  }
-
-  /**
-   * Estimates the number of single-use pre-keys available for a given device.
-
-   * @param identifier the identifier for the account/identity with which the target device is associated
-   * @param deviceId the identifier for the device within the given account/identity
-
-   * @return a future that yields the approximate number of single-use pre-keys currently available for the target
-   * device
-   */
-  public CompletableFuture<Integer> getCount(final UUID identifier, final byte deviceId) {
-    final Timer.Sample sample = Timer.start();
-
-    return dynamoDbAsyncClient.query(QueryRequest.builder()
-            .tableName(tableName)
-            .consistentRead(false)
-            .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
-            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
-            .expressionAttributeValues(Map.of(
-                ":uuid", getPartitionKey(identifier),
-                ":sortprefix", getSortKeyPrefix(deviceId)))
-            .projectionExpression(ATTR_REMAINING_KEYS)
-            .limit(1)
-            .build())
-        .thenApply(response -> {
-          if (response.count() > 0) {
-            final Map<String, AttributeValue> item = response.items().getFirst();
-
-            if (item.containsKey(ATTR_REMAINING_KEYS)) {
-              return Integer.parseInt(item.get(ATTR_REMAINING_KEYS).n());
-            } else {
-              // Some legacy keys sets may not have pre-counted keys; in that case, we'll tell the owners of those key
-              // sets that they have none remaining, prompting an upload of a fresh set that we'll pre-count. This has
-              // no effect on consumers of keys, which will still be able to take keys if any are actually present.
-              noKeyCountAvailableCounter.increment();
-              return 0;
-            }
-          } else {
-            return 0;
-          }
-        })
-        .whenComplete((keyCount, throwable) -> {
-          sample.stop(getKeyCountTimer);
-
-          if (throwable == null && keyCount != null) {
-            availableKeyCountDistributionSummary.record(keyCount);
-          }
-        });
-  }
-
-  /**
-   * Removes all single-use pre-keys for all devices associated with the given account/identity.
-   *
-   * @param identifier the identifier for the account/identity for which to remove single-use pre-keys
-   *
-   * @return a future that completes when all single-use pre-keys have been removed for all devices associated with the
-   * given account/identity
-   */
-  public CompletableFuture<Void> delete(final UUID identifier) {
-    final Timer.Sample sample = Timer.start();
-
-    return deleteItems(getPartitionKey(identifier), Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
-            .tableName(tableName)
-            .keyConditionExpression("#uuid = :uuid")
-            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
-            .expressionAttributeValues(Map.of(":uuid", getPartitionKey(identifier)))
-            .projectionExpression(KEY_DEVICE_ID_KEY_ID)
-            .consistentRead(true)
-            .build())
-        .items()))
-        .thenRun(() -> sample.stop(deleteForAccountTimer));
-  }
-
-  /**
-   * Removes all single-use pre-keys for a specific device.
-   *
-   * @param identifier the identifier for the account/identity with which the target device is associated
-   * @param deviceId the identifier for the device within the given account/identity
-
-   * @return a future that completes when all single-use pre-keys have been removed for the target device
-   */
-  public CompletableFuture<Void> delete(final UUID identifier, final byte deviceId) {
-    final Timer.Sample sample = Timer.start();
-
-    return deleteItems(getPartitionKey(identifier), Flux.from(dynamoDbAsyncClient.queryPaginator(QueryRequest.builder()
-            .tableName(tableName)
-            .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
-            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
-            .expressionAttributeValues(Map.of(
-                ":uuid", getPartitionKey(identifier),
-                ":sortprefix", getSortKeyPrefix(deviceId)))
-            .projectionExpression(KEY_DEVICE_ID_KEY_ID)
-            .consistentRead(true)
-            .build())
-        .items()))
-        .thenRun(() -> sample.stop(deleteForDeviceTimer));
-  }
-
-  private CompletableFuture<Void> deleteItems(final AttributeValue partitionKey, final Flux<Map<String, AttributeValue>> items) {
-    return items
-        .map(item -> DeleteItemRequest.builder()
-            .tableName(tableName)
-            .key(Map.of(
-                KEY_ACCOUNT_UUID, partitionKey,
-                KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)
-            ))
-            .build())
-        .flatMap(deleteItemRequest -> Mono.fromFuture(() -> dynamoDbAsyncClient.deleteItem(deleteItemRequest)), DYNAMO_DB_MAX_BATCH_SIZE)
-        .then()
-        .toFuture()
-        .thenRun(Util.NOOP);
-  }
-
-  protected static AttributeValue getPartitionKey(final UUID accountUuid) {
-    return AttributeValues.fromUUID(accountUuid);
-  }
-
-  protected static AttributeValue getSortKey(final byte deviceId, final long keyId) {
-    final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
-    byteBuffer.putLong(deviceId);
-    byteBuffer.putLong(keyId);
-    return AttributeValues.fromByteBuffer(byteBuffer.flip());
-  }
-
-  private static AttributeValue getSortKeyPrefix(final byte deviceId) {
-    final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
-    byteBuffer.putLong(deviceId);
-    return AttributeValues.fromByteBuffer(byteBuffer.flip());
-  }
-
-  protected abstract Map<String, AttributeValue> getItemFromPreKey(final UUID identifier,
-      final byte deviceId,
-      final K preKey,
-      final int remainingKeys);
-
-  protected abstract K getPreKeyFromItem(final Map<String, AttributeValue> item);
-}

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

@@ -72,7 +72,6 @@ import org.whispersystems.textsecuregcm.storage.RepeatedUseKEMSignedPreKeyStore;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
 import org.whispersystems.textsecuregcm.storage.SingleUseECPreKeyStore;
-import org.whispersystems.textsecuregcm.storage.SingleUseKEMPreKeyStore;
 import org.whispersystems.textsecuregcm.storage.SubscriptionManager;
 import org.whispersystems.textsecuregcm.storage.Subscriptions;
 import org.whispersystems.textsecuregcm.subscriptions.AppleAppStoreClient;
@@ -238,7 +237,6 @@ public record CommandDependencies(
         configuration.getPagedSingleUseKEMPreKeyStore().bucket());
     KeysManager keys = new KeysManager(
         new SingleUseECPreKeyStore(dynamoDbAsyncClient, configuration.getDynamoDbTables().getEcKeys().getTableName()),
-        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, configuration.getDynamoDbTables().getKemKeys().getTableName()),
         pagedSingleUseKEMPreKeyStore,
         new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient,
             configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName()),

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,
+        SingleUseECPreKeyStore.KEY_ACCOUNT_UUID,
-        SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID,
+        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)
+                .attributeName(SingleUseECPreKeyStore.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)
                 .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
+  private ECPreKey generatePreKey(final long keyId) {
-  protected SingleUsePreKeyStore<ECPreKey> getPreKeyStore() {
-    return preKeyStore;
-  }
-
-  @Override
-  protected ECPreKey generatePreKey(final long keyId) {
     return new ECPreKey(keyId, ECKeyPair.generate().getPublicKey());
   }
 
-  @Override
+  private void clearKeyCountAttributes() {
-  protected 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),
+                SingleUseECPreKeyStore.KEY_ACCOUNT_UUID, item.get(SingleUseECPreKeyStore.KEY_ACCOUNT_UUID),
-                SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID, item.get(SingleUsePreKeyStore.KEY_DEVICE_ID_KEY_ID)))
+                SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID, item.get(SingleUseECPreKeyStore.KEY_DEVICE_ID_KEY_ID)))
-            .updateExpression("REMOVE " + SingleUsePreKeyStore.ATTR_REMAINING_KEYS)
+            .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();
-  }
-}

service/src/test/resources/config/test.yml

@@ -108,8 +108,6 @@ dynamoDbTables:
     tableName: keys_test
   ecSignedPreKeys:
     tableName: repeated_use_signed_ec_pre_keys_test
-  pqKeys:
-    tableName: pq_keys_test
   pagedPqKeys:
     tableName: paged_pq_keys_test
   pqLastResortKeys: