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

Drop the old Postgres-based pre-key store.

This commit is contained in:
Jon Chambers 2021-02-05 16:16:27 -05:00 committed by Jon Chambers
parent 6865cdfce3
commit 04728ea4bc
11 changed files with 40 additions and 735 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -125,7 +125,6 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.FeatureFlags; import org.whispersystems.textsecuregcm.storage.FeatureFlags;
import org.whispersystems.textsecuregcm.storage.FeatureFlagsManager; import org.whispersystems.textsecuregcm.storage.FeatureFlagsManager;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.MessagePersister; import org.whispersystems.textsecuregcm.storage.MessagePersister;
import org.whispersystems.textsecuregcm.storage.Messages; import org.whispersystems.textsecuregcm.storage.Messages;
@ -293,7 +292,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
Usernames usernames = new Usernames(accountDatabase); Usernames usernames = new Usernames(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase); ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase); Profiles profiles = new Profiles(accountDatabase);
Keys keys = new Keys(accountDatabase, config.getAccountsDatabaseConfiguration().getKeyOperationRetryConfiguration());
KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName()); KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName());
Messages messages = new Messages(messageDatabase); Messages messages = new Messages(messageDatabase);
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, config.getMessageDynamoDbConfiguration().getTableName(), config.getMessageDynamoDbConfiguration().getTimeToLive()); MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, config.getMessageDynamoDbConfiguration().getTableName(), config.getMessageDynamoDbConfiguration().getTimeToLive());
@ -349,7 +347,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
MessagesCache messagesCache = new MessagesCache(messagesCluster, messagesCluster, keyspaceNotificationDispatchExecutor); MessagesCache messagesCache = new MessagesCache(messagesCluster, messagesCluster, keyspaceNotificationDispatchExecutor);
PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster); PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster);
MessagesManager messagesManager = new MessagesManager(messages, messagesDynamoDb, messagesCache, pushLatencyManager, experimentEnrollmentManager); MessagesManager messagesManager = new MessagesManager(messages, messagesDynamoDb, messagesCache, pushLatencyManager, experimentEnrollmentManager);
AccountsManager accountsManager = new AccountsManager(accounts, directory, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directory, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs); RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
FeatureFlagsManager featureFlagsManager = new FeatureFlagsManager(featureFlags, recurringJobExecutor); FeatureFlagsManager featureFlagsManager = new FeatureFlagsManager(featureFlags, recurringJobExecutor);
DeadLetterHandler deadLetterHandler = new DeadLetterHandler(accountsManager, messagesManager); DeadLetterHandler deadLetterHandler = new DeadLetterHandler(accountsManager, messagesManager);
@ -420,7 +418,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
AttachmentControllerV1 attachmentControllerV1 = new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket()); 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()); 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()); 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, keysDynamoDb, accountsManager, directoryQueue, experimentEnrollmentManager); KeysController keysController = new KeysController(rateLimiters, keysDynamoDb, accountsManager, directoryQueue);
MessageController messageController = new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, messagesManager, apnFallbackManager, featureFlagsManager); 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); 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()); StickerController stickerController = new StickerController(rateLimiters, config.getCdnConfiguration().getAccessKey(), config.getCdnConfiguration().getAccessSecret(), config.getCdnConfiguration().getRegion(), config.getCdnConfiguration().getBucket());

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

@ -16,16 +16,13 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem; import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
import org.whispersystems.textsecuregcm.entities.PreKeyState; import org.whispersystems.textsecuregcm.entities.PreKeyState;
import org.whispersystems.textsecuregcm.entities.SignedPreKey; import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiters; import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue; import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyRecord; import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.PreKeyStore;
import javax.validation.Valid; import javax.validation.Valid;
import javax.ws.rs.Consumes; import javax.ws.rs.Consumes;
@ -47,28 +44,21 @@ import java.util.Optional;
public class KeysController { public class KeysController {
private final RateLimiters rateLimiters; private final RateLimiters rateLimiters;
private final Keys keys;
private final KeysDynamoDb keysDynamoDb; private final KeysDynamoDb keysDynamoDb;
private final AccountsManager accounts; private final AccountsManager accounts;
private final DirectoryQueue directoryQueue; private final DirectoryQueue directoryQueue;
private final ExperimentEnrollmentManager experimentEnrollmentManager;
private static final String DYNAMODB_CONSUMER_EXPERIMENT = "keys_dynamodb_consumer"; public KeysController(RateLimiters rateLimiters, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue) {
private static final String DYNAMODB_PRODUCER_EXPERIMENT = "keys_dynamodb_producer";
public KeysController(RateLimiters rateLimiters, Keys keys, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue, final ExperimentEnrollmentManager experimentEnrollmentManager) {
this.rateLimiters = rateLimiters; this.rateLimiters = rateLimiters;
this.keys = keys;
this.keysDynamoDb = keysDynamoDb; this.keysDynamoDb = keysDynamoDb;
this.accounts = accounts; this.accounts = accounts;
this.directoryQueue = directoryQueue; this.directoryQueue = directoryQueue;
this.experimentEnrollmentManager = experimentEnrollmentManager;
} }
@GET @GET
@Produces(MediaType.APPLICATION_JSON) @Produces(MediaType.APPLICATION_JSON)
public PreKeyCount getStatus(@Auth Account account) { public PreKeyCount getStatus(@Auth Account account) {
int count = getPreKeyStoreForProducer(account).getCount(account, account.getAuthenticatedDevice().get().getId()); int count = keysDynamoDb.getCount(account, account.getAuthenticatedDevice().get().getId());
if (count > 0) { if (count > 0) {
count = count - 1; count = count - 1;
@ -104,7 +94,7 @@ public class KeysController {
} }
} }
getPreKeyStoreForProducer(account).store(account, device.getId(), preKeys.getPreKeys()); keysDynamoDb.store(account, device.getId(), preKeys.getPreKeys());
} }
@Timed @Timed
@ -185,22 +175,14 @@ public class KeysController {
private List<KeyRecord> getLocalKeys(Account destination, String deviceIdSelector) { private List<KeyRecord> getLocalKeys(Account destination, String deviceIdSelector) {
try { try {
if (deviceIdSelector.equals("*")) { if (deviceIdSelector.equals("*")) {
return getPreKeyStoreForConsumer(destination).take(destination); return keysDynamoDb.take(destination);
} }
long deviceId = Long.parseLong(deviceIdSelector); long deviceId = Long.parseLong(deviceIdSelector);
return getPreKeyStoreForConsumer(destination).take(destination, deviceId); return keysDynamoDb.take(destination, deviceId);
} catch (NumberFormatException e) { } catch (NumberFormatException e) {
throw new WebApplicationException(Response.status(422).build()); 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;
}
} }

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

@ -55,7 +55,6 @@ public class AccountsManager {
private final FaultTolerantRedisCluster cacheCluster; private final FaultTolerantRedisCluster cacheCluster;
private final DirectoryManager directory; private final DirectoryManager directory;
private final DirectoryQueue directoryQueue; private final DirectoryQueue directoryQueue;
private final Keys keys;
private final KeysDynamoDb keysDynamoDb; private final KeysDynamoDb keysDynamoDb;
private final MessagesManager messagesManager; private final MessagesManager messagesManager;
private final UsernamesManager usernamesManager; private final UsernamesManager usernamesManager;
@ -74,12 +73,11 @@ public class AccountsManager {
} }
} }
public AccountsManager(Accounts accounts, DirectoryManager directory, FaultTolerantRedisCluster cacheCluster, final DirectoryQueue directoryQueue, final Keys keys, final KeysDynamoDb keysDynamoDb, final MessagesManager messagesManager, final UsernamesManager usernamesManager, final ProfilesManager profilesManager) { public AccountsManager(Accounts accounts, DirectoryManager directory, FaultTolerantRedisCluster cacheCluster, final DirectoryQueue directoryQueue, final KeysDynamoDb keysDynamoDb, final MessagesManager messagesManager, final UsernamesManager usernamesManager, final ProfilesManager profilesManager) {
this.accounts = accounts; this.accounts = accounts;
this.directory = directory; this.directory = directory;
this.cacheCluster = cacheCluster; this.cacheCluster = cacheCluster;
this.directoryQueue = directoryQueue; this.directoryQueue = directoryQueue;
this.keys = keys;
this.keysDynamoDb = keysDynamoDb; this.keysDynamoDb = keysDynamoDb;
this.messagesManager = messagesManager; this.messagesManager = messagesManager;
this.usernamesManager = usernamesManager; this.usernamesManager = usernamesManager;
@ -152,7 +150,6 @@ public class AccountsManager {
directoryQueue.deleteAccount(account); directoryQueue.deleteAccount(account);
directory.remove(account.getNumber()); directory.remove(account.getNumber());
profilesManager.deleteAll(account.getUuid()); profilesManager.deleteAll(account.getUuid());
keys.delete(account);
keysDynamoDb.delete(account); keysDynamoDb.delete(account);
messagesManager.clear(account.getNumber(), account.getUuid()); messagesManager.clear(account.getNumber(), account.getUuid());
redisDelete(account); redisDelete(account);

164
service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java
View File

@ -1,164 +0,0 @@
/*
* Copyright 2013-2020 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import com.codahale.metrics.Counter;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SharedMetricRegistries;
import com.codahale.metrics.Timer;
import io.github.resilience4j.retry.Retry;
import org.jdbi.v3.core.JdbiException;
import org.jdbi.v3.core.statement.PreparedBatch;
import org.jdbi.v3.core.statement.UnableToExecuteStatementException;
import org.jdbi.v3.core.transaction.SerializableTransactionRunner;
import org.jdbi.v3.core.transaction.TransactionIsolationLevel;
import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
import org.whispersystems.textsecuregcm.entities.PreKey;
import org.whispersystems.textsecuregcm.storage.mappers.KeyRecordRowMapper;
import org.whispersystems.textsecuregcm.util.Constants;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Supplier;
import static com.codahale.metrics.MetricRegistry.name;
public class Keys implements PreKeyStore {
private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
private final Meter fallbackMeter = metricRegistry.meter(name(Keys.class, "fallback"));
private final Timer storeTimer = metricRegistry.timer(name(Keys.class, "store" ));
private final Timer getDevicetTimer = metricRegistry.timer(name(Keys.class, "getDevice"));
private final Timer getTimer = metricRegistry.timer(name(Keys.class, "get" ));
private final Timer getCountTimer = metricRegistry.timer(name(Keys.class, "getCount" ));
private final Timer vacuumTimer = metricRegistry.timer(name(Keys.class, "vacuum" ));
private final FaultTolerantDatabase database;
private final Retry retry;
public Keys(FaultTolerantDatabase database, RetryConfiguration retryConfiguration) {
this.database = database;
this.database.getDatabase().registerRowMapper(new KeyRecordRowMapper());
this.database.getDatabase().setTransactionHandler(new SerializableTransactionRunner());
this.database.getDatabase().getConfig(SerializableTransactionRunner.Configuration.class).setMaxRetries(10);
this.retry = Retry.of("keys", retryConfiguration.toRetryConfigBuilder().build());
}
@Override
public void store(Account account, long deviceId, List<PreKey> keys) {
final String number = account.getNumber();
retry.executeRunnable(() -> {
database.use(jdbi -> jdbi.useTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> {
try (Timer.Context ignored = storeTimer.time()) {
PreparedBatch preparedBatch = handle.prepareBatch("INSERT INTO keys (number, device_id, key_id, public_key) VALUES (:number, :device_id, :key_id, :public_key)");
for (PreKey key : keys) {
preparedBatch.bind("number", number)
.bind("device_id", deviceId)
.bind("key_id", key.getKeyId())
.bind("public_key", key.getPublicKey())
.add();
}
handle.createUpdate("DELETE FROM keys WHERE number = :number AND device_id = :device_id")
.bind("number", number)
.bind("device_id", deviceId)
.execute();
preparedBatch.execute();
}
}));
});
}
@Override
public List<KeyRecord> take(Account account, long deviceId) {
/*
final String number = account.getNumber();
try {
return database.with(jdbi -> jdbi.inTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> {
try (Timer.Context ignored = getDevicetTimer.time()) {
return handle.createQuery("DELETE FROM keys WHERE id IN (SELECT id FROM keys WHERE number = :number AND device_id = :device_id ORDER BY key_id ASC LIMIT 1) RETURNING *")
.bind("number", number)
.bind("device_id", deviceId)
.mapTo(KeyRecord.class)
.list();
}
}));
} catch (JdbiException e) {
// TODO 2021-01-13 Replace this with a retry once desktop clients better handle HTTP/500 responses
fallbackMeter.mark();
return Collections.emptyList();
} */
// 2021-01-15 Emergency service recovery measure
return new LinkedList<>();
}
@Override
public List<KeyRecord> take(Account account) {
/*
final String number = account.getNumber();
try {
return database.with(jdbi -> jdbi.inTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> {
try (Timer.Context ignored = getTimer.time()) {
return handle.createQuery("DELETE FROM keys WHERE id IN (SELECT DISTINCT ON (number, device_id) id FROM keys WHERE number = :number ORDER BY number, device_id, key_id ASC) RETURNING *")
.bind("number", number)
.mapTo(KeyRecord.class)
.list();
}
}));
} catch (JdbiException e) {
// TODO 2021-01-13 Replace this with a retry once desktop clients better handle HTTP/500 responses
fallbackMeter.mark();
return Collections.emptyList();
} */
// 2021-01-15 Emergency service recovery measure
return new LinkedList<>();
}
@Override
public int getCount(Account account, long deviceId) {
final String number = account.getNumber();
return database.with(jdbi -> jdbi.withHandle(handle -> {
try (Timer.Context ignored = getCountTimer.time()) {
return handle.createQuery("SELECT COUNT(*) FROM keys WHERE number = :number AND device_id = :device_id")
.bind("number", number)
.bind("device_id", deviceId)
.mapTo(Integer.class)
.findOnly();
}
}));
}
public void delete(final Account account) {
final String number = account.getNumber();
database.use(jdbi -> jdbi.useHandle(handle -> {
try (Timer.Context ignored = getCountTimer.time()) {
handle.createUpdate("DELETE FROM keys WHERE number = :number")
.bind("number", number)
.execute();
}
}));
}
public void vacuum() {
database.use(jdbi -> jdbi.useHandle(handle -> {
try (Timer.Context ignored = vacuumTimer.time()) {
handle.execute("VACUUM keys");
}
}));
}
}

7
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java
View File

@ -31,7 +31,7 @@ import java.util.UUID;
import static com.codahale.metrics.MetricRegistry.name; import static com.codahale.metrics.MetricRegistry.name;
public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { public class KeysDynamoDb extends AbstractDynamoDbStore {
private final Table table; private final Table table;
@ -53,7 +53,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
this.table = dynamoDB.getTable(tableName); this.table = dynamoDB.getTable(tableName);
} }
@Override
public void store(final Account account, final long deviceId, final List<PreKey> keys) { public void store(final Account account, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> { STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId); delete(account, deviceId);
@ -70,7 +69,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
}); });
} }
@Override
public List<KeyRecord> take(final Account account, final long deviceId) { public List<KeyRecord> take(final Account account, final long deviceId) {
return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> { return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> {
final byte[] partitionKey = getPartitionKey(account.getUuid()); final byte[] partitionKey = getPartitionKey(account.getUuid());
@ -106,7 +104,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
}); });
} }
@Override
public List<KeyRecord> take(final Account account) { public List<KeyRecord> take(final Account account) {
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final List<KeyRecord> keyRecords = new ArrayList<>(); final List<KeyRecord> keyRecords = new ArrayList<>();
@ -119,7 +116,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
}); });
} }
@Override
public int getCount(final Account account, final long deviceId) { public int getCount(final Account account, final long deviceId) {
return GET_KEY_COUNT_TIMER.record(() -> { return GET_KEY_COUNT_TIMER.record(() -> {
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)") final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
@ -133,7 +129,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
}); });
} }
@Override
public void delete(final Account account) { public void delete(final Account account) {
DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid") final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid")

23
service/src/main/java/org/whispersystems/textsecuregcm/storage/PreKeyStore.java
View File

@ -1,23 +0,0 @@
/*
* Copyright 2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import org.whispersystems.textsecuregcm.entities.PreKey;
import java.util.List;
public interface PreKeyStore {
void store(Account account, long deviceId, List<PreKey> keys);
int getCount(Account account, long deviceId);
List<KeyRecord> take(Account account, long deviceId);
List<KeyRecord> take(Account account);
void delete(Account account);
}

9
service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
View File

@ -33,7 +33,6 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.DirectoryManager; import org.whispersystems.textsecuregcm.storage.DirectoryManager;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Messages; import org.whispersystems.textsecuregcm.storage.Messages;
import org.whispersystems.textsecuregcm.storage.MessagesCache; import org.whispersystems.textsecuregcm.storage.MessagesCache;
@ -107,21 +106,19 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
.withCredentials(InstanceProfileCredentialsProvider.getInstance()); .withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDB messageDynamoDb = new DynamoDB(clientBuilder.build()); DynamoDB messageDynamoDb = new DynamoDB(clientBuilder.build());
DynamoDB preKeyDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build()); DynamoDB preKeysDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources); FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources);
ExecutorService keyspaceNotificationDispatchExecutor = environment.lifecycle().executorService(name(getClass(), "keyspaceNotification-%d")).maxThreads(4).build(); ExecutorService keyspaceNotificationDispatchExecutor = environment.lifecycle().executorService(name(getClass(), "keyspaceNotification-%d")).maxThreads(4).build();
DynamicConfigurationManager dynamicConfigurationManager = new DynamicConfigurationManager(configuration.getAppConfig().getApplication(), configuration.getAppConfig().getEnvironment(), configuration.getAppConfig().getConfigurationName()); DynamicConfigurationManager dynamicConfigurationManager = new DynamicConfigurationManager(configuration.getAppConfig().getApplication(), configuration.getAppConfig().getEnvironment(), configuration.getAppConfig().getConfigurationName());
ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
Accounts accounts = new Accounts(accountDatabase); Accounts accounts = new Accounts(accountDatabase);
Usernames usernames = new Usernames(accountDatabase); Usernames usernames = new Usernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase); Profiles profiles = new Profiles(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase); ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
Keys keys = new Keys(accountDatabase, configuration.getAccountsDatabaseConfiguration().getKeyOperationRetryConfiguration()); KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeysDynamoDb, configuration.getKeysDynamoDbConfiguration().getTableName());
KeysDynamoDb keysDynamoDb = new KeysDynamoDb(messageDynamoDb, configuration.getKeysDynamoDbConfiguration().getTableName());
Messages messages = new Messages(messageDatabase); Messages messages = new Messages(messageDatabase);
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, configuration.getMessageDynamoDbConfiguration().getTableName(), configuration.getMessageDynamoDbConfiguration().getTimeToLive()); MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, configuration.getMessageDynamoDbConfiguration().getTableName(), configuration.getMessageDynamoDbConfiguration().getTimeToLive());
ReplicatedJedisPool redisClient = new RedisClientFactory("directory_cache_delete_command", configuration.getDirectoryConfiguration().getRedisConfiguration().getUrl(), configuration.getDirectoryConfiguration().getRedisConfiguration().getReplicaUrls(), configuration.getDirectoryConfiguration().getRedisConfiguration().getCircuitBreakerConfiguration()).getRedisClientPool(); ReplicatedJedisPool redisClient = new RedisClientFactory("directory_cache_delete_command", configuration.getDirectoryConfiguration().getRedisConfiguration().getUrl(), configuration.getDirectoryConfiguration().getRedisConfiguration().getReplicaUrls(), configuration.getDirectoryConfiguration().getRedisConfiguration().getCircuitBreakerConfiguration()).getRedisClientPool();
@ -135,7 +132,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
UsernamesManager usernamesManager = new UsernamesManager(usernames, reservedUsernames, cacheCluster); UsernamesManager usernamesManager = new UsernamesManager(usernames, reservedUsernames, cacheCluster);
ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster); ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster);
MessagesManager messagesManager = new MessagesManager(messages, messagesDynamoDb, messagesCache, pushLatencyManager, new ExperimentEnrollmentManager(dynamicConfigurationManager)); MessagesManager messagesManager = new MessagesManager(messages, messagesDynamoDb, messagesCache, pushLatencyManager, new ExperimentEnrollmentManager(dynamicConfigurationManager));
AccountsManager accountsManager = new AccountsManager(accounts, directory, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directory, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
for (String user: users) { for (String user: users) {
Optional<Account> account = accountsManager.get(user); Optional<Account> account = accountsManager.get(user);

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

@ -14,7 +14,6 @@ import org.whispersystems.textsecuregcm.configuration.DatabaseConfiguration;
import org.whispersystems.textsecuregcm.storage.Accounts; import org.whispersystems.textsecuregcm.storage.Accounts;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.FeatureFlags; import org.whispersystems.textsecuregcm.storage.FeatureFlags;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.Messages; import org.whispersystems.textsecuregcm.storage.Messages;
import org.whispersystems.textsecuregcm.storage.PendingAccounts; import org.whispersystems.textsecuregcm.storage.PendingAccounts;
@ -46,7 +45,6 @@ public class VacuumCommand extends ConfiguredCommand<WhisperServerConfiguration>
FaultTolerantDatabase messageDatabase = new FaultTolerantDatabase("message_database_vacuum", messageJdbi, messageDbConfig.getCircuitBreakerConfiguration()); FaultTolerantDatabase messageDatabase = new FaultTolerantDatabase("message_database_vacuum", messageJdbi, messageDbConfig.getCircuitBreakerConfiguration());
Accounts accounts = new Accounts(accountDatabase); Accounts accounts = new Accounts(accountDatabase);
Keys keys = new Keys(accountDatabase, config.getAccountsDatabaseConfiguration().getKeyOperationRetryConfiguration());
PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase); PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase);
Messages messages = new Messages(messageDatabase); Messages messages = new Messages(messageDatabase);
FeatureFlags featureFlags = new FeatureFlags(accountDatabase); FeatureFlags featureFlags = new FeatureFlags(accountDatabase);
@ -57,9 +55,6 @@ public class VacuumCommand extends ConfiguredCommand<WhisperServerConfiguration>
logger.info("Vacuuming pending_accounts..."); logger.info("Vacuuming pending_accounts...");
pendingAccounts.vacuum(); pendingAccounts.vacuum();
logger.info("Vacuuming keys...");
keys.vacuum();
logger.info("Vacuuming messages..."); logger.info("Vacuuming messages...");
messages.vacuum(); messages.vacuum();

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

@ -29,7 +29,6 @@ import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyRecord; import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper; import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
@ -71,12 +70,10 @@ public class KeysControllerTest {
private final SignedPreKey SAMPLE_SIGNED_KEY3 = new SignedPreKey( 3333, "barfoo", "sig33" ); 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 SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid");
private final Keys keys = mock(Keys.class );
private final KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class ); private final KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class );
private final AccountsManager accounts = mock(AccountsManager.class ); private final AccountsManager accounts = mock(AccountsManager.class );
private final DirectoryQueue directoryQueue = mock(DirectoryQueue.class ); private final DirectoryQueue directoryQueue = mock(DirectoryQueue.class );
private final Account existsAccount = mock(Account.class ); private final Account existsAccount = mock(Account.class );
private final ExperimentEnrollmentManager experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class);
private RateLimiters rateLimiters = mock(RateLimiters.class); private RateLimiters rateLimiters = mock(RateLimiters.class);
private RateLimiter rateLimiter = mock(RateLimiter.class ); private RateLimiter rateLimiter = mock(RateLimiter.class );
@ -86,7 +83,7 @@ public class KeysControllerTest {
.addProvider(AuthHelper.getAuthFilter()) .addProvider(AuthHelper.getAuthFilter())
.addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(Account.class, DisabledPermittedAccount.class))) .addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(Account.class, DisabledPermittedAccount.class)))
.setTestContainerFactory(new GrizzlyWebTestContainerFactory()) .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
.addResource(new KeysController(rateLimiters, keys, keysDynamoDb, accounts, directoryQueue, experimentEnrollmentManager)) .addResource(new KeysController(rateLimiters, keysDynamoDb, accounts, directoryQueue))
.build(); .build();
@Before @Before
@ -145,16 +142,16 @@ public class KeysControllerTest {
List<KeyRecord> singleDevice = new LinkedList<>(); List<KeyRecord> singleDevice = new LinkedList<>();
singleDevice.add(SAMPLE_KEY); singleDevice.add(SAMPLE_KEY);
when(keys.take(eq(existsAccount), eq(1L))).thenReturn(singleDevice); when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(singleDevice);
List<KeyRecord> multiDevice = new LinkedList<>(); List<KeyRecord> multiDevice = new LinkedList<>();
multiDevice.add(SAMPLE_KEY); multiDevice.add(SAMPLE_KEY);
multiDevice.add(SAMPLE_KEY2); multiDevice.add(SAMPLE_KEY2);
multiDevice.add(SAMPLE_KEY3); multiDevice.add(SAMPLE_KEY3);
multiDevice.add(SAMPLE_KEY4); multiDevice.add(SAMPLE_KEY4);
when(keys.take(existsAccount)).thenReturn(multiDevice); when(keysDynamoDb.take(existsAccount)).thenReturn(multiDevice);
when(keys.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5); when(keysDynamoDb.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5);
when(AuthHelper.VALID_DEVICE.getSignedPreKey()).thenReturn(VALID_DEVICE_SIGNED_KEY); when(AuthHelper.VALID_DEVICE.getSignedPreKey()).thenReturn(VALID_DEVICE_SIGNED_KEY);
when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null); when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null);
@ -171,7 +168,7 @@ public class KeysControllerTest {
assertThat(result.getCount()).isEqualTo(4); assertThat(result.getCount()).isEqualTo(4);
verify(keys).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L));
} }
@Test @Test
@ -185,7 +182,7 @@ public class KeysControllerTest {
assertThat(result.getCount()).isEqualTo(4); assertThat(result.getCount()).isEqualTo(4);
verify(keys).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L));
} }
@ -285,8 +282,8 @@ public class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keys).take(eq(existsAccount), eq(1L)); verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@Test @Test
@ -303,8 +300,8 @@ public class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keys).take(eq(existsAccount), eq(1L)); verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@ -322,8 +319,8 @@ public class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keys).take(eq(existsAccount), eq(1L)); verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@Test @Test
@ -340,8 +337,8 @@ public class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keys).take(eq(existsAccount), eq(1L)); verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@ -354,7 +351,7 @@ public class KeysControllerTest {
.get(); .get();
assertThat(response.getStatus()).isEqualTo(401); assertThat(response.getStatus()).isEqualTo(401);
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@Test @Test
@ -366,7 +363,7 @@ public class KeysControllerTest {
.get(); .get();
assertThat(response.getStatus()).isEqualTo(401); assertThat(response.getStatus()).isEqualTo(401);
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@ -416,8 +413,8 @@ public class KeysControllerTest {
assertThat(signedPreKey).isNull(); assertThat(signedPreKey).isNull();
assertThat(deviceId).isEqualTo(4); assertThat(deviceId).isEqualTo(4);
verify(keys).take(eq(existsAccount)); verify(keysDynamoDb).take(eq(existsAccount));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@Test @Test
@ -466,8 +463,8 @@ public class KeysControllerTest {
assertThat(signedPreKey).isNull(); assertThat(signedPreKey).isNull();
assertThat(deviceId).isEqualTo(4); assertThat(deviceId).isEqualTo(4);
verify(keys).take(eq(existsAccount)); verify(keysDynamoDb).take(eq(existsAccount));
verifyNoMoreInteractions(keys); verifyNoMoreInteractions(keysDynamoDb);
} }
@ -535,7 +532,7 @@ public class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204); assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class); ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class);
verify(keys).store(eq(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture()); verify(keysDynamoDb).store(eq(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture());
List<PreKey> capturedList = listCaptor.getValue(); List<PreKey> capturedList = listCaptor.getValue();
assertThat(capturedList.size()).isEqualTo(1); assertThat(capturedList.size()).isEqualTo(1);
@ -569,7 +566,7 @@ public class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204); assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class); ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class);
verify(keys).store(eq(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture()); verify(keysDynamoDb).store(eq(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture());
List<PreKey> capturedList = listCaptor.getValue(); List<PreKey> capturedList = listCaptor.getValue();
assertThat(capturedList.size()).isEqualTo(1); assertThat(capturedList.size()).isEqualTo(1);

19
service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java
View File

@ -14,7 +14,6 @@ import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.Accounts; import org.whispersystems.textsecuregcm.storage.Accounts;
import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.DirectoryManager; import org.whispersystems.textsecuregcm.storage.DirectoryManager;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.MessagesManager; import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.ProfilesManager; import org.whispersystems.textsecuregcm.storage.ProfilesManager;
@ -45,7 +44,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -56,7 +54,7 @@ public class AccountsManagerTest {
when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(uuid.toString()); when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(uuid.toString());
when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}"); when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}");
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> account = accountsManager.get("+14152222222"); Optional<Account> account = accountsManager.get("+14152222222");
assertTrue(account.isPresent()); assertTrue(account.isPresent());
@ -76,7 +74,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -86,7 +83,7 @@ public class AccountsManagerTest {
when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}"); when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}");
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> account = accountsManager.get(uuid); Optional<Account> account = accountsManager.get(uuid);
assertTrue(account.isPresent()); assertTrue(account.isPresent());
@ -107,7 +104,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -118,7 +114,7 @@ public class AccountsManagerTest {
when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(null); when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(null);
when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account)); when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account));
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> retrieved = accountsManager.get("+14152222222"); Optional<Account> retrieved = accountsManager.get("+14152222222");
assertTrue(retrieved.isPresent()); assertTrue(retrieved.isPresent());
@ -140,7 +136,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -151,7 +146,7 @@ public class AccountsManagerTest {
when(commands.get(eq("Account3::" + uuid))).thenReturn(null); when(commands.get(eq("Account3::" + uuid))).thenReturn(null);
when(accounts.get(eq(uuid))).thenReturn(Optional.of(account)); when(accounts.get(eq(uuid))).thenReturn(Optional.of(account));
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> retrieved = accountsManager.get(uuid); Optional<Account> retrieved = accountsManager.get(uuid);
assertTrue(retrieved.isPresent()); assertTrue(retrieved.isPresent());
@ -173,7 +168,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -184,7 +178,7 @@ public class AccountsManagerTest {
when(commands.get(eq("AccountMap::+14152222222"))).thenThrow(new RedisException("Connection lost!")); when(commands.get(eq("AccountMap::+14152222222"))).thenThrow(new RedisException("Connection lost!"));
when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account)); when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account));
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> retrieved = accountsManager.get("+14152222222"); Optional<Account> retrieved = accountsManager.get("+14152222222");
assertTrue(retrieved.isPresent()); assertTrue(retrieved.isPresent());
@ -206,7 +200,6 @@ public class AccountsManagerTest {
Accounts accounts = mock(Accounts.class); Accounts accounts = mock(Accounts.class);
DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryManager directoryManager = mock(DirectoryManager.class);
DirectoryQueue directoryQueue = mock(DirectoryQueue.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class);
Keys keys = mock(Keys.class);
KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class);
MessagesManager messagesManager = mock(MessagesManager.class); MessagesManager messagesManager = mock(MessagesManager.class);
UsernamesManager usernamesManager = mock(UsernamesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class);
@ -217,7 +210,7 @@ public class AccountsManagerTest {
when(commands.get(eq("Account3::" + uuid))).thenThrow(new RedisException("Connection lost!")); when(commands.get(eq("Account3::" + uuid))).thenThrow(new RedisException("Connection lost!"));
when(accounts.get(eq(uuid))).thenReturn(Optional.of(account)); when(accounts.get(eq(uuid))).thenReturn(Optional.of(account));
AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager);
Optional<Account> retrieved = accountsManager.get(uuid); Optional<Account> retrieved = accountsManager.get(uuid);
assertTrue(retrieved.isPresent()); assertTrue(retrieved.isPresent());

462
service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java
View File

@ -1,462 +0,0 @@
/*
* Copyright 2013-2020 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.tests.storage;
import com.opentable.db.postgres.embedded.LiquibasePreparer;
import com.opentable.db.postgres.junit.EmbeddedPostgresRules;
import com.opentable.db.postgres.junit.PreparedDbRule;
import io.github.resilience4j.circuitbreaker.CallNotPermittedException;
import org.jdbi.v3.core.HandleCallback;
import org.jdbi.v3.core.HandleConsumer;
import org.jdbi.v3.core.Jdbi;
import org.jdbi.v3.core.statement.UnableToExecuteStatementException;
import org.jdbi.v3.core.transaction.SerializableTransactionRunner;
import org.jdbi.v3.core.transaction.TransactionException;
import org.jdbi.v3.core.transaction.TransactionIsolationLevel;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.configuration.AccountsDatabaseConfiguration;
import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
import org.whispersystems.textsecuregcm.entities.PreKey;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
public class KeysTest {
@Rule
public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(LiquibasePreparer.forClasspathLocation("accountsdb.xml"));
private Account firstAccount;
private Account secondAccount;
private Keys keys;
@Before
public void setup() {
FaultTolerantDatabase faultTolerantDatabase = new FaultTolerantDatabase("keysTest",
Jdbi.create(db.getTestDatabase()),
new CircuitBreakerConfiguration());
this.keys = new Keys(faultTolerantDatabase, new RetryConfiguration());
this.firstAccount = mock(Account.class);
this.secondAccount = mock(Account.class);
when(firstAccount.getNumber()).thenReturn("+14152222222");
when(secondAccount.getNumber()).thenReturn("+14151111111");
}
@Test
public void testPopulateKeys() throws SQLException {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
List<PreKey> deviceTwoPreKeys = new LinkedList<>();
List<PreKey> oldAnotherDeviceOnePrKeys = new LinkedList<>();
List<PreKey> anotherDeviceOnePreKeys = new LinkedList<>();
List<PreKey> anotherDeviceTwoPreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i));
}
for (int i=1;i<=100;i++) {
oldAnotherDeviceOnePrKeys.add(new PreKey(i, "OldPublicKey" + i));
anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i));
anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
keys.store(firstAccount, 2, deviceTwoPreKeys);
keys.store(secondAccount, 1, oldAnotherDeviceOnePrKeys);
keys.store(secondAccount, 1, anotherDeviceOnePreKeys);
keys.store(secondAccount, 2, anotherDeviceTwoPreKeys);
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("SELECT * FROM keys WHERE number = ? AND device_id = ? ORDER BY key_id");
verifyStoredState(statement, firstAccount, 1);
verifyStoredState(statement, firstAccount, 2);
verifyStoredState(statement, secondAccount, 1);
verifyStoredState(statement, secondAccount, 2);
}
@Test
public void testKeyCount() {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100);
}
@Ignore
@Test
public void testGetForDevice() {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
List<PreKey> deviceTwoPreKeys = new LinkedList<>();
List<PreKey> anotherDeviceOnePreKeys = new LinkedList<>();
List<PreKey> anotherDeviceTwoPreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i));
}
for (int i=1;i<=100;i++) {
anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i));
anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
keys.store(firstAccount, 2, deviceTwoPreKeys);
keys.store(secondAccount, 1, anotherDeviceOnePreKeys);
keys.store(secondAccount, 2, anotherDeviceTwoPreKeys);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100);
List<KeyRecord> records = keys.take(firstAccount, 1);
assertThat(records.size()).isEqualTo(1);
assertThat(records.get(0).getKeyId()).isEqualTo(1);
assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device1PublicKey1");
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(99);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100);
records = keys.take(firstAccount, 1);
assertThat(records.size()).isEqualTo(1);
assertThat(records.get(0).getKeyId()).isEqualTo(2);
assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device1PublicKey2");
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100);
records = keys.take(firstAccount, 2);
assertThat(records.size()).isEqualTo(1);
assertThat(records.get(0).getKeyId()).isEqualTo(1);
assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device2PublicKey1");
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(99);
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100);
}
@Ignore
@Test
public void testGetForAllDevices() {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
List<PreKey> deviceTwoPreKeys = new LinkedList<>();
List<PreKey> anotherDeviceOnePreKeys = new LinkedList<>();
List<PreKey> anotherDeviceTwoPreKeys = new LinkedList<>();
List<PreKey> anotherDeviceThreePreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i));
}
for (int i=1;i<=100;i++) {
anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i));
anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i));
anotherDeviceThreePreKeys.add(new PreKey(i, "+14151111111Device3PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
keys.store(firstAccount, 2, deviceTwoPreKeys);
keys.store(secondAccount, 1, anotherDeviceOnePreKeys);
keys.store(secondAccount, 2, anotherDeviceTwoPreKeys);
keys.store(secondAccount, 3, anotherDeviceThreePreKeys);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100);
List<KeyRecord> records = keys.take(firstAccount);
assertThat(records.size()).isEqualTo(2);
assertThat(records.get(0).getKeyId()).isEqualTo(1);
assertThat(records.get(1).getKeyId()).isEqualTo(1);
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device1PublicKey1"))).isTrue();
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device2PublicKey1"))).isTrue();
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(99);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(99);
records = keys.take(firstAccount);
assertThat(records.size()).isEqualTo(2);
assertThat(records.get(0).getKeyId()).isEqualTo(2);
assertThat(records.get(1).getKeyId()).isEqualTo(2);
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device1PublicKey2"))).isTrue();
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device2PublicKey2"))).isTrue();
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(98);
records = keys.take(secondAccount);
assertThat(records.size()).isEqualTo(3);
assertThat(records.get(0).getKeyId()).isEqualTo(1);
assertThat(records.get(1).getKeyId()).isEqualTo(1);
assertThat(records.get(2).getKeyId()).isEqualTo(1);
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device1PublicKey1"))).isTrue();
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device2PublicKey1"))).isTrue();
assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device3PublicKey1"))).isTrue();
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(99);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(99);
assertThat(keys.getCount(secondAccount, 3)).isEqualTo(99);
}
@Ignore
@Test
public void testGetForAllDevicesParallel() throws InterruptedException {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
List<PreKey> deviceTwoPreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
keys.store(firstAccount, 2, deviceTwoPreKeys);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100);
List<Thread> threads = new LinkedList<>();
for (int i=0;i<20;i++) {
Thread thread = new Thread(() -> {
List<KeyRecord> results = null;
final int MAX_RETRIES = 5;
for (int retryAttempt = 0; results == null && retryAttempt < MAX_RETRIES; ++retryAttempt) {
try {
results = keys.take(firstAccount);
} catch (UnableToExecuteStatementException e) {
if (retryAttempt == MAX_RETRIES - 1) {
throw e;
}
}
}
assertThat(results).isNotNull();
assertThat(results.size()).isEqualTo(2);
});
thread.start();
threads.add(thread);
}
for (Thread thread : threads) {
thread.join();
}
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(80);
assertThat(keys.getCount(firstAccount,2)).isEqualTo(80);
}
@Test
public void testDelete() {
List<PreKey> deviceOnePreKeys = new LinkedList<>();
List<PreKey> deviceTwoPreKeys = new LinkedList<>();
List<PreKey> anotherDeviceOnePreKeys = new LinkedList<>();
List<PreKey> anotherDeviceTwoPreKeys = new LinkedList<>();
List<PreKey> anotherDeviceThreePreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i));
}
for (int i=1;i<=100;i++) {
anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i));
anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i));
anotherDeviceThreePreKeys.add(new PreKey(i, "+14151111111Device3PublicKey" + i));
}
keys.store(firstAccount, 1, deviceOnePreKeys);
keys.store(firstAccount, 2, deviceTwoPreKeys);
keys.store(secondAccount, 1, anotherDeviceOnePreKeys);
keys.store(secondAccount, 2, anotherDeviceTwoPreKeys);
keys.store(secondAccount, 3, anotherDeviceThreePreKeys);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 3)).isEqualTo(100);
keys.delete(firstAccount);
assertThat(keys.getCount(firstAccount, 1)).isEqualTo(0);
assertThat(keys.getCount(firstAccount, 2)).isEqualTo(0);
assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100);
assertThat(keys.getCount(secondAccount, 3)).isEqualTo(100);
}
@Test
public void testEmptyKeyGet() {
List<KeyRecord> records = keys.take(firstAccount);
assertThat(records.isEmpty()).isTrue();
}
@Test
public void testVacuum() {
keys.vacuum();
}
@Test
public void testBreaker() throws InterruptedException {
Jdbi jdbi = mock(Jdbi.class);
doThrow(new TransactionException("Database error!")).when(jdbi).useTransaction(any(TransactionIsolationLevel.class), any(HandleConsumer.class));
when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class));
CircuitBreakerConfiguration configuration = new CircuitBreakerConfiguration();
configuration.setWaitDurationInOpenStateInSeconds(1);
configuration.setRingBufferSizeInHalfOpenState(1);
configuration.setRingBufferSizeInClosedState(2);
configuration.setFailureRateThreshold(50);
RetryConfiguration retryConfiguration = new RetryConfiguration();
retryConfiguration.setMaxAttempts(1);
Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, configuration), retryConfiguration);
List<PreKey> deviceOnePreKeys = new LinkedList<>();
for (int i=1;i<=100;i++) {
deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i));
}
try {
keys.store(firstAccount, 1, deviceOnePreKeys);
throw new AssertionError();
} catch (TransactionException e) {
// good
}
try {
keys.store(firstAccount, 1, deviceOnePreKeys);
throw new AssertionError();
} catch (TransactionException e) {
// good
}
try {
keys.store(firstAccount, 1, deviceOnePreKeys);
throw new AssertionError();
} catch (CallNotPermittedException e) {
// good
}
Thread.sleep(1100);
try {
keys.store(firstAccount, 1, deviceOnePreKeys);
throw new AssertionError();
} catch (TransactionException e) {
// good
}
}
@Test
public void testRetry() {
Jdbi jdbi = mock(Jdbi.class);
doThrow(new TransactionException("Database error!")).doNothing().when(jdbi).useTransaction(any(TransactionIsolationLevel.class), any(HandleConsumer.class));
when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class));
Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, new CircuitBreakerConfiguration()), new RetryConfiguration());
// We're happy as long as nothing throws an exception
Account account = mock(Account.class);
when(account.getNumber()).thenReturn("+18005551234");
keys.store(account, 1, Collections.emptyList());
}
@Test
public void testGetKeysWithException() {
Jdbi jdbi = mock(Jdbi.class);
when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class));
when(jdbi.inTransaction(any(TransactionIsolationLevel.class), any(HandleCallback.class)))
.thenThrow(new TransactionException("Database error!"));
Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, new CircuitBreakerConfiguration()), new RetryConfiguration());
Account account = mock(Account.class);
when(account.getNumber()).thenReturn("+18005551234");
assertThat(keys.take(account)).isEqualTo(Collections.emptyList());
assertThat(keys.take(account, 1)).isEqualTo(Collections.emptyList());
}
private void verifyStoredState(PreparedStatement statement, Account account, int deviceId) throws SQLException {
statement.setString(1, account.getNumber());
statement.setInt(2, deviceId);
ResultSet resultSet = statement.executeQuery();
int rowCount = 1;
while (resultSet.next()) {
long keyId = resultSet.getLong("key_id");
String publicKey = resultSet.getString("public_key");
assertThat(keyId).isEqualTo(rowCount);
assertThat(publicKey).isEqualTo(account.getNumber() + "Device" + deviceId + "PublicKey" + rowCount);
rowCount++;
}
resultSet.close();
assertThat(rowCount).isEqualTo(101);
}
}