diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index c1c3e9cb9..6a47e6716 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -344,8 +344,8 @@ public class WhisperServerService extends Application d.setPhoneNumberIdentitySignedPreKey(setKeysRequest.signedPreKey()); } }, - d -> List.of(keys.buildWriteItemForEcSignedPreKey(identifier, d.getId(), setKeysRequest.signedPreKey()))) + d -> keys.buildWriteItemForEcSignedPreKey(identifier, d.getId(), setKeysRequest.signedPreKey()) + .map(List::of) + .orElseGet(Collections::emptyList)) .toCompletableFuture(); } else { updateAccountFuture = CompletableFuture.completedFuture(account); @@ -297,7 +299,9 @@ public class KeysController { case PNI -> d.setPhoneNumberIdentitySignedPreKey(signedPreKey); } }, - d -> List.of(keys.buildWriteItemForEcSignedPreKey(identifier, d.getId(), signedPreKey))) + d -> keys.buildWriteItemForEcSignedPreKey(identifier, d.getId(), signedPreKey) + .map(List::of) + .orElseGet(Collections::emptyList)) .toCompletableFuture() .thenApply(Util.ASYNC_EMPTY_RESPONSE); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index cc37a9190..fbe3105a6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -430,7 +430,8 @@ public class AccountsManager { if (pniSignedPreKeys != null) { pniSignedPreKeys.forEach((deviceId, signedPreKey) -> - keyWriteItems.add(keysManager.buildWriteItemForEcSignedPreKey(phoneNumberIdentifier, deviceId, signedPreKey))); + keysManager.buildWriteItemForEcSignedPreKey(phoneNumberIdentifier, deviceId, signedPreKey) + .ifPresent(keyWriteItems::add)); } if (pniPqLastResortPreKeys != null) { 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 bcc387b06..2cb9fed12 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java @@ -11,6 +11,7 @@ import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.CompletableFuture; +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; @@ -19,6 +20,8 @@ import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; public class KeysManager { + private final DynamicConfigurationManager dynamicConfigurationManager; + private final SingleUseECPreKeyStore ecPreKeys; private final SingleUseKEMPreKeyStore pqPreKeys; private final RepeatedUseECSignedPreKeyStore ecSignedPreKeys; @@ -29,18 +32,22 @@ 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 TransactWriteItem buildWriteItemForEcSignedPreKey(final UUID identifier, + public Optional buildWriteItemForEcSignedPreKey(final UUID identifier, final byte deviceId, final ECSignedPreKey ecSignedPreKey) { - return ecSignedPreKeys.buildTransactWriteItem(identifier, deviceId, ecSignedPreKey); + return dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys() + ? Optional.of(ecSignedPreKeys.buildTransactWriteItem(identifier, deviceId, ecSignedPreKey)) + : Optional.empty(); } public TransactWriteItem buildWriteItemForLastResortKey(final UUID identifier, @@ -65,7 +72,11 @@ public class KeysManager { } public CompletableFuture storeEcSignedPreKeys(final UUID identifier, final Map keys) { - return ecSignedPreKeys.store(identifier, keys); + if (dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys()) { + return ecSignedPreKeys.store(identifier, keys); + } else { + return CompletableFuture.completedFuture(null); + } } public CompletableFuture storePqLastResort(final UUID identifier, final Map keys) { @@ -122,15 +133,19 @@ public class KeysManager { return CompletableFuture.allOf( ecPreKeys.delete(identifier), pqPreKeys.delete(identifier), - ecSignedPreKeys.delete(identifier, excludePrimaryDevice), + dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys() + ? ecSignedPreKeys.delete(identifier, excludePrimaryDevice) + : CompletableFuture.completedFuture(null), pqLastResortKeys.delete(identifier, excludePrimaryDevice)); } public CompletableFuture delete(final UUID accountUuid, final byte deviceId) { return CompletableFuture.allOf( - ecPreKeys.delete(accountUuid, deviceId), - pqPreKeys.delete(accountUuid, deviceId), - ecSignedPreKeys.delete(accountUuid, deviceId), - pqLastResortKeys.delete(accountUuid, deviceId)); + ecPreKeys.delete(accountUuid, deviceId), + pqPreKeys.delete(accountUuid, deviceId), + dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys() + ? ecSignedPreKeys.delete(accountUuid, deviceId) + : CompletableFuture.completedFuture(null), + pqLastResortKeys.delete(accountUuid, deviceId)); } } 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 80befaaa0..a76bdce91 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java @@ -168,8 +168,8 @@ public class AssignUsernameCommand extends EnvironmentCommand 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); + when(ecPreKeyMigrationConfiguration.deleteEcSignedPreKeys()).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 @@ -230,6 +246,22 @@ class KeysManagerTest { Set.copyOf(keysManager.getPqEnabledDevices(ACCOUNT_UUID).join())); } + @Test + void testStoreEcSignedPreKeyDisabled() { + when(ecPreKeyMigrationConfiguration.storeEcSignedPreKeys()).thenReturn(false); + + keysManager.storeEcSignedPreKeys(ACCOUNT_UUID, Map.of(DEVICE_ID, generateTestECSignedPreKey(1))).join(); + assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent()); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void buildWriteItemForEcSignedPreKey(final boolean enableSignedPreKeyWrite) { + when(ecPreKeyMigrationConfiguration.storeEcSignedPreKeys()).thenReturn(enableSignedPreKeyWrite); + assertEquals(enableSignedPreKeyWrite, + keysManager.buildWriteItemForEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID, generateTestECSignedPreKey(1)).isPresent()); + } + private static ECPreKey generateTestPreKey(final long keyId) { return new ECPreKey(keyId, Curve.generateKeyPair().getPublicKey()); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/LinkDeviceIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/LinkDeviceIntegrationTest.java index 4984179ac..349db1aaf 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/LinkDeviceIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/LinkDeviceIntegrationTest.java @@ -72,8 +72,8 @@ public class LinkDeviceIntegrationTest { 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() - ); + DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName(), + dynamicConfigurationManager); final Accounts accounts = new Accounts( DYNAMO_DB_EXTENSION.getDynamoDbClient(),