Add paged prekey store

2026-04-21 23:18:05 +01:00 · 2025-05-20 10:47:45 -05:00
parent 6d8701665e
commit 2bb14892af
23 changed files with 1125 additions and 54 deletions
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java
@@ -50,6 +50,7 @@ import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;

 public class AccountCreationDeletionIntegrationTest {

@@ -71,6 +72,9 @@ public class AccountCreationDeletionIntegrationTest {
  @RegisterExtension
  static final RedisClusterExtension CACHE_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();

+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension("testbucket");
+
  private static final Clock CLOCK = Clock.fixed(Instant.now(), ZoneId.systemDefault());

  private ScheduledExecutorService executor;
@@ -90,13 +94,18 @@ public class AccountCreationDeletionIntegrationTest {
    final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
    when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);

+    final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
    keysManager = new KeysManager(
-        DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        DynamoDbExtensionSchema.Tables.EC_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
-    );
+        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(),
+            S3_EXTENSION.getBucketName()),
+        new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
+        new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));

    final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
        DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java
@@ -44,6 +44,7 @@ import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
 import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;

 class AccountsManagerChangeNumberIntegrationTest {

@@ -65,6 +66,9 @@ class AccountsManagerChangeNumberIntegrationTest {
  @RegisterExtension
  static final RedisClusterExtension CACHE_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();

+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension("testbucket");
+
  private KeysManager keysManager;
  private DisconnectionRequestManager disconnectionRequestManager;
  private ScheduledExecutorService executor;
@@ -81,13 +85,18 @@ class AccountsManagerChangeNumberIntegrationTest {
      DynamicConfiguration dynamicConfiguration = new DynamicConfiguration();
      when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);

+      final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
      keysManager = new KeysManager(
-          DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-          Tables.EC_KEYS.tableName(),
-          Tables.PQ_KEYS.tableName(),
-          Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
-          Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
-      );
+          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(),
+              S3_EXTENSION.getBucketName()),
+          new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient,
+              DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
+          new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient,
+              DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));

      final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
          DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java
@@ -47,6 +47,7 @@ import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil;
 import org.whispersystems.textsecuregcm.util.TestRandomUtil;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
@@ -78,6 +79,9 @@ class AccountsManagerUsernameIntegrationTest {
  @RegisterExtension
  static RedisClusterExtension CACHE_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();

+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension("testbucket");
+
  private AccountsManager accountsManager;
  private Accounts accounts;

@@ -94,13 +98,18 @@ class AccountsManagerUsernameIntegrationTest {
    DynamicConfiguration dynamicConfiguration = new DynamicConfiguration();
    when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);

+    final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
    final KeysManager keysManager = new KeysManager(
-        DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        Tables.EC_KEYS.tableName(),
-        Tables.PQ_KEYS.tableName(),
-        Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
-        Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
-    );
+        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(),
+            S3_EXTENSION.getBucketName()),
+        new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
+        new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));

    accounts = Mockito.spy(new Accounts(
        Clock.systemUTC(),
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java
@@ -45,6 +45,7 @@ import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.TestClock;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;

 public class AddRemoveDeviceIntegrationTest {

@@ -70,6 +71,9 @@ public class AddRemoveDeviceIntegrationTest {
  @RegisterExtension
  static final RedisServerExtension PUBSUB_SERVER_EXTENSION = RedisServerExtension.builder().build();

+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension("testbucket");
+
  private ExecutorService accountLockExecutor;
  private ScheduledExecutorService messagePollExecutor;

@@ -89,13 +93,18 @@ public class AddRemoveDeviceIntegrationTest {

    clock = TestClock.pinned(Instant.now());

+    final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
    keysManager = new KeysManager(
-        DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        DynamoDbExtensionSchema.Tables.EC_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.PQ_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
-    );
+        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(),
+            S3_EXTENSION.getBucketName()),
+        new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
+        new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient,
+            DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));

    final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
        DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java
@@ -143,6 +143,20 @@ public final class DynamoDbExtensionSchema {
                .build()),
        List.of(), List.of()),

+    PAGED_PQ_KEYS("paged_pq_keys_test",
+        PagedSingleUseKEMPreKeyStore.KEY_ACCOUNT_UUID,
+        PagedSingleUseKEMPreKeyStore.KEY_DEVICE_ID,
+        List.of(
+            AttributeDefinition.builder()
+                .attributeName(PagedSingleUseKEMPreKeyStore.KEY_ACCOUNT_UUID)
+                .attributeType(ScalarAttributeType.B)
+                .build(),
+            AttributeDefinition.builder()
+                .attributeName(PagedSingleUseKEMPreKeyStore.KEY_DEVICE_ID)
+                .attributeType(ScalarAttributeType.N)
+                .build()),
+        List.of(), List.of()),
+
    PUSH_NOTIFICATION_EXPERIMENT_SAMPLES("push_notification_experiment_samples_test",
        PushNotificationExperimentSamples.KEY_EXPERIMENT_NAME,
        PushNotificationExperimentSamples.ATTR_ACI_AND_DEVICE_ID,
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KEMPreKeyPageTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KEMPreKeyPageTest.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.storage;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.signal.libsignal.protocol.InvalidKeyException;
+import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.protocol.ecc.ECKeyPair;
+import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
+import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+
+class KEMPreKeyPageTest {
+
+  private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
+
+  @Test
+  void serializeSinglePreKey() {
+    final ByteBuffer page = KEMPreKeyPage.serialize(KEMPreKeyPage.FORMAT, List.of(generatePreKey(5)));
+    final int actualMagic = page.getInt();
+    assertEquals(KEMPreKeyPage.HEADER_MAGIC, actualMagic);
+    final int version = page.getInt();
+    assertEquals(version, 1);
+    assertEquals(KEMPreKeyPage.SERIALIZED_PREKEY_LENGTH, page.remaining());
+  }
+
+  @Test
+  void emptyPreKeys() {
+    assertThrows(IllegalArgumentException.class, () -> KEMPreKeyPage.serialize(KEMPreKeyPage.FORMAT, Collections.emptyList()));
+  }
+
+  @Test
+  void roundTripSingleton() throws InvalidKeyException {
+    final KEMSignedPreKey preKey = generatePreKey(5);
+    final ByteBuffer buffer = KEMPreKeyPage.serialize(KEMPreKeyPage.FORMAT, List.of(preKey));
+    final long serializedLength = buffer.remaining();
+    assertEquals(KEMPreKeyPage.HEADER_SIZE + KEMPreKeyPage.SERIALIZED_PREKEY_LENGTH, serializedLength);
+
+    final KEMPreKeyPage.KeyLocation keyLocation = KEMPreKeyPage.keyLocation(1, 0);
+    assertEquals(KEMPreKeyPage.HEADER_SIZE, keyLocation.getStartInclusive());
+    assertEquals(serializedLength, KEMPreKeyPage.HEADER_SIZE + keyLocation.length());
+
+    buffer.position(keyLocation.getStartInclusive());
+    final KEMSignedPreKey deserializedPreKey = KEMPreKeyPage.deserializeKey(1, buffer);
+
+    assertEquals(5L, deserializedPreKey.keyId());
+    assertEquals(preKey, deserializedPreKey);
+  }
+
+  @Test
+  void roundTripMultiple() throws InvalidKeyException {
+    final List<KEMSignedPreKey> keys = Arrays.asList(generatePreKey(1), generatePreKey(2), generatePreKey(5));
+    final ByteBuffer page = KEMPreKeyPage.serialize(KEMPreKeyPage.FORMAT, keys);
+
+    assertEquals(KEMPreKeyPage.HEADER_SIZE + KEMPreKeyPage.SERIALIZED_PREKEY_LENGTH * 3, page.remaining());
+
+    for (int i = 0; i < keys.size(); i++) {
+      final KEMPreKeyPage.KeyLocation keyLocation = KEMPreKeyPage.keyLocation(1, i);
+      assertEquals(
+          KEMPreKeyPage.HEADER_SIZE + KEMPreKeyPage.SERIALIZED_PREKEY_LENGTH * i,
+          keyLocation.getStartInclusive());
+      final ByteBuffer buf = page.slice(keyLocation.getStartInclusive(), keyLocation.length());
+      final KEMSignedPreKey actual = KEMPreKeyPage.deserializeKey(1, buf);
+      assertEquals(keys.get(i), actual);
+    }
+  }
+
+  @Test
+  void wrongFormat() {
+    assertThrows(IllegalArgumentException.class, () ->
+        KEMPreKeyPage.deserializeKey(2,
+            ByteBuffer.allocate(KEMPreKeyPage.HEADER_SIZE + KEMPreKeyPage.SERIALIZED_PREKEY_LENGTH)));
+  }
+
+  @Test
+  void wrongSize() {
+    assertThrows(IllegalArgumentException.class, () -> KEMPreKeyPage.deserializeKey(1, ByteBuffer.allocate(100)));
+  }
+
+
+  @Test
+  void negativeKeyId() throws InvalidKeyException {
+    final KEMSignedPreKey preKey = generatePreKey(-1);
+    ByteBuffer page = KEMPreKeyPage.serialize(KEMPreKeyPage.FORMAT, List.of(preKey));
+    page.position(KEMPreKeyPage.HEADER_SIZE);
+    KEMSignedPreKey deserializedPreKey = KEMPreKeyPage.deserializeKey(1, page);
+    assertEquals(-1L, deserializedPreKey.keyId());
+  }
+
+  private static KEMSignedPreKey generatePreKey(long keyId) {
+    return KeysHelper.signedKEMPreKey((int) keyId, IDENTITY_KEY_PAIR);
+  }
+
+}
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java
@@ -22,6 +22,7 @@ import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;

 class KeysManagerTest {

@@ -31,6 +32,9 @@ class KeysManagerTest {
  static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
      Tables.EC_KEYS, Tables.PQ_KEYS, Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);

+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension("testbucket");
+
  private static final UUID ACCOUNT_UUID = UUID.randomUUID();
  private static final byte DEVICE_ID = 1;

@@ -38,13 +42,16 @@ class KeysManagerTest {

  @BeforeEach
  void setup() {
+    final DynamoDbAsyncClient dynamoDbAsyncClient = DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient();
    keysManager = new KeysManager(
-        DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        Tables.EC_KEYS.tableName(),
-        Tables.PQ_KEYS.tableName(),
-        Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
-        Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
-    );
+        new SingleUseECPreKeyStore(dynamoDbAsyncClient, Tables.EC_KEYS.tableName()),
+        new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, Tables.PQ_KEYS.tableName()),
+        new PagedSingleUseKEMPreKeyStore(dynamoDbAsyncClient,
+            S3_EXTENSION.getS3Client(),
+            DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),
+            S3_EXTENSION.getBucketName()),
+        new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient, Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName()),
+        new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()));
  }

  @Test
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/PagedSingleUseKEMPreKeyStoreTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/PagedSingleUseKEMPreKeyStoreTest.java
@@ -0,0 +1,218 @@
+/*
+ * 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 static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
+
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.RepeatedTest;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.protocol.ecc.ECKeyPair;
+import org.testcontainers.containers.localstack.LocalStackContainer;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.utility.DockerImageName;
+import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
+import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+import reactor.core.publisher.Flux;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
+import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
+import software.amazon.awssdk.services.s3.model.S3Object;
+
+class PagedSingleUseKEMPreKeyStoreTest {
+
+  private static final int KEY_COUNT = 100;
+  private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
+  private static final String BUCKET_NAME = "testbucket";
+
+  private PagedSingleUseKEMPreKeyStore keyStore;
+
+  @RegisterExtension
+  static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
+      DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS);
+
+  @RegisterExtension
+  static final S3LocalStackExtension S3_EXTENSION = new S3LocalStackExtension(BUCKET_NAME);
+
+  @BeforeEach
+  void setUp() {
+    keyStore = new PagedSingleUseKEMPreKeyStore(
+        DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
+        S3_EXTENSION.getS3Client(),
+        DynamoDbExtensionSchema.Tables.PAGED_PQ_KEYS.tableName(),
+        BUCKET_NAME);
+  }
+
+  @Test
+  void storeTake() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(Optional.empty(), keyStore.take(accountIdentifier, deviceId).join());
+
+    final List<KEMSignedPreKey> preKeys = generateRandomPreKeys();
+    assertDoesNotThrow(() -> keyStore.store(accountIdentifier, deviceId, preKeys).join());
+
+    final List<KEMSignedPreKey> sortedPreKeys = preKeys.stream()
+        .sorted(Comparator.comparing(preKey -> preKey.keyId()))
+        .toList();
+
+    assertEquals(Optional.of(sortedPreKeys.get(0)), keyStore.take(accountIdentifier, deviceId).join());
+    assertEquals(Optional.of(sortedPreKeys.get(1)), keyStore.take(accountIdentifier, deviceId).join());
+  }
+
+  @Test
+  void storeTwice() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    final List<KEMSignedPreKey> preKeys1 = generateRandomPreKeys();
+    keyStore.store(accountIdentifier, deviceId, preKeys1).join();
+    List<String> oldPages = listPages(accountIdentifier).stream().map(S3Object::key).collect(Collectors.toList());
+    assertEquals(1, oldPages.size());
+
+    final List<KEMSignedPreKey> preKeys2 = generateRandomPreKeys();
+    keyStore.store(accountIdentifier, deviceId, preKeys2).join();
+    List<String> newPages = listPages(accountIdentifier).stream().map(S3Object::key).collect(Collectors.toList());
+    assertEquals(1, newPages.size());
+
+    assertNotEquals(oldPages.getFirst(), newPages.getFirst());
+
+    assertEquals(
+        preKeys2.stream().sorted(Comparator.comparing(preKey -> preKey.keyId())).toList(),
+
+        IntStream.range(0, preKeys2.size())
+            .mapToObj(i -> keyStore.take(accountIdentifier, deviceId).join())
+            .map(Optional::orElseThrow)
+            .toList());
+
+    assertTrue(keyStore.take(accountIdentifier, deviceId).join().isEmpty());
+
+  }
+
+  @Test
+  void takeAll() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    final List<KEMSignedPreKey> preKeys = generateRandomPreKeys();
+    assertDoesNotThrow(() -> keyStore.store(accountIdentifier, deviceId, preKeys).join());
+
+    final List<KEMSignedPreKey> sortedPreKeys = preKeys.stream()
+        .sorted(Comparator.comparing(preKey -> preKey.keyId()))
+        .toList();
+
+    for (int i = 0; i < KEY_COUNT; i++) {
+      assertEquals(Optional.of(sortedPreKeys.get(i)), keyStore.take(accountIdentifier, deviceId).join());
+    }
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+    assertTrue(keyStore.take(accountIdentifier, deviceId).join().isEmpty());
+  }
+
+  @Test
+  void getCount() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+
+    final List<KEMSignedPreKey> preKeys = generateRandomPreKeys();
+
+    keyStore.store(accountIdentifier, deviceId, preKeys).join();
+
+    assertEquals(KEY_COUNT, keyStore.getCount(accountIdentifier, deviceId).join());
+
+    for (int i = 0; i < KEY_COUNT; i++) {
+      keyStore.take(accountIdentifier, deviceId).join();
+      assertEquals(KEY_COUNT - (i + 1), keyStore.getCount(accountIdentifier, deviceId).join());
+    }
+  }
+
+  @Test
+  void deleteSingleDevice() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+    assertDoesNotThrow(() -> keyStore.delete(accountIdentifier, deviceId).join());
+
+    final List<KEMSignedPreKey> preKeys = generateRandomPreKeys();
+
+    keyStore.store(accountIdentifier, deviceId, preKeys).join();
+    keyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
+
+    assertDoesNotThrow(() -> keyStore.delete(accountIdentifier, deviceId).join());
+
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+    assertEquals(KEY_COUNT, keyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
+
+    final List<S3Object> pages = listPages(accountIdentifier);
+    assertEquals(1, pages.size());
+    assertTrue(pages.get(0).key().startsWith("%s/%s".formatted(accountIdentifier, deviceId + 1)));
+  }
+
+  @Test
+  void deleteAllDevices() {
+    final UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceId = 1;
+
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+    assertDoesNotThrow(() -> keyStore.delete(accountIdentifier).join());
+
+    final List<KEMSignedPreKey> preKeys = generateRandomPreKeys();
+
+    keyStore.store(accountIdentifier, deviceId, preKeys).join();
+    keyStore.store(accountIdentifier, (byte) (deviceId + 1), preKeys).join();
+
+    assertDoesNotThrow(() -> keyStore.delete(accountIdentifier).join());
+
+    assertEquals(0, keyStore.getCount(accountIdentifier, deviceId).join());
+    assertEquals(0, keyStore.getCount(accountIdentifier, (byte) (deviceId + 1)).join());
+    assertEquals(0, listPages(accountIdentifier).size());
+  }
+
+  private List<S3Object> listPages(final UUID identifier) {
+    return Flux.from(S3_EXTENSION.getS3Client().listObjectsV2Paginator(ListObjectsV2Request.builder()
+            .bucket(BUCKET_NAME)
+            .prefix(identifier.toString())
+            .build()))
+        .concatMap(response -> Flux.fromIterable(response.contents()))
+        .collectList()
+        .block();
+  }
+
+  private List<KEMSignedPreKey> 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(keyId -> KeysHelper.signedKEMPreKey(keyId, IDENTITY_KEY_PAIR))
+        .toList();
+  }
+}
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/S3LocalStackExtension.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/S3LocalStackExtension.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2021-2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
+
+import java.util.Objects;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.testcontainers.containers.localstack.LocalStackContainer;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.utility.DockerImageName;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
+import software.amazon.awssdk.services.s3.model.DeleteBucketRequest;
+import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
+import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
+
+@Testcontainers
+public class S3LocalStackExtension implements BeforeEachCallback, AfterEachCallback, BeforeAllCallback,
+    AfterAllCallback {
+
+  private final static DockerImageName LOCAL_STACK_IMAGE =
+      DockerImageName.parse(Objects.requireNonNull(
+          System.getProperty("localstackImage"),
+          "Local stack image not found; must provide localstackImage system property"));
+
+  private static LocalStackContainer LOCAL_STACK = new LocalStackContainer(LOCAL_STACK_IMAGE).withServices(S3);
+
+  private final String bucketName;
+  private S3AsyncClient s3Client;
+
+  public S3LocalStackExtension(final String bucketName) {
+    this.bucketName = bucketName;
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context) {
+    Flux.from(s3Client.listObjectsV2Paginator(ListObjectsV2Request.builder()
+                .bucket(bucketName)
+                .build())
+            .contents())
+        .flatMap(obj -> Mono.fromFuture(() -> s3Client.deleteObject(DeleteObjectRequest.builder()
+            .bucket(bucketName)
+            .key(obj.key())
+            .build())), 100)
+        .then()
+        .block();
+    s3Client.deleteBucket(DeleteBucketRequest.builder().bucket(bucketName).build()).join();
+  }
+
+
+  @Override
+  public void beforeEach(ExtensionContext context) throws Exception {
+    s3Client.createBucket(CreateBucketRequest.builder().bucket(bucketName).build()).join();
+  }
+
+  public S3AsyncClient getS3Client() {
+    return s3Client;
+  }
+
+  @Override
+  public void afterAll(final ExtensionContext context) throws Exception {
+    s3Client.close();
+    LOCAL_STACK.close();
+  }
+
+  @Override
+  public void beforeAll(final ExtensionContext context) throws Exception {
+    LOCAL_STACK.start();
+    s3Client = S3AsyncClient.builder()
+        .endpointOverride(LOCAL_STACK.getEndpoint())
+        .credentialsProvider(StaticCredentialsProvider
+            .create(AwsBasicCredentials.create(LOCAL_STACK.getAccessKey(), LOCAL_STACK.getSecretKey())))
+        .region(Region.of(LOCAL_STACK.getRegion()))
+        .build();
+  }
+
+  public String getBucketName() {
+    return bucketName;
+  }
+}