MirrorMirror
/
Signal-Server
mirror of https://github.com/signalapp/Signal-Server.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
Jon Chambers 2023-06-20 13:48:16 -04:00 committed by Jon Chambers
parent a3e82dfae8
commit 2d154eb0cf
8 changed files with 124 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
View File

@ -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(),

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

@ -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;
}
}

9
service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicECPreKeyMigrationConfiguration.java Normal file
View File

@ -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) {
}

17
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java
View File

@ -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) {

3
service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java
View File

@ -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(),

3
service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
View File

@ -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(),

53
service/src/main/java/org/whispersystems/textsecuregcm/workers/MigrateSignedECPreKeysCommand.java Normal file
View File

@ -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();
}
}

34
service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java
View File

@ -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());
}