Store signed EC pre-keys in a dedicated table

service/config/sample.yml

@@ -98,6 +98,8 @@ dynamoDbTables:
     generator: abcdefg12345678= # random base64-encoded binary sequence
   ecKeys:
     tableName: Example_Keys
+  ecSignedPreKeys:
+    tableName: Example_EC_Signed_Pre_Keys
   pqKeys:
     tableName: Example_PQ_Keys
   pqLastResortKeys:

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

@@ -316,6 +316,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     KeysManager keys = new KeysManager(
             dynamoDbAsyncClient,
         config.getDynamoDbTables().getEcKeys().getTableName(),
+        config.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         config.getDynamoDbTables().getKemKeys().getTableName(),
         config.getDynamoDbTables().getKemLastResortKeys().getTableName());
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

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

@@ -51,6 +51,7 @@ public class DynamoDbTables {
   private final Table deletedAccountsLock;
   private final IssuedReceiptsTableConfiguration issuedReceipts;
   private final Table ecKeys;
+  private final Table ecSignedPreKeys;
   private final Table kemKeys;
   private final Table kemLastResortKeys;
   private final TableWithExpiration messages;
@@ -72,6 +73,7 @@ public class DynamoDbTables {
       @JsonProperty("deletedAccountsLock") final Table deletedAccountsLock,
       @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("messages") final TableWithExpiration messages,
@@ -92,6 +94,7 @@ public class DynamoDbTables {
     this.deletedAccountsLock = deletedAccountsLock;
     this.issuedReceipts = issuedReceipts;
     this.ecKeys = ecKeys;
+    this.ecSignedPreKeys = ecSignedPreKeys;
     this.kemKeys = kemKeys;
     this.kemLastResortKeys = kemLastResortKeys;
     this.messages = messages;
@@ -138,6 +141,12 @@ public class DynamoDbTables {
     return ecKeys;
   }
 
+  @NotNull
+  @Valid
+  public Table getEcSignedPreKeys() {
+    return ecSignedPreKeys;
+  }
+
   @NotNull
   @Valid
   public Table getKemKeys() {

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

@@ -341,7 +341,9 @@ public class DeviceController {
       keys.delete(a.getPhoneNumberIdentifier(), device.getId());
 
       maybeDeviceActivationRequest.ifPresent(deviceActivationRequest -> {
+        keys.storeEcSignedPreKeys(a.getUuid(), Map.of(device.getId(), deviceActivationRequest.aciSignedPreKey().get()));
         keys.storePqLastResort(a.getUuid(), Map.of(device.getId(), deviceActivationRequest.aciPqLastResortPreKey().get()));
+        keys.storeEcSignedPreKeys(a.getPhoneNumberIdentifier(), Map.of(device.getId(), deviceActivationRequest.pniSignedPreKey().get()));
         keys.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(device.getId(), deviceActivationRequest.pniPqLastResortPreKey().get()));
       });
 

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

@@ -167,7 +167,7 @@ public class KeysController {
 
     keys.store(
         getIdentifier(account, identityType), device.getId(),
-        preKeys.getPreKeys(), preKeys.getPqPreKeys(), preKeys.getPqLastResortPreKey());
+        preKeys.getPreKeys(), preKeys.getPqPreKeys(), preKeys.getSignedPreKey(), preKeys.getPqLastResortPreKey());
   }
 
   @Timed

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

@@ -176,7 +176,9 @@ public class RegistrationController {
         registrationRequest.deviceActivationRequest().gcmToken().ifPresent(gcmRegistrationId ->
             device.setGcmId(gcmRegistrationId.gcmRegistrationId()));
 
+        keysManager.storeEcSignedPreKeys(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciSignedPreKey().get()));
         keysManager.storePqLastResort(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciPqLastResortPreKey().get()));
+        keysManager.storeEcSignedPreKeys(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniSignedPreKey().get()));
         keysManager.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniPqLastResortPreKey().get()));
       });
     }

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

@@ -312,7 +312,10 @@ public class AccountsManager {
 
       numberChangedAccount = updateWithRetries(
           account,
-          a -> { setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds); return true; },
+          a -> {
+            setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds);
+            return true;
+          },
           a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier),
           () -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
           AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
@@ -322,6 +325,8 @@ public class AccountsManager {
       keysManager.delete(phoneNumberIdentifier);
       keysManager.delete(originalPhoneNumberIdentifier);
 
+      keysManager.storeEcSignedPreKeys(phoneNumberIdentifier, pniSignedPreKeys);
+
       if (pniPqLastResortPreKeys != null) {
         keysManager.storePqLastResort(
             phoneNumberIdentifier,
@@ -362,6 +367,7 @@ public class AccountsManager {
 
     final List<Long> pqEnabledDeviceIDs = keysManager.getPqEnabledDevices(pni);
     keysManager.delete(pni);
+    keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys);
     if (pniPqLastResortPreKeys != null) {
       keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
     }

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

@@ -14,6 +14,7 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import javax.annotation.Nullable;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 
@@ -21,26 +22,30 @@ public class KeysManager {
 
   private final SingleUseECPreKeyStore ecPreKeys;
   private final SingleUseKEMPreKeyStore pqPreKeys;
+  private final RepeatedUseECSignedPreKeyStore ecSignedPreKeys;
   private final RepeatedUseKEMSignedPreKeyStore pqLastResortKeys;
 
   public KeysManager(
       final DynamoDbAsyncClient dynamoDbAsyncClient,
       final String ecTableName,
       final String pqTableName,
+      final String ecSignedPreKeysTableName,
       final String pqLastResortTableName) {
     this.ecPreKeys = new SingleUseECPreKeyStore(dynamoDbAsyncClient, ecTableName);
     this.pqPreKeys = new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, pqTableName);
+    this.ecSignedPreKeys = new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient, ecSignedPreKeysTableName);
     this.pqLastResortKeys = new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient, pqLastResortTableName);
   }
 
   public void store(final UUID identifier, final long deviceId, final List<ECPreKey> keys) {
-    store(identifier, deviceId, keys, null, null);
+    store(identifier, deviceId, keys, null, null, null);
   }
 
   public void store(
       final UUID identifier, final long deviceId,
       @Nullable final List<ECPreKey> ecKeys,
       @Nullable final List<KEMSignedPreKey> pqKeys,
+      @Nullable final ECSignedPreKey ecSignedPreKey,
       @Nullable final KEMSignedPreKey pqLastResortKey) {
 
     final List<CompletableFuture<Void>> storeFutures = new ArrayList<>();
@@ -53,6 +58,10 @@ public class KeysManager {
       storeFutures.add(pqPreKeys.store(identifier, deviceId, pqKeys));
     }
 
+    if (ecSignedPreKey != null) {
+      storeFutures.add(ecSignedPreKeys.store(identifier, deviceId, ecSignedPreKey));
+    }
+
     if (pqLastResortKey != null) {
       storeFutures.add(pqLastResortKeys.store(identifier, deviceId, pqLastResortKey));
     }
@@ -60,6 +69,10 @@ public class KeysManager {
     CompletableFuture.allOf(storeFutures.toArray(new CompletableFuture[0])).join();
   }
 
+  public void storeEcSignedPreKeys(final UUID identifier, final Map<Long, ECSignedPreKey> keys) {
+    ecSignedPreKeys.store(identifier, keys).join();
+  }
+
   public void storePqLastResort(final UUID identifier, final Map<Long, KEMSignedPreKey> keys) {
     pqLastResortKeys.store(identifier, keys).join();
   }
@@ -80,6 +93,10 @@ public class KeysManager {
     return pqLastResortKeys.find(identifier, deviceId).join();
   }
 
+  public CompletableFuture<Optional<ECSignedPreKey>> getEcSignedPreKey(final UUID identifier, final long deviceId) {
+    return ecSignedPreKeys.find(identifier, deviceId);
+  }
+
   public List<Long> getPqEnabledDevices(final UUID identifier) {
     return pqLastResortKeys.getDeviceIdsWithKeys(identifier).collectList().block();
   }
@@ -96,6 +113,7 @@ public class KeysManager {
     CompletableFuture.allOf(
             ecPreKeys.delete(accountUuid),
             pqPreKeys.delete(accountUuid),
+            ecSignedPreKeys.delete(accountUuid),
             pqLastResortKeys.delete(accountUuid))
         .join();
   }
@@ -104,6 +122,7 @@ public class KeysManager {
     CompletableFuture.allOf(
             ecPreKeys.delete(accountUuid, deviceId),
             pqPreKeys.delete(accountUuid, deviceId),
+            ecSignedPreKeys.delete(accountUuid, deviceId),
             pqLastResortKeys.delete(accountUuid, deviceId))
         .join();
   }

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

@@ -165,6 +165,7 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
     KeysManager keys = new KeysManager(
             dynamoDbAsyncClient,
         configuration.getDynamoDbTables().getEcKeys().getTableName(),
+        configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         configuration.getDynamoDbTables().getKemKeys().getTableName(),
         configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

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

@@ -149,6 +149,7 @@ record CommandDependencies(
     KeysManager keys = new KeysManager(
             dynamoDbAsyncClient,
         configuration.getDynamoDbTables().getEcKeys().getTableName(),
+        configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         configuration.getDynamoDbTables().getKemKeys().getTableName(),
         configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

service/src/test/java/org/whispersystems/textsecuregcm/controllers/RegistrationControllerTest.java

@@ -671,6 +671,8 @@ class RegistrationControllerTest {
     verify(device).setSignedPreKey(expectedAciSignedPreKey);
     verify(device).setPhoneNumberIdentitySignedPreKey(expectedPniSignedPreKey);
 
+    verify(keysManager).storeEcSignedPreKeys(accountIdentifier, Map.of(Device.MASTER_ID, expectedAciSignedPreKey));
+    verify(keysManager).storeEcSignedPreKeys(phoneNumberIdentifier, Map.of(Device.MASTER_ID, expectedPniSignedPreKey));
     verify(keysManager).storePqLastResort(accountIdentifier, Map.of(Device.MASTER_ID, expectedAciPqLastResortPreKey));
     verify(keysManager).storePqLastResort(phoneNumberIdentifier, Map.of(Device.MASTER_ID, expectedPniPqLastResortPreKey));
 

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

@@ -709,6 +709,7 @@ class AccountsManagerTest {
     verify(keysManager).delete(originalPni);
     verify(keysManager, atLeastOnce()).delete(targetPni);
     verify(keysManager).delete(newPni);
+    verify(keysManager).storeEcSignedPreKeys(eq(newPni), any());
     verifyNoMoreInteractions(keysManager);
   }
 
@@ -748,6 +749,7 @@ class AccountsManagerTest {
     verify(keysManager).delete(newPni);
     verify(keysManager).delete(originalPni);
     verify(keysManager).getPqEnabledDevices(uuid);
+    verify(keysManager).storeEcSignedPreKeys(newPni, newSignedKeys);
     verify(keysManager).storePqLastResort(eq(newPni), eq(Map.of(1L, newSignedPqKeys.get(1L))));
     verifyNoMoreInteractions(keysManager);
   }
@@ -852,6 +854,7 @@ class AccountsManagerTest {
     verifyNoInteractions(deletedAccounts);
 
     verify(keysManager).delete(oldPni);
+    verify(keysManager).storeEcSignedPreKeys(oldPni, newSignedKeys);
 
     // only the pq key for the already-pq-enabled device should be saved
     verify(keysManager).storePqLastResort(eq(oldPni), eq(Map.of(1L, newSignedPqKeys.get(1L))));

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

@@ -115,7 +115,21 @@ public final class DynamoDbExtensionSchema {
                 .build()),
         List.of(), List.of()),
 
-    REPEATED_USE_SIGNED_PRE_KEYS("repeated_use_signed_pre_keys_test",
+    REPEATED_USE_EC_SIGNED_PRE_KEYS("repeated_use_signed_ec_pre_keys_test",
+        RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID,
+        RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID,
+        List.of(
+            AttributeDefinition.builder()
+                .attributeName(RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID)
+                .attributeType(ScalarAttributeType.B)
+                .build(),
+            AttributeDefinition.builder()
+                .attributeName(RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID)
+                .attributeType(ScalarAttributeType.N)
+                .build()),
+        List.of(), List.of()),
+
+    REPEATED_USE_KEM_SIGNED_PRE_KEYS("repeated_use_signed_kem_pre_keys_test",
         RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID,
         RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID,
         List.of(

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

@@ -32,7 +32,7 @@ class KeysManagerTest {
 
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
-      Tables.EC_KEYS, Tables.PQ_KEYS, Tables.REPEATED_USE_SIGNED_PRE_KEYS);
+      Tables.EC_KEYS, Tables.PQ_KEYS, Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
 
   private static final UUID ACCOUNT_UUID = UUID.randomUUID();
   private static final long DEVICE_ID = 1L;
@@ -45,7 +45,8 @@ class KeysManagerTest {
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
         Tables.EC_KEYS.tableName(),
         Tables.PQ_KEYS.tableName(),
-        Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
+        Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
+        Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName());
   }
 
   @Test
@@ -64,25 +65,25 @@ class KeysManagerTest {
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Repeatedly storing same key should have no effect");
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(generateTestKEMSignedPreKey(1)), null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(generateTestKEMSignedPreKey(1)), null, null);
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Uploading new PQ prekeys should have no effect on EC prekeys");
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, null, generateTestKEMSignedPreKey(1001));
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, null, null, generateTestKEMSignedPreKey(1001));
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Uploading new PQ last-resort prekey should have no effect on EC prekeys");
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
         "Uploading new PQ last-resort prekey should have no effect on one-time PQ prekeys");
     assertEquals(1001, keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).get().keyId());
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(2)), null, null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(2)), null, null, null);
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Inserting a new key should overwrite all prior keys of the same type for the given account/device");
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
         "Uploading new EC prekeys should have no effect on PQ prekeys");
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(3)), List.of(generateTestKEMSignedPreKey(2)), null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(3)), List.of(generateTestKEMSignedPreKey(2)), null, null);
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Inserting a new key should overwrite all prior keys of the same type for the given account/device");
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
@@ -90,7 +91,7 @@ class KeysManagerTest {
 
     keysManager.store(ACCOUNT_UUID, DEVICE_ID,
         List.of(generateTestPreKey(4), generateTestPreKey(5)),
-        List.of(generateTestKEMSignedPreKey(6), generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(1002));
+        List.of(generateTestKEMSignedPreKey(6), generateTestKEMSignedPreKey(7)), null, generateTestKEMSignedPreKey(1002));
     assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
         "Inserting multiple new keys should overwrite all prior keys for the given account/device");
     assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
@@ -119,7 +120,7 @@ class KeysManagerTest {
     final KEMSignedPreKey preKey2 = generateTestKEMSignedPreKey(2);
     final KEMSignedPreKey preKeyLast = generateTestKEMSignedPreKey(1001);
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(preKey1, preKey2), preKeyLast);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(preKey1, preKey2), null, preKeyLast);
 
     assertEquals(Optional.of(preKey1), keysManager.takePQ(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
@@ -139,7 +140,7 @@ class KeysManagerTest {
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(1)), List.of(generateTestKEMSignedPreKey(1)), null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(1)), List.of(generateTestKEMSignedPreKey(1)), null, null);
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
   }
@@ -148,53 +149,69 @@ class KeysManagerTest {
   void testDeleteByAccount() {
     keysManager.store(ACCOUNT_UUID, DEVICE_ID,
         List.of(generateTestPreKey(1), generateTestPreKey(2)),
-        List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)), generateTestKEMSignedPreKey(5));
+        List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)),
+        generateTestECSignedPreKey(5),
+        generateTestKEMSignedPreKey(6));
 
     keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1,
-        List.of(generateTestPreKey(6)),
+        List.of(generateTestPreKey(7)),
-        List.of(generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(8));
+        List.of(generateTestKEMSignedPreKey(8)),
+        generateTestECSignedPreKey(9),
+        generateTestKEMSignedPreKey(10));
 
     assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
+    assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
+    assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
 
     keysManager.delete(ACCOUNT_UUID);
 
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
+    assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
     assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
+    assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
   }
 
   @Test
   void testDeleteByAccountAndDevice() {
     keysManager.store(ACCOUNT_UUID, DEVICE_ID,
-        List.of(generateTestPreKey(1), generateTestPreKey(2)),
+        List.of(generateTestPreKey(1),generateTestPreKey(2)),
-        List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)), generateTestKEMSignedPreKey(5));
+        List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)),
+        generateTestECSignedPreKey(5),
+        generateTestKEMSignedPreKey(6));
 
     keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1,
-        List.of(generateTestPreKey(6)),
+        List.of(generateTestPreKey(7)),
-        List.of(generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(8));
+        List.of(generateTestKEMSignedPreKey(8)),
+        generateTestECSignedPreKey(9),
+        generateTestKEMSignedPreKey(10));
 
     assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
+    assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
+    assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
 
     keysManager.delete(ACCOUNT_UUID, DEVICE_ID);
 
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
+    assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
     assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
     assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
     assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
+    assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
   }
 
@@ -225,10 +242,10 @@ class KeysManagerTest {
   void testGetPqEnabledDevices() {
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
 
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(KeysHelper.signedKEMPreKey(1, identityKeyPair)), null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(KeysHelper.signedKEMPreKey(1, identityKeyPair)), null, null);
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1, null, null, KeysHelper.signedKEMPreKey(2, identityKeyPair));
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1, null, null, null, KeysHelper.signedKEMPreKey(2, identityKeyPair));
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 2, null, List.of(KeysHelper.signedKEMPreKey(3, identityKeyPair)), KeysHelper.signedKEMPreKey(4, identityKeyPair));
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 2, null, List.of(KeysHelper.signedKEMPreKey(3, identityKeyPair)), null, KeysHelper.signedKEMPreKey(4, identityKeyPair));
-    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 3, null, null, null);
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID + 3, null, null, null, null);
     assertIterableEquals(
         Set.of(DEVICE_ID + 1, DEVICE_ID + 2),
         Set.copyOf(keysManager.getPqEnabledDevices(ACCOUNT_UUID)));

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

@@ -20,14 +20,14 @@ class RepeatedUseECSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTes
 
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION =
-      new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS);
+      new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS);
 
   private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
 
   @BeforeEach
   void setUp() {
     keyStore = new RepeatedUseECSignedPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
+        DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName());
   }
 
   @Override

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

@@ -12,8 +12,6 @@ import org.signal.libsignal.protocol.ecc.ECKeyPair;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
 
-import static org.junit.jupiter.api.Assertions.*;
-
 class RepeatedUseKEMSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTest<KEMSignedPreKey> {
 
   private RepeatedUseKEMSignedPreKeyStore keyStore;
@@ -22,14 +20,14 @@ class RepeatedUseKEMSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTe
 
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION =
-      new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS);
+      new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
 
   private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
 
   @BeforeEach
   void setUp() {
     keyStore = new RepeatedUseKEMSignedPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
+        DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName());
   }
 
   @Override

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java

@@ -315,6 +315,8 @@ class DeviceControllerTest {
     verify(pendingDevicesManager).remove(AuthHelper.VALID_NUMBER);
     verify(messagesManager).clear(eq(AuthHelper.VALID_UUID), eq(42L));
     verify(clientPresenceManager).disconnectPresence(AuthHelper.VALID_UUID, Device.MASTER_ID);
+    verify(keysManager).storeEcSignedPreKeys(AuthHelper.VALID_UUID, Map.of(response.getDeviceId(), aciSignedPreKey.get()));
+    verify(keysManager).storeEcSignedPreKeys(AuthHelper.VALID_PNI, Map.of(response.getDeviceId(), pniSignedPreKey.get()));
     verify(keysManager).storePqLastResort(AuthHelper.VALID_UUID, Map.of(response.getDeviceId(), aciPqLastResortPreKey.get()));
     verify(keysManager).storePqLastResort(AuthHelper.VALID_PNI, Map.of(response.getDeviceId(), pniPqLastResortPreKey.get()));
   }

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java

@@ -687,7 +687,7 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
-    verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), listCaptor.capture(), isNull(), isNull());
+    verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
 
     assertThat(listCaptor.getValue()).containsExactly(preKey);
 
@@ -718,7 +718,7 @@ class KeysControllerTest {
 
     ArgumentCaptor<List<ECPreKey>> ecCaptor = ArgumentCaptor.forClass(List.class);
     ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
-    verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(pqLastResortPreKey));
+    verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(signedPreKey), eq(pqLastResortPreKey));
 
     assertThat(ecCaptor.getValue()).containsExactly(preKey);
     assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);
@@ -820,7 +820,7 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
-    verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), listCaptor.capture(), isNull(), isNull());
+    verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
 
     assertThat(listCaptor.getValue()).containsExactly(preKey);
 
@@ -852,7 +852,7 @@ class KeysControllerTest {
 
     ArgumentCaptor<List<ECPreKey>> ecCaptor = ArgumentCaptor.forClass(List.class);
     ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
-    verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(pqLastResortPreKey));
+    verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(signedPreKey), eq(pqLastResortPreKey));
 
     assertThat(ecCaptor.getValue()).containsExactly(preKey);
     assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);
@@ -896,7 +896,7 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
-    verify(KEYS).store(eq(AuthHelper.DISABLED_UUID), eq(1L), listCaptor.capture(), isNull(), isNull());
+    verify(KEYS).store(eq(AuthHelper.DISABLED_UUID), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
 
     List<ECPreKey> capturedList = listCaptor.getValue();
     assertThat(capturedList.size()).isEqualTo(1);