Add a command to copy signed pre-keys from `Account` records to their own table

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

@@ -198,6 +198,7 @@ import org.whispersystems.textsecuregcm.workers.CheckDynamicConfigurationCommand
 import org.whispersystems.textsecuregcm.workers.CrawlAccountsCommand;
 import org.whispersystems.textsecuregcm.workers.DeleteUserCommand;
 import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand;
+import org.whispersystems.textsecuregcm.workers.MigrateSignedECPreKeysCommand;
 import org.whispersystems.textsecuregcm.workers.ScheduledApnPushNotificationSenderServiceCommand;
 import org.whispersystems.textsecuregcm.workers.ServerVersionCommand;
 import org.whispersystems.textsecuregcm.workers.SetRequestLoggingEnabledTask;
@@ -252,6 +253,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     bootstrap.addCommand(new CrawlAccountsCommand());
     bootstrap.addCommand(new ScheduledApnPushNotificationSenderServiceCommand());
     bootstrap.addCommand(new MessagePersisterServiceCommand());
+    bootstrap.addCommand(new MigrateSignedECPreKeysCommand());
   }
 
   @Override
@@ -318,7 +320,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         config.getDynamoDbTables().getEcKeys().getTableName(),
         config.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         config.getDynamoDbTables().getKemKeys().getTableName(),
-        config.getDynamoDbTables().getKemLastResortKeys().getTableName());
+        config.getDynamoDbTables().getKemLastResortKeys().getTableName(),
+        dynamicConfigurationManager);
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         config.getDynamoDbTables().getMessages().getTableName(),
         config.getDynamoDbTables().getMessages().getExpiration(),

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java

@@ -55,6 +55,10 @@ public class DynamicConfiguration {
   @Valid
   DynamicRateLimitPolicy rateLimitPolicy = new DynamicRateLimitPolicy(false);
 
+  @JsonProperty
+  @Valid
+  DynamicECPreKeyMigrationConfiguration ecPreKeyMigration = new DynamicECPreKeyMigrationConfiguration(false);
+
   public Optional<DynamicExperimentEnrollmentConfiguration> getExperimentEnrollmentConfiguration(
       final String experimentName) {
     return Optional.ofNullable(experiments.get(experimentName));
@@ -97,4 +101,7 @@ public class DynamicConfiguration {
     return rateLimitPolicy;
   }
 
+  public DynamicECPreKeyMigrationConfiguration getEcPreKeyMigrationConfiguration() {
+    return ecPreKeyMigration;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicECPreKeyMigrationConfiguration.java

@@ -0,0 +1,9 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.configuration.dynamic;
+
+public record DynamicECPreKeyMigrationConfiguration(boolean storeEcSignedPreKeys) {
+}

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

@@ -13,6 +13,7 @@ import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import javax.annotation.Nullable;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
@@ -20,6 +21,8 @@ import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 
 public class KeysManager {
 
+  private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
+
   private final SingleUseECPreKeyStore ecPreKeys;
   private final SingleUseKEMPreKeyStore pqPreKeys;
   private final RepeatedUseECSignedPreKeyStore ecSignedPreKeys;
@@ -30,11 +33,13 @@ public class KeysManager {
       final String ecTableName,
       final String pqTableName,
       final String ecSignedPreKeysTableName,
-      final String pqLastResortTableName) {
+      final String pqLastResortTableName,
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager) {
     this.ecPreKeys = new SingleUseECPreKeyStore(dynamoDbAsyncClient, ecTableName);
     this.pqPreKeys = new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, pqTableName);
     this.ecSignedPreKeys = new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient, ecSignedPreKeysTableName);
     this.pqLastResortKeys = new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient, pqLastResortTableName);
+    this.dynamicConfigurationManager = dynamicConfigurationManager;
   }
 
   public void store(final UUID identifier, final long deviceId, final List<ECPreKey> keys) {
@@ -58,7 +63,7 @@ public class KeysManager {
       storeFutures.add(pqPreKeys.store(identifier, deviceId, pqKeys));
     }
 
-    if (ecSignedPreKey != null) {
+    if (ecSignedPreKey != null && dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys()) {
       storeFutures.add(ecSignedPreKeys.store(identifier, deviceId, ecSignedPreKey));
     }
 
@@ -70,7 +75,13 @@ public class KeysManager {
   }
 
   public void storeEcSignedPreKeys(final UUID identifier, final Map<Long, ECSignedPreKey> keys) {
-    ecSignedPreKeys.store(identifier, keys).join();
+    if (dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys()) {
+      ecSignedPreKeys.store(identifier, keys).join();
+    }
+  }
+
+  public CompletableFuture<Boolean> storeEcSignedPreKeyIfAbsent(final UUID identifier, final long deviceId, final ECSignedPreKey signedPreKey) {
+    return ecSignedPreKeys.storeIfAbsent(identifier, deviceId, signedPreKey);
   }
 
   public void storePqLastResort(final UUID identifier, final Map<Long, KEMSignedPreKey> keys) {

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

@@ -167,7 +167,8 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
         configuration.getDynamoDbTables().getEcKeys().getTableName(),
         configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         configuration.getDynamoDbTables().getKemKeys().getTableName(),
-        configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
+        configuration.getDynamoDbTables().getKemLastResortKeys().getTableName(),
+        dynamicConfigurationManager);
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         configuration.getDynamoDbTables().getMessages().getTableName(),
         configuration.getDynamoDbTables().getMessages().getExpiration(),

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

@@ -151,7 +151,8 @@ record CommandDependencies(
         configuration.getDynamoDbTables().getEcKeys().getTableName(),
         configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
         configuration.getDynamoDbTables().getKemKeys().getTableName(),
-        configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
+        configuration.getDynamoDbTables().getKemLastResortKeys().getTableName(),
+        dynamicConfigurationManager);
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         configuration.getDynamoDbTables().getMessages().getTableName(),
         configuration.getDynamoDbTables().getMessages().getExpiration(),

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

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import io.micrometer.core.instrument.Metrics;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple3;
+import reactor.util.function.Tuples;
+
+public class MigrateSignedECPreKeysCommand extends AbstractSinglePassCrawlAccountsCommand {
+
+  private static final String STORE_KEY_ATTEMPT_COUNTER_NAME =
+      MetricsUtil.name(MigrateSignedECPreKeysCommand.class, "storeKeyAttempt");
+
+  public MigrateSignedECPreKeysCommand() {
+    super("migrate-signed-ec-pre-keys", "Migrate signed EC pre-keys from Account records to a dedicated table");
+  }
+
+  @Override
+  protected void crawlAccounts(final Flux<Account> accounts) {
+    final KeysManager keysManager = getCommandDependencies().keysManager();
+
+    accounts.flatMap(account -> Flux.fromIterable(account.getDevices())
+            .flatMap(device -> {
+              final List<Tuple3<UUID, Long, ECSignedPreKey>> keys = new ArrayList<>(2);
+
+              if (device.getSignedPreKey() != null) {
+                keys.add(Tuples.of(account.getUuid(), device.getId(), device.getSignedPreKey()));
+              }
+
+              if (device.getPhoneNumberIdentitySignedPreKey() != null) {
+                keys.add(Tuples.of(account.getPhoneNumberIdentifier(), device.getId(), device.getPhoneNumberIdentitySignedPreKey()));
+              }
+
+              return Flux.fromIterable(keys);
+            }))
+        .flatMap(keyTuple -> Mono.fromFuture(
+            keysManager.storeEcSignedPreKeyIfAbsent(keyTuple.getT1(), keyTuple.getT2(), keyTuple.getT3())))
+        .doOnNext(keyStored -> Metrics.counter(STORE_KEY_ATTEMPT_COUNTER_NAME, "stored", String.valueOf(keyStored)).increment())
+        .blockLast();
+  }
+}

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

@@ -9,6 +9,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.util.List;
 import java.util.Map;
@@ -20,6 +22,8 @@ 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.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicECPreKeyMigrationConfiguration;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
@@ -28,6 +32,7 @@ import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
 
 class KeysManagerTest {
 
+  private DynamicECPreKeyMigrationConfiguration ecPreKeyMigrationConfiguration;
   private KeysManager keysManager;
 
   @RegisterExtension
@@ -41,12 +46,21 @@ class KeysManagerTest {
 
   @BeforeEach
   void setup() {
+    final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
+    final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
+    ecPreKeyMigrationConfiguration = mock(DynamicECPreKeyMigrationConfiguration.class);
+
+    when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
+    when(dynamicConfiguration.getEcPreKeyMigrationConfiguration()).thenReturn(ecPreKeyMigrationConfiguration);
+    when(ecPreKeyMigrationConfiguration.storeEcSignedPreKeys()).thenReturn(true);
+
     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());
+        Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName(),
+        dynamicConfigurationManager);
   }
 
   @Test
@@ -251,6 +265,24 @@ class KeysManagerTest {
         Set.copyOf(keysManager.getPqEnabledDevices(ACCOUNT_UUID)));
   }
 
+  @Test
+  void testStoreEcSignedPreKeyDisabled() {
+    when(ecPreKeyMigrationConfiguration.storeEcSignedPreKeys()).thenReturn(false);
+
+    final ECKeyPair identityKeyPair = Curve.generateKeyPair();
+
+    keysManager.store(ACCOUNT_UUID, DEVICE_ID,
+        List.of(generateTestPreKey(1)),
+        List.of(KeysHelper.signedKEMPreKey(2, identityKeyPair)),
+        KeysHelper.signedECPreKey(3, identityKeyPair),
+        KeysHelper.signedKEMPreKey(4, identityKeyPair));
+
+    assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
+    assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
+    assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
+    assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
+  }
+
   private static ECPreKey generateTestPreKey(final long keyId) {
     return new ECPreKey(keyId, Curve.generateKeyPair().getPublicKey());
   }