From 2d154eb0cfdf662aa0cbace02f934f89258a18d4 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Tue, 20 Jun 2023 13:48:16 -0400 Subject: [PATCH] Add a command to copy signed pre-keys from `Account` records to their own table --- .../textsecuregcm/WhisperServerService.java | 5 +- .../dynamic/DynamicConfiguration.java | 7 +++ ...DynamicECPreKeyMigrationConfiguration.java | 9 ++++ .../textsecuregcm/storage/KeysManager.java | 17 ++++-- .../workers/AssignUsernameCommand.java | 3 +- .../workers/CommandDependencies.java | 3 +- .../MigrateSignedECPreKeysCommand.java | 53 +++++++++++++++++++ .../storage/KeysManagerTest.java | 34 +++++++++++- 8 files changed, 124 insertions(+), 7 deletions(-) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicECPreKeyMigrationConfiguration.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/MigrateSignedECPreKeysCommand.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index c4c094842..fd0177b00 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/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 getExperimentEnrollmentConfiguration( final String experimentName) { return Optional.ofNullable(experiments.get(experimentName)); @@ -97,4 +101,7 @@ public class DynamicConfiguration { return rateLimitPolicy; } + public DynamicECPreKeyMigrationConfiguration getEcPreKeyMigrationConfiguration() { + return ecPreKeyMigration; + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicECPreKeyMigrationConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicECPreKeyMigrationConfiguration.java new file mode 100644 index 000000000..2836d2397 --- /dev/null +++ b/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) { +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java index a767711cc..b9e68a86d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java +++ b/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 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 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 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 keys) { - ecSignedPreKeys.store(identifier, keys).join(); + if (dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys()) { + ecSignedPreKeys.store(identifier, keys).join(); + } + } + + public CompletableFuture storeEcSignedPreKeyIfAbsent(final UUID identifier, final long deviceId, final ECSignedPreKey signedPreKey) { + return ecSignedPreKeys.storeIfAbsent(identifier, deviceId, signedPreKey); } public void storePqLastResort(final UUID identifier, final Map keys) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java index 7285487a6..06829405e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java @@ -167,7 +167,8 @@ public class AssignUsernameCommand extends EnvironmentCommand accounts) { + final KeysManager keysManager = getCommandDependencies().keysManager(); + + accounts.flatMap(account -> Flux.fromIterable(account.getDevices()) + .flatMap(device -> { + final List> 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(); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java index 45078d65c..cb848e97c 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java +++ b/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 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()); }