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

Split traffic between the Postgres and Dynamo pre-key stores.

This commit is contained in:
Jon Chambers 2021-01-26 11:32:15 -05:00 committed by Jon Chambers
parent d4d9403829
commit 3298db8683
3 changed files with 40 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -420,7 +420,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
AttachmentControllerV1 attachmentControllerV1 = new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket());
AttachmentControllerV2 attachmentControllerV2 = new AttachmentControllerV2(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getRegion(), config.getAwsAttachmentsConfiguration().getBucket());
AttachmentControllerV3 attachmentControllerV3 = new AttachmentControllerV3(rateLimiters, config.getGcpAttachmentsConfiguration().getDomain(), config.getGcpAttachmentsConfiguration().getEmail(), config.getGcpAttachmentsConfiguration().getMaxSizeInBytes(), config.getGcpAttachmentsConfiguration().getPathPrefix(), config.getGcpAttachmentsConfiguration().getRsaSigningKey());
KeysController keysController = new KeysController(rateLimiters, keys, accountsManager, directoryQueue);
KeysController keysController = new KeysController(rateLimiters, keys, keysDynamoDb, accountsManager, directoryQueue, experimentEnrollmentManager);
MessageController messageController = new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, messagesManager, apnFallbackManager, featureFlagsManager);
ProfileController profileController = new ProfileController(rateLimiters, accountsManager, profilesManager, usernamesManager, cdnS3Client, profileCdnPolicyGenerator, profileCdnPolicySigner, config.getCdnConfiguration().getBucket(), zkProfileOperations, isZkEnabled);
StickerController stickerController = new StickerController(rateLimiters, config.getCdnConfiguration().getAccessKey(), config.getCdnConfiguration().getAccessSecret(), config.getCdnConfiguration().getRegion(), config.getCdnConfiguration().getBucket());

46
service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
View File

@ -6,8 +6,6 @@ package org.whispersystems.textsecuregcm.controllers;
import com.codahale.metrics.annotation.Timed;
import io.dropwizard.auth.Auth;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.auth.AmbiguousIdentifier;
import org.whispersystems.textsecuregcm.auth.Anonymous;
import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
@ -18,6 +16,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
import org.whispersystems.textsecuregcm.entities.PreKeyState;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.storage.Account;
@ -25,6 +24,8 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.PreKeyStore;
import javax.validation.Valid;
import javax.ws.rs.Consumes;
@ -45,24 +46,29 @@ import java.util.Optional;
@Path("/v2/keys")
public class KeysController {
private static final Logger logger = LoggerFactory.getLogger(KeysController.class);
private final RateLimiters rateLimiters;
private final Keys keys;
private final KeysDynamoDb keysDynamoDb;
private final AccountsManager accounts;
private final DirectoryQueue directoryQueue;
private final ExperimentEnrollmentManager experimentEnrollmentManager;
private final RateLimiters rateLimiters;
private final Keys keys;
private final AccountsManager accounts;
private final DirectoryQueue directoryQueue;
private static final String DYNAMODB_CONSUMER_EXPERIMENT = "keys_dynamodb_consumer";
private static final String DYNAMODB_PRODUCER_EXPERIMENT = "keys_dynamodb_producer";
public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts, DirectoryQueue directoryQueue) {
this.rateLimiters = rateLimiters;
this.keys = keys;
this.accounts = accounts;
this.directoryQueue = directoryQueue;
public KeysController(RateLimiters rateLimiters, Keys keys, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue, final ExperimentEnrollmentManager experimentEnrollmentManager) {
this.rateLimiters = rateLimiters;
this.keys = keys;
this.keysDynamoDb = keysDynamoDb;
this.accounts = accounts;
this.directoryQueue = directoryQueue;
this.experimentEnrollmentManager = experimentEnrollmentManager;
}
@GET
@Produces(MediaType.APPLICATION_JSON)
public PreKeyCount getStatus(@Auth Account account) {
int count = keys.getCount(account, account.getAuthenticatedDevice().get().getId());
int count = getPreKeyStoreForProducer(account).getCount(account, account.getAuthenticatedDevice().get().getId());
if (count > 0) {
count = count - 1;
@ -98,7 +104,7 @@ public class KeysController {
}
}
keys.store(account, device.getId(), preKeys.getPreKeys());
getPreKeyStoreForProducer(account).store(account, device.getId(), preKeys.getPreKeys());
}
@Timed
@ -179,14 +185,22 @@ public class KeysController {
private List<KeyRecord> getLocalKeys(Account destination, String deviceIdSelector) {
try {
if (deviceIdSelector.equals("*")) {
return keys.take(destination);
return getPreKeyStoreForConsumer(destination).take(destination);
}
long deviceId = Long.parseLong(deviceIdSelector);
return keys.take(destination, deviceId);
return getPreKeyStoreForConsumer(destination).take(destination, deviceId);
} catch (NumberFormatException e) {
throw new WebApplicationException(Response.status(422).build());
}
}
private PreKeyStore getPreKeyStoreForProducer(final Account account) {
return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_PRODUCER_EXPERIMENT) ? keysDynamoDb : keys;
}
private PreKeyStore getPreKeyStoreForConsumer(final Account account) {
return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_CONSUMER_EXPERIMENT) ? keysDynamoDb : keys;
}
}

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

@ -21,6 +21,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyCount;
import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
import org.whispersystems.textsecuregcm.entities.PreKeyState;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiter;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
@ -29,6 +30,7 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
import javax.ws.rs.client.Entity;
@ -69,10 +71,12 @@ public class KeysControllerTest {
private final SignedPreKey SAMPLE_SIGNED_KEY3 = new SignedPreKey( 3333, "barfoo", "sig33" );
private final SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid");
private final Keys keys = mock(Keys.class );
private final AccountsManager accounts = mock(AccountsManager.class);
private final DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
private final Account existsAccount = mock(Account.class );
private final Keys keys = mock(Keys.class );
private final KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class );
private final AccountsManager accounts = mock(AccountsManager.class );
private final DirectoryQueue directoryQueue = mock(DirectoryQueue.class );
private final Account existsAccount = mock(Account.class );
private final ExperimentEnrollmentManager experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class);
private RateLimiters rateLimiters = mock(RateLimiters.class);
private RateLimiter rateLimiter = mock(RateLimiter.class );
@ -82,7 +86,7 @@ public class KeysControllerTest {
.addProvider(AuthHelper.getAuthFilter())
.addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(Account.class, DisabledPermittedAccount.class)))
.setTestContainerFactory(new GrizzlyWebTestContainerFactory())
.addResource(new KeysController(rateLimiters, keys, accounts, directoryQueue))
.addResource(new KeysController(rateLimiters, keys, keysDynamoDb, accounts, directoryQueue, experimentEnrollmentManager))
.build();
@Before