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

Switch DynamoDB to AWSv2. 

Switch from using com.amazonaws.services.dynamodbv2 to using
software.amazon.awssdk.services.dynamodb for all current DynamoDB uses.
This commit is contained in:
Graeme Connell 2021-05-24 16:43:56 -06:00 committed by gram-signal
parent cbd9681e3e
commit c545cff1b3
31 changed files with 1114 additions and 876 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
service/pom.xml
View File

@ -256,6 +256,10 @@
<groupId>software.amazon.awssdk</groupId> <groupId>software.amazon.awssdk</groupId>
<artifactId>s3</artifactId> <artifactId>s3</artifactId>
</dependency> </dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>dynamodb</artifactId>
</dependency>
<dependency> <dependency>
<groupId>com.amazonaws</groupId> <groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-core</artifactId> <artifactId>aws-java-sdk-core</artifactId>
@ -272,10 +276,6 @@
<groupId>com.amazonaws</groupId> <groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-appconfig</artifactId> <artifactId>aws-java-sdk-appconfig</artifactId>
</dependency> </dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-dynamodb</artifactId>
</dependency>
<dependency> <dependency>
<groupId>redis.clients</groupId> <groupId>redis.clients</groupId>

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

@ -12,11 +12,6 @@ import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSStaticCredentialsProvider; import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials; import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.auth.InstanceProfileCredentialsProvider; import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.s3.AmazonS3; import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3Client; import com.amazonaws.services.s3.AmazonS3Client;
import com.codahale.metrics.SharedMetricRegistries; import com.codahale.metrics.SharedMetricRegistries;
@ -194,6 +189,7 @@ import org.whispersystems.textsecuregcm.storage.Usernames;
import org.whispersystems.textsecuregcm.storage.UsernamesManager; import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.util.AsnManager; import org.whispersystems.textsecuregcm.util.AsnManager;
import org.whispersystems.textsecuregcm.util.Constants; import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
import org.whispersystems.textsecuregcm.util.TorExitNodeManager; import org.whispersystems.textsecuregcm.util.TorExitNodeManager;
import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener; import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener;
import org.whispersystems.textsecuregcm.websocket.DeadLetterHandler; import org.whispersystems.textsecuregcm.websocket.DeadLetterHandler;
@ -210,6 +206,14 @@ import org.whispersystems.textsecuregcm.workers.VacuumCommand;
import org.whispersystems.textsecuregcm.workers.ZkParamsCommand; import org.whispersystems.textsecuregcm.workers.ZkParamsCommand;
import org.whispersystems.websocket.WebSocketResourceProviderFactory; import org.whispersystems.websocket.WebSocketResourceProviderFactory;
import org.whispersystems.websocket.setup.WebSocketEnvironment; import org.whispersystems.websocket.setup.WebSocketEnvironment;
import software.amazon.awssdk.core.client.config.ClientAsyncConfiguration;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
import software.amazon.awssdk.http.SdkHttpClient;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder;
public class WhisperServerService extends Application<WhisperServerConfiguration> { public class WhisperServerService extends Application<WhisperServerConfiguration> {
@ -316,80 +320,40 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("accounts_database", accountJdbi, config.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration()); FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("accounts_database", accountJdbi, config.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration());
FaultTolerantDatabase abuseDatabase = new FaultTolerantDatabase("abuse_database", abuseJdbi, config.getAbuseDatabaseConfiguration().getCircuitBreakerConfiguration()); FaultTolerantDatabase abuseDatabase = new FaultTolerantDatabase("abuse_database", abuseJdbi, config.getAbuseDatabaseConfiguration().getCircuitBreakerConfiguration());
AmazonDynamoDBClientBuilder messageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient messageDynamoDb = DynamoDbFromConfig.client(config.getMessageDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder keysDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient preKeyDynamoDb = DynamoDbFromConfig.client(config.getKeysDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getKeysDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getKeysDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getKeysDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder accountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient accountsDynamoDbClient = DynamoDbFromConfig.client(config.getAccountsDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
// The thread pool core & max sizes are set via dynamic configuration within AccountsDynamoDb // The thread pool core & max sizes are set via dynamic configuration within AccountsDynamoDb
ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS,
new LinkedBlockingDeque<>()); new LinkedBlockingDeque<>());
AmazonDynamoDBAsyncClientBuilder accountsDynamoDbAsyncClientBuilder = AmazonDynamoDBAsyncClientBuilder DynamoDbAsyncClient accountsDynamoDbAsyncClient = DynamoDbFromConfig.asyncClient(config.getAccountsDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create(),
.withRegion(accountsDynamoDbClientBuilder.getRegion()) accountsDynamoDbMigrationThreadPool);
.withClientConfiguration(accountsDynamoDbClientBuilder.getClientConfiguration())
.withCredentials(accountsDynamoDbClientBuilder.getCredentials())
.withExecutorFactory(() -> accountsDynamoDbMigrationThreadPool);
AmazonDynamoDBClientBuilder migrationDeletedAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient recentlyDeletedAccountsDynamoDb = DynamoDbFromConfig.client(config.getMigrationDeletedAccountsDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getMigrationDeletedAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMigrationDeletedAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMigrationDeletedAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder migrationRetryAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient pushChallengeDynamoDbClient = DynamoDbFromConfig.client(config.getPushChallengeDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getMigrationRetryAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMigrationRetryAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMigrationRetryAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder pushChallengeDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient reportMessageDynamoDbClient = DynamoDbFromConfig.client(config.getReportMessageDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getPushChallengeDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getPushChallengeDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getPushChallengeDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder reportMessageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder DynamoDbClient migrationRetryAccountsDynamoDb = DynamoDbFromConfig.client(config.getMigrationRetryAccountsDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(config.getReportMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getReportMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getReportMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDB messageDynamoDb = new DynamoDB(messageDynamoDbClientBuilder.build());
DynamoDB preKeyDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
AmazonDynamoDBAsync accountsDynamodbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
DynamoDB recentlyDeletedAccountsDynamoDb = new DynamoDB(migrationDeletedAccountsDynamoDbClientBuilder.build());
DynamoDB migrationRetryAccountsDynamoDb = new DynamoDB(migrationRetryAccountsDynamoDbClientBuilder.build());
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(recentlyDeletedAccountsDynamoDb, config.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName()); MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(recentlyDeletedAccountsDynamoDb, config.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, config.getMigrationRetryAccountsDynamoDbConfiguration().getTableName()); MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, config.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
Accounts accounts = new Accounts(accountDatabase); Accounts accounts = new Accounts(accountDatabase);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamodbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts); AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase); PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase);
PendingDevices pendingDevices = new PendingDevices (accountDatabase); PendingDevices pendingDevices = new PendingDevices (accountDatabase);
Usernames usernames = new Usernames(accountDatabase); Usernames usernames = new Usernames(accountDatabase);
@ -399,8 +363,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, config.getMessageDynamoDbConfiguration().getTableName(), config.getMessageDynamoDbConfiguration().getTimeToLive()); MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, config.getMessageDynamoDbConfiguration().getTableName(), config.getMessageDynamoDbConfiguration().getTimeToLive());
AbusiveHostRules abusiveHostRules = new AbusiveHostRules(abuseDatabase); AbusiveHostRules abusiveHostRules = new AbusiveHostRules(abuseDatabase);
RemoteConfigs remoteConfigs = new RemoteConfigs(accountDatabase); RemoteConfigs remoteConfigs = new RemoteConfigs(accountDatabase);
PushChallengeDynamoDb pushChallengeDynamoDb = new PushChallengeDynamoDb(new DynamoDB(pushChallengeDynamoDbClientBuilder.build()), config.getPushChallengeDynamoDbConfiguration().getTableName()); PushChallengeDynamoDb pushChallengeDynamoDb = new PushChallengeDynamoDb(pushChallengeDynamoDbClient, config.getPushChallengeDynamoDbConfiguration().getTableName());
ReportMessageDynamoDb reportMessageDynamoDb = new ReportMessageDynamoDb(new DynamoDB(reportMessageDynamoDbClientBuilder.build()), config.getReportMessageDynamoDbConfiguration().getTableName()); ReportMessageDynamoDb reportMessageDynamoDb = new ReportMessageDynamoDb(reportMessageDynamoDbClient, config.getReportMessageDynamoDbConfiguration().getTableName());
RedisClientFactory pubSubClientFactory = new RedisClientFactory("pubsub_cache", config.getPubsubCacheConfiguration().getUrl(), config.getPubsubCacheConfiguration().getReplicaUrls(), config.getPubsubCacheConfiguration().getCircuitBreakerConfiguration()); RedisClientFactory pubSubClientFactory = new RedisClientFactory("pubsub_cache", config.getPubsubCacheConfiguration().getUrl(), config.getPubsubCacheConfiguration().getReplicaUrls(), config.getPubsubCacheConfiguration().getCircuitBreakerConfiguration());
ReplicatedJedisPool pubsubClient = pubSubClientFactory.getRedisClientPool(); ReplicatedJedisPool pubsubClient = pubSubClientFactory.getRedisClientPool();

58
service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java
View File

@ -5,21 +5,18 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.BatchWriteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.QueryOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Timer; import io.micrometer.core.instrument.Timer;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer; import java.util.function.Consumer;
@ -29,7 +26,7 @@ import static io.micrometer.core.instrument.Metrics.timer;
public class AbstractDynamoDbStore { public class AbstractDynamoDbStore {
private final DynamoDB dynamoDb; private final DynamoDbClient dynamoDbClient;
private final Timer batchWriteItemsFirstPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "true"); private final Timer batchWriteItemsFirstPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "true");
private final Timer batchWriteItemsRetryPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "false"); private final Timer batchWriteItemsRetryPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "false");
@ -41,44 +38,31 @@ public class AbstractDynamoDbStore {
public static final int DYNAMO_DB_MAX_BATCH_SIZE = 25; // This limit comes from Amazon Dynamo DB itself. It will reject batch writes larger than this. public static final int DYNAMO_DB_MAX_BATCH_SIZE = 25; // This limit comes from Amazon Dynamo DB itself. It will reject batch writes larger than this.
public static final int RESULT_SET_CHUNK_SIZE = 100; public static final int RESULT_SET_CHUNK_SIZE = 100;
public AbstractDynamoDbStore(final DynamoDB dynamoDb) { public AbstractDynamoDbStore(final DynamoDbClient dynamoDbClient) {
this.dynamoDb = dynamoDb; this.dynamoDbClient = dynamoDbClient;
} }
protected DynamoDB getDynamoDb() { protected DynamoDbClient db() {
return dynamoDb; return dynamoDbClient;
} }
protected void executeTableWriteItemsUntilComplete(final TableWriteItems items) { protected void executeTableWriteItemsUntilComplete(final Map<String,List<WriteRequest>> items) {
AtomicReference<BatchWriteItemOutcome> outcome = new AtomicReference<>(); AtomicReference<BatchWriteItemResponse> outcome = new AtomicReference<>();
batchWriteItemsFirstPass.record(() -> outcome.set(dynamoDb.batchWriteItem(items))); batchWriteItemsFirstPass.record(() -> outcome.set(dynamoDbClient.batchWriteItem(BatchWriteItemRequest.builder().requestItems(items).build())));
int attemptCount = 0; int attemptCount = 0;
while (!outcome.get().getUnprocessedItems().isEmpty() && attemptCount < MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE) { while (!outcome.get().unprocessedItems().isEmpty() && attemptCount < MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE) {
batchWriteItemsRetryPass.record(() -> outcome.set(dynamoDb.batchWriteItemUnprocessed(outcome.get().getUnprocessedItems()))); batchWriteItemsRetryPass.record(() -> outcome.set(dynamoDbClient.batchWriteItem(BatchWriteItemRequest.builder()
.requestItems(outcome.get().unprocessedItems())
.build())));
++attemptCount; ++attemptCount;
} }
if (!outcome.get().getUnprocessedItems().isEmpty()) { if (!outcome.get().unprocessedItems().isEmpty()) {
logger.error("Attempt count ({}) reached max ({}}) before applying all batch writes to dynamo. {} unprocessed items remain.", attemptCount, MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE, outcome.get().getUnprocessedItems().size()); int totalItems = outcome.get().unprocessedItems().values().stream().mapToInt(List::size).sum();
batchWriteItemsUnprocessed.increment(outcome.get().getUnprocessedItems().size()); logger.error("Attempt count ({}) reached max ({}}) before applying all batch writes to dynamo. {} unprocessed items remain.", attemptCount, MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE, totalItems);
batchWriteItemsUnprocessed.increment(totalItems);
} }
} }
protected long countItemsMatchingQuery(final Table table, final QuerySpec querySpec) {
// This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
// - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count
long matchingItems = 0;
for (final Page<Item, QueryOutcome> page : table.query(querySpec).pages()) {
matchingItems += page.getLowLevelResult().getQueryResult().getCount();
}
return matchingItems;
}
static <T> void writeInBatches(final Iterable<T> items, final Consumer<List<T>> action) { static <T> void writeInBatches(final Iterable<T> items, final Consumer<List<T>> action) {
final List<T> batch = new ArrayList<>(DYNAMO_DB_MAX_BATCH_SIZE); final List<T> batch = new ArrayList<>(DYNAMO_DB_MAX_BATCH_SIZE);

286
service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java
View File

@ -2,25 +2,6 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name; import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.handlers.AsyncHandler;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec;
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
import com.amazonaws.services.dynamodbv2.model.CancellationReason;
import com.amazonaws.services.dynamodbv2.model.Delete;
import com.amazonaws.services.dynamodbv2.model.GetItemResult;
import com.amazonaws.services.dynamodbv2.model.Put;
import com.amazonaws.services.dynamodbv2.model.ReturnValuesOnConditionCheckFailure;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItem;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItemsRequest;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItemsResult;
import com.amazonaws.services.dynamodbv2.model.TransactionCanceledException;
import com.amazonaws.services.dynamodbv2.model.UpdateItemRequest;
import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.Counter;
@ -33,12 +14,27 @@ import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.ThreadPoolExecutor;
import java.util.stream.Collectors; import java.util.stream.Collectors;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.SystemMapper; import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
import software.amazon.awssdk.services.dynamodb.model.Delete;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.Put;
import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountStore { public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountStore {
@ -51,9 +47,8 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
static final String ATTR_MIGRATION_VERSION = "V"; static final String ATTR_MIGRATION_VERSION = "V";
private final AmazonDynamoDB client; private final DynamoDbClient client;
private final Table accountsTable; private final DynamoDbAsyncClient asyncClient;
private final AmazonDynamoDBAsync asyncClient;
private final ThreadPoolExecutor migrationThreadPool; private final ThreadPoolExecutor migrationThreadPool;
@ -61,6 +56,7 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
private final MigrationRetryAccounts migrationRetryAccounts; private final MigrationRetryAccounts migrationRetryAccounts;
private final String phoneNumbersTableName; private final String phoneNumbersTableName;
private final String accountsTableName;
private static final Timer CREATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "create")); private static final Timer CREATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "create"));
private static final Timer UPDATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "update")); private static final Timer UPDATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "update"));
@ -70,18 +66,17 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
private final Logger logger = LoggerFactory.getLogger(AccountsDynamoDb.class); private final Logger logger = LoggerFactory.getLogger(AccountsDynamoDb.class);
public AccountsDynamoDb(AmazonDynamoDB client, AmazonDynamoDBAsync asyncClient, public AccountsDynamoDb(DynamoDbClient client, DynamoDbAsyncClient asyncClient,
ThreadPoolExecutor migrationThreadPool, DynamoDB dynamoDb, String accountsTableName, String phoneNumbersTableName, ThreadPoolExecutor migrationThreadPool, String accountsTableName, String phoneNumbersTableName,
MigrationDeletedAccounts migrationDeletedAccounts, MigrationDeletedAccounts migrationDeletedAccounts,
MigrationRetryAccounts accountsMigrationErrors) { MigrationRetryAccounts accountsMigrationErrors) {
super(dynamoDb); super(client);
this.client = client; this.client = client;
this.accountsTable = dynamoDb.getTable(accountsTableName);
this.phoneNumbersTableName = phoneNumbersTableName;
this.asyncClient = asyncClient; this.asyncClient = asyncClient;
this.phoneNumbersTableName = phoneNumbersTableName;
this.accountsTableName = accountsTableName;
this.migrationThreadPool = migrationThreadPool; this.migrationThreadPool = migrationThreadPool;
this.migrationDeletedAccounts = migrationDeletedAccounts; this.migrationDeletedAccounts = migrationDeletedAccounts;
@ -90,32 +85,34 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
@Override @Override
public boolean create(Account account) { public boolean create(Account account) {
return CREATE_TIMER.record(() -> { return CREATE_TIMER.record(() -> {
try { try {
TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid()); TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid(), Put.builder()
.conditionExpression("attribute_not_exists(#number) OR #number = :number")
.expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
.expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber()))));
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid()); final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.transactItems(phoneNumberConstraintPut, accountPut)
final TransactWriteItemsRequest request = new TransactWriteItemsRequest() .build();
.withTransactItems(phoneNumberConstraintPut, accountPut);
try { try {
client.transactWriteItems(request); client.transactWriteItems(request);
} catch (TransactionCanceledException e) { } catch (TransactionCanceledException e) {
final CancellationReason accountCancellationReason = e.getCancellationReasons().get(1); final CancellationReason accountCancellationReason = e.cancellationReasons().get(1);
if ("ConditionalCheckFailed".equals(accountCancellationReason.getCode())) { if ("ConditionalCheckFailed".equals(accountCancellationReason.code())) {
throw new IllegalArgumentException("uuid present with different phone number"); throw new IllegalArgumentException("uuid present with different phone number");
} }
final CancellationReason phoneNumberConstraintCancellationReason = e.getCancellationReasons().get(0); final CancellationReason phoneNumberConstraintCancellationReason = e.cancellationReasons().get(0);
if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.getCode())) { if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.code())) {
ByteBuffer actualAccountUuid = phoneNumberConstraintCancellationReason.getItem().get(KEY_ACCOUNT_UUID).getB(); ByteBuffer actualAccountUuid = phoneNumberConstraintCancellationReason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer();
account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid)); account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid));
update(account); update(account);
@ -134,39 +131,37 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
}); });
} }
private TransactWriteItem buildPutWriteItemForAccount(Account account, UUID uuid) throws JsonProcessingException { private TransactWriteItem buildPutWriteItemForAccount(Account account, UUID uuid, Put.Builder putBuilder) throws JsonProcessingException {
return new TransactWriteItem() return TransactWriteItem.builder()
.withPut( .put(putBuilder
new Put() .tableName(accountsTableName)
.withTableName(accountsTable.getTableName()) .item(Map.of(
.withItem(Map.of( KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid)), ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber()), ATTR_ACCOUNT_DATA, AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
ATTR_ACCOUNT_DATA, new AttributeValue() ATTR_MIGRATION_VERSION, AttributeValues.fromInt(account.getDynamoDbMigrationVersion())))
.withB(ByteBuffer.wrap(SystemMapper.getMapper().writeValueAsBytes(account))), .build())
ATTR_MIGRATION_VERSION, new AttributeValue().withN( .build();
String.valueOf(account.getDynamoDbMigrationVersion()))))
.withConditionExpression("attribute_not_exists(#number) OR #number = :number")
.withExpressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
.withExpressionAttributeValues(Map.of(":number", new AttributeValue(account.getNumber()))));
} }
private TransactWriteItem buildPutWriteItemForPhoneNumberConstraint(Account account, UUID uuid) { private TransactWriteItem buildPutWriteItemForPhoneNumberConstraint(Account account, UUID uuid) {
return new TransactWriteItem() return TransactWriteItem.builder()
.withPut( .put(
new Put() Put.builder()
.withTableName(phoneNumbersTableName) .tableName(phoneNumbersTableName)
.withItem(Map.of( .item(Map.of(
ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber()), ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid)))) KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.withConditionExpression( .conditionExpression(
"attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)") "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
.withExpressionAttributeNames( .expressionAttributeNames(
Map.of("#uuid", KEY_ACCOUNT_UUID, Map.of("#uuid", KEY_ACCOUNT_UUID,
"#number", ATTR_ACCOUNT_E164)) "#number", ATTR_ACCOUNT_E164))
.withExpressionAttributeValues( .expressionAttributeValues(
Map.of(":uuid", new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid)))) Map.of(":uuid", AttributeValues.fromUUID(uuid)))
.withReturnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)); .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
.build())
.build();
} }
@Override @Override
@ -174,16 +169,18 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
UPDATE_TIMER.record(() -> { UPDATE_TIMER.record(() -> {
UpdateItemRequest updateItemRequest; UpdateItemRequest updateItemRequest;
try { try {
updateItemRequest = new UpdateItemRequest() updateItemRequest = UpdateItemRequest.builder()
.withTableName(accountsTable.getTableName()) .tableName(accountsTableName)
.withKey(Map.of(KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(account.getUuid())))) .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
.withUpdateExpression("SET #data = :data, #version = :version") .updateExpression("SET #data = :data, #version = :version")
.withConditionExpression("attribute_exists(#number)") .conditionExpression("attribute_exists(#number)")
.withExpressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164, .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
"#data", ATTR_ACCOUNT_DATA, "#data", ATTR_ACCOUNT_DATA,
"#version", ATTR_MIGRATION_VERSION)) "#version", ATTR_MIGRATION_VERSION))
.withExpressionAttributeValues(Map.of(":data", new AttributeValue().withB(ByteBuffer.wrap(SystemMapper.getMapper().writeValueAsBytes(account))), .expressionAttributeValues(Map.of(
":version", new AttributeValue().withN(String.valueOf(account.getDynamoDbMigrationVersion())))); ":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
":version", AttributeValues.fromInt(account.getDynamoDbMigrationVersion())))
.build();
} catch (JsonProcessingException e) { } catch (JsonProcessingException e) {
throw new IllegalArgumentException(e); throw new IllegalArgumentException(e);
@ -193,37 +190,42 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
}); });
} }
@Override @Override
public Optional<Account> get(String number) { public Optional<Account> get(String number) {
return GET_BY_NUMBER_TIMER.record(() -> { return GET_BY_NUMBER_TIMER.record(() -> {
final GetItemResult phoneNumberAndUuid = client.getItem(phoneNumbersTableName, final GetItemResponse response = client.getItem(GetItemRequest.builder()
Map.of(ATTR_ACCOUNT_E164, new AttributeValue(number)), true); .tableName(phoneNumbersTableName)
.key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
.build());
return Optional.ofNullable(phoneNumberAndUuid.getItem()) return Optional.ofNullable(response.item())
.map(item -> item.get(KEY_ACCOUNT_UUID).getB()) .map(item -> item.get(KEY_ACCOUNT_UUID))
.map(uuid -> accountsTable.getItem(new GetItemSpec() .map(uuid -> accountByUuid(uuid))
.withPrimaryKey(KEY_ACCOUNT_UUID, uuid.array())
.withConsistentRead(true)))
.map(AccountsDynamoDb::fromItem); .map(AccountsDynamoDb::fromItem);
}); });
} }
private Map<String, AttributeValue> accountByUuid(AttributeValue uuid) {
GetItemResponse r = client.getItem(GetItemRequest.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, uuid))
.consistentRead(true)
.build());
return r.item().isEmpty() ? null : r.item();
}
@Override @Override
public Optional<Account> get(UUID uuid) { public Optional<Account> get(UUID uuid) {
Optional<Item> maybeItem = GET_BY_UUID_TIMER.record(() -> return GET_BY_UUID_TIMER.record(() ->
Optional.ofNullable(accountsTable.getItem(new GetItemSpec(). Optional.ofNullable(accountByUuid(AttributeValues.fromUUID(uuid)))
withPrimaryKey(new PrimaryKey(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(uuid))) .map(AccountsDynamoDb::fromItem));
.withConsistentRead(true))));
return maybeItem.map(AccountsDynamoDb::fromItem);
} }
@Override @Override
public void delete(UUID uuid) { public void delete(UUID uuid) {
DELETE_TIMER.record(() -> { DELETE_TIMER.record(() -> {
delete(uuid, true); delete(uuid, true);
}); });
} }
@ -238,18 +240,22 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
maybeAccount.ifPresent(account -> { maybeAccount.ifPresent(account -> {
TransactWriteItem phoneNumberDelete = new TransactWriteItem() TransactWriteItem phoneNumberDelete = TransactWriteItem.builder()
.withDelete(new Delete() .delete(Delete.builder()
.withTableName(phoneNumbersTableName) .tableName(phoneNumbersTableName)
.withKey(Map.of(ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber())))); .key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())))
.build())
.build();
TransactWriteItem accountDelete = new TransactWriteItem().withDelete( TransactWriteItem accountDelete = TransactWriteItem.builder()
new Delete() .delete(Delete.builder()
.withTableName(accountsTable.getTableName()) .tableName(accountsTableName)
.withKey(Map.of(KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid))))); .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.build())
.build();
TransactWriteItemsRequest request = new TransactWriteItemsRequest() TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.withTransactItems(phoneNumberDelete, accountDelete); .transactItems(phoneNumberDelete, accountDelete).build();
client.transactWriteItems(request); client.transactWriteItems(request);
}); });
@ -299,64 +305,62 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
try { try {
TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid()); TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid()); TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid(), Put.builder()
accountPut.getPut() .conditionExpression("attribute_not_exists(#uuid) OR (attribute_exists(#uuid) AND #version < :version)")
.setConditionExpression("attribute_not_exists(#uuid) OR (attribute_exists(#uuid) AND #version < :version)"); .expressionAttributeNames(Map.of(
accountPut.getPut() "#uuid", KEY_ACCOUNT_UUID,
.setExpressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#version", ATTR_MIGRATION_VERSION))
"#version", ATTR_MIGRATION_VERSION)); .expressionAttributeValues(Map.of(
accountPut.getPut() ":version", AttributeValues.fromInt(account.getDynamoDbMigrationVersion()))));
.setExpressionAttributeValues(
Map.of(":version", new AttributeValue().withN(String.valueOf(account.getDynamoDbMigrationVersion()))));
final TransactWriteItemsRequest request = new TransactWriteItemsRequest() final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.withTransactItems(phoneNumberConstraintPut, accountPut); .transactItems(phoneNumberConstraintPut, accountPut).build();
final CompletableFuture<Boolean> resultFuture = new CompletableFuture<>(); final CompletableFuture<Boolean> resultFuture = new CompletableFuture<>();
asyncClient.transactWriteItems(request).whenCompleteAsync((result, exception) -> {
asyncClient.transactWriteItemsAsync(request, if (result != null) {
new AsyncHandler<>() { resultFuture.complete(true);
@Override return;
public void onError(Exception exception) { }
if (exception instanceof TransactionCanceledException) { if (exception instanceof CompletionException) {
// account is already migrated // whenCompleteAsync can wrap exceptions in a CompletionException; unwrap it to get to the root cause.
resultFuture.complete(false); exception = exception.getCause();
} else { }
try { if (exception instanceof TransactionCanceledException) {
migrationRetryAccounts.put(account.getUuid()); // account is already migrated
} catch (final Exception e) { resultFuture.complete(false);
logger.error("Could not store account {}", account.getUuid()); return;
} }
resultFuture.completeExceptionally(exception); try {
} migrationRetryAccounts.put(account.getUuid());
} } catch (final Exception e) {
logger.error("Could not store account {}", account.getUuid());
@Override }
public void onSuccess(TransactWriteItemsRequest request, TransactWriteItemsResult transactWriteItemsResult) { resultFuture.completeExceptionally(exception);
resultFuture.complete(true); });
}
});
return resultFuture; return resultFuture;
} catch (Exception e) { } catch (Exception e) {
return CompletableFuture.failedFuture(e); return CompletableFuture.failedFuture(e);
} }
} }
private static String extractCancellationReasonCodes(final TransactionCanceledException exception) { private static String extractCancellationReasonCodes(final TransactionCanceledException exception) {
return exception.getCancellationReasons().stream() return exception.cancellationReasons().stream()
.map(CancellationReason::getCode) .map(CancellationReason::code)
.collect(Collectors.joining(", ")); .collect(Collectors.joining(", "));
} }
@VisibleForTesting @VisibleForTesting
static Account fromItem(Item item) { static Account fromItem(Map<String, AttributeValue> item) {
if (!item.containsKey(ATTR_ACCOUNT_DATA) ||
!item.containsKey(ATTR_ACCOUNT_E164) ||
!item.containsKey(KEY_ACCOUNT_UUID)) {
throw new RuntimeException("item missing values");
}
try { try {
Account account = SystemMapper.getMapper().readValue(item.getBinary(ATTR_ACCOUNT_DATA), Account.class); Account account = SystemMapper.getMapper().readValue(item.get(ATTR_ACCOUNT_DATA).b().asByteArray(), Account.class);
account.setNumber(item.get(ATTR_ACCOUNT_E164).s());
account.setNumber(item.getString(ATTR_ACCOUNT_E164)); account.setUuid(UUIDUtil.fromByteBuffer(item.get(KEY_ACCOUNT_UUID).b().asByteBuffer()));
account.setUuid(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_ACCOUNT_UUID)));
return account; return account;

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

@ -7,7 +7,6 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name; import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import com.codahale.metrics.MetricRegistry; import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SharedMetricRegistries; import com.codahale.metrics.SharedMetricRegistries;
import com.codahale.metrics.Timer; import com.codahale.metrics.Timer;
@ -42,6 +41,7 @@ import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.util.Constants; import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.SystemMapper; import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.Util; import org.whispersystems.textsecuregcm.util.Util;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
public class AccountsManager { public class AccountsManager {

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

@ -7,195 +7,230 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name; import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import com.amazonaws.services.dynamodbv2.model.ReturnValue;
import com.amazonaws.services.dynamodbv2.model.Select;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import io.micrometer.core.instrument.DistributionSummary; import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.Metrics; import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Timer; import io.micrometer.core.instrument.Timer;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap; import java.util.HashMap;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.UUID; import java.util.UUID;
import org.whispersystems.textsecuregcm.entities.PreKey; import org.whispersystems.textsecuregcm.entities.PreKey;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import software.amazon.awssdk.services.dynamodb.model.Select;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class KeysDynamoDb extends AbstractDynamoDbStore { public class KeysDynamoDb extends AbstractDynamoDbStore {
private final Table table; private final String tableName;
static final String KEY_ACCOUNT_UUID = "U"; static final String KEY_ACCOUNT_UUID = "U";
static final String KEY_DEVICE_ID_KEY_ID = "DK"; static final String KEY_DEVICE_ID_KEY_ID = "DK";
static final String KEY_PUBLIC_KEY = "P"; static final String KEY_PUBLIC_KEY = "P";
private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(KeysDynamoDb.class, "storeKeys")); private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(KeysDynamoDb.class, "storeKeys"));
private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice")); private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice"));
private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForAccount")); private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForAccount"));
private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "getKeyCount")); private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "getKeyCount"));
private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForDevice")); private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForDevice"));
private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForAccount")); private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForAccount"));
private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "contestedKeys")); private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "contestedKeys"));
private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "keyCount")); private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "keyCount"));
public KeysDynamoDb(final DynamoDB dynamoDB, final String tableName) { public KeysDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
super(dynamoDB); super(dynamoDB);
this.tableName = tableName;
}
this.table = dynamoDB.getTable(tableName); public void store(final Account account, final long deviceId, final List<PreKey> keys) {
} STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId);
public void store(final Account account, final long deviceId, final List<PreKey> keys) { writeInBatches(keys, batch -> {
STORE_KEYS_TIMER.record(() -> { List<WriteRequest> items = new ArrayList<>();
delete(account, deviceId); for (final PreKey preKey : batch) {
items.add(WriteRequest.builder()
.putRequest(PutRequest.builder()
.item(getItemFromPreKey(account.getUuid(), deviceId, preKey))
.build())
.build());
}
executeTableWriteItemsUntilComplete(Map.of(tableName, items));
});
});
}
writeInBatches(keys, batch -> { public Optional<PreKey> take(final Account account, final long deviceId) {
final TableWriteItems items = new TableWriteItems(table.getTableName()); return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> {
final AttributeValue partitionKey = getPartitionKey(account.getUuid());
QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", partitionKey,
":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(false)
.build();
for (final PreKey preKey : batch) { int contestedKeys = 0;
items.addItemToPut(getItemFromPreKey(account.getUuid(), deviceId, preKey));
}
executeTableWriteItemsUntilComplete(items); try {
}); QueryResponse response = db().query(queryRequest);
}); for (Map<String, AttributeValue> candidate : response.items()) {
} DeleteItemRequest deleteItemRequest = DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(
KEY_ACCOUNT_UUID, partitionKey,
KEY_DEVICE_ID_KEY_ID, candidate.get(KEY_DEVICE_ID_KEY_ID)))
.returnValues(ReturnValue.ALL_OLD)
.build();
DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest);
if (deleteItemResponse.attributes() != null) {
return Optional.of(getPreKeyFromItem(deleteItemResponse.attributes()));
}
public Optional<PreKey> take(final Account account, final long deviceId) { contestedKeys++;
return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> { }
final byte[] partitionKey = getPartitionKey(account.getUuid());
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)") return Optional.empty();
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID)) } finally {
.withValueMap(Map.of(":uuid", partitionKey, CONTESTED_KEY_DISTRIBUTION.record(contestedKeys);
":sortprefix", getSortKeyPrefix(deviceId))) }
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID) });
.withConsistentRead(false); }
int contestedKeys = 0; public Map<Long, PreKey> take(final Account account) {
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final Map<Long, PreKey> preKeysByDeviceId = new HashMap<>();
try { for (final Device device : account.getDevices()) {
for (final Item candidate : table.query(querySpec)) { take(account, device.getId()).ifPresent(preKey -> preKeysByDeviceId.put(device.getId(), preKey));
final DeleteItemSpec deleteItemSpec = new DeleteItemSpec().withPrimaryKey(KEY_ACCOUNT_UUID, partitionKey, KEY_DEVICE_ID_KEY_ID, candidate.getBinary(KEY_DEVICE_ID_KEY_ID)) }
.withReturnValues(ReturnValue.ALL_OLD);
final DeleteItemOutcome outcome = table.deleteItem(deleteItemSpec); return preKeysByDeviceId;
});
}
if (outcome.getItem() != null) { public int getCount(final Account account, final long deviceId) {
return Optional.of(getPreKeyFromItem(outcome.getItem())); return GET_KEY_COUNT_TIMER.record(() -> {
} QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.select(Select.COUNT)
.consistentRead(false)
.build();
contestedKeys++; int keyCount = 0;
} // This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
// - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count
for (final QueryResponse page : db().queryPaginator(queryRequest)) {
keyCount += page.count();
}
KEY_COUNT_DISTRIBUTION.record(keyCount);
return keyCount;
});
}
return Optional.empty(); public void delete(final Account account) {
} finally { DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
CONTESTED_KEY_DISTRIBUTION.record(contestedKeys); final QueryRequest queryRequest = QueryRequest.builder()
} .tableName(tableName)
}); .keyConditionExpression("#uuid = :uuid")
} .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid())))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
public Map<Long, PreKey> take(final Account account) { deleteItemsForAccountMatchingQuery(account, queryRequest);
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { });
final Map<Long, PreKey> preKeysByDeviceId = new HashMap<>(); }
for (final Device device : account.getDevices()) { @VisibleForTesting
take(account, device.getId()).ifPresent(preKey -> preKeysByDeviceId.put(device.getId(), preKey)); void delete(final Account account, final long deviceId) {
} DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
return preKeysByDeviceId; deleteItemsForAccountMatchingQuery(account, queryRequest);
}); });
} }
public int getCount(final Account account, final long deviceId) { private void deleteItemsForAccountMatchingQuery(final Account account, final QueryRequest querySpec) {
return GET_KEY_COUNT_TIMER.record(() -> { final AttributeValue partitionKey = getPartitionKey(account.getUuid());
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.withSelect(Select.COUNT)
.withConsistentRead(false);
final int keyCount = (int)countItemsMatchingQuery(table, querySpec); writeInBatches(db().query(querySpec).items(), batch -> {
List<WriteRequest> deletes = new ArrayList<>();
for (final Map<String, AttributeValue> item : batch) {
deletes.add(WriteRequest.builder()
.deleteRequest(DeleteRequest.builder()
.key(Map.of(
KEY_ACCOUNT_UUID, partitionKey,
KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)))
.build())
.build());
}
executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
});
}
KEY_COUNT_DISTRIBUTION.record(keyCount); private static AttributeValue getPartitionKey(final UUID accountUuid) {
return keyCount; return AttributeValues.fromUUID(accountUuid);
}); }
}
public void delete(final Account account) { private static AttributeValue getSortKey(final long deviceId, final long keyId) {
DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid") byteBuffer.putLong(deviceId);
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID)) byteBuffer.putLong(keyId);
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid()))) return AttributeValues.fromByteBuffer(byteBuffer.flip());
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID) }
.withConsistentRead(true);
deleteItemsForAccountMatchingQuery(account, querySpec); @VisibleForTesting
}); static AttributeValue getSortKeyPrefix(final long deviceId) {
} final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(deviceId);
return AttributeValues.fromByteBuffer(byteBuffer.flip());
}
@VisibleForTesting private Map<String, AttributeValue> getItemFromPreKey(final UUID accountUuid, final long deviceId, final PreKey preKey) {
void delete(final Account account, final long deviceId) { return Map.of(
DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> { KEY_ACCOUNT_UUID, getPartitionKey(accountUuid),
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)") KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, preKey.getKeyId()),
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID)) KEY_PUBLIC_KEY, AttributeValues.fromString(preKey.getPublicKey()));
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid()), }
":sortprefix", getSortKeyPrefix(deviceId)))
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID)
.withConsistentRead(true);
deleteItemsForAccountMatchingQuery(account, querySpec); private PreKey getPreKeyFromItem(Map<String, AttributeValue> item) {
}); final long keyId = item.get(KEY_DEVICE_ID_KEY_ID).b().asByteBuffer().getLong(8);
} return new PreKey(keyId, item.get(KEY_PUBLIC_KEY).s());
}
private void deleteItemsForAccountMatchingQuery(final Account account, final QuerySpec querySpec) {
final byte[] partitionKey = getPartitionKey(account.getUuid());
writeInBatches(table.query(querySpec), batch -> {
final TableWriteItems writeItems = new TableWriteItems(table.getTableName());
for (final Item item : batch) {
writeItems.addPrimaryKeyToDelete(new PrimaryKey(KEY_ACCOUNT_UUID, partitionKey, KEY_DEVICE_ID_KEY_ID, item.getBinary(KEY_DEVICE_ID_KEY_ID)));
}
executeTableWriteItemsUntilComplete(writeItems);
});
}
private static byte[] getPartitionKey(final UUID accountUuid) {
return UUIDUtil.toBytes(accountUuid);
}
private static byte[] getSortKey(final long deviceId, final long keyId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
byteBuffer.putLong(deviceId);
byteBuffer.putLong(keyId);
return byteBuffer.array();
}
private static byte[] getSortKeyPrefix(final long deviceId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(deviceId);
return byteBuffer.array();
}
private Item getItemFromPreKey(final UUID accountUuid, final long deviceId, final PreKey preKey) {
return new Item().withBinary(KEY_ACCOUNT_UUID, getPartitionKey(accountUuid))
.withBinary(KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, preKey.getKeyId()))
.withString(KEY_PUBLIC_KEY, preKey.getPublicKey());
}
private PreKey getPreKeyFromItem(final Item item) {
final long keyId = ByteBuffer.wrap(item.getBinary(KEY_DEVICE_ID_KEY_ID)).getLong(8);
return new PreKey(keyId, item.getString(KEY_PUBLIC_KEY));
}
} }

256
service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
View File

@ -8,17 +8,7 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name; import static com.codahale.metrics.MetricRegistry.name;
import static io.micrometer.core.instrument.Metrics.timer; import static io.micrometer.core.instrument.Metrics.timer;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome; import com.google.common.collect.ImmutableMap;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Index;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.api.QueryApi;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import com.amazonaws.services.dynamodbv2.model.ReturnValue;
import io.micrometer.core.instrument.Timer; import io.micrometer.core.instrument.Timer;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.time.Duration; import java.time.Duration;
@ -27,11 +17,22 @@ import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.UUID; import java.util.UUID;
import java.util.stream.Collectors;
import javax.annotation.Nonnull; import javax.annotation.Nonnull;
import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.StringUtils;
import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity; import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class MessagesDynamoDb extends AbstractDynamoDbStore { public class MessagesDynamoDb extends AbstractDynamoDbStore {
@ -60,7 +61,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
private final String tableName; private final String tableName;
private final Duration timeToLive; private final Duration timeToLive;
public MessagesDynamoDb(DynamoDB dynamoDb, String tableName, Duration timeToLive) { public MessagesDynamoDb(DynamoDbClient dynamoDb, String tableName, Duration timeToLive) {
super(dynamoDb); super(dynamoDb);
this.tableName = tableName; this.tableName = tableName;
@ -76,54 +77,61 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
throw new IllegalArgumentException("Maximum batch size of " + DYNAMO_DB_MAX_BATCH_SIZE + " execeeded with " + messages.size() + " messages"); throw new IllegalArgumentException("Maximum batch size of " + DYNAMO_DB_MAX_BATCH_SIZE + " execeeded with " + messages.size() + " messages");
} }
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
TableWriteItems items = new TableWriteItems(tableName); List<WriteRequest> writeItems = new ArrayList<>();
for (MessageProtos.Envelope message : messages) { for (MessageProtos.Envelope message : messages) {
final UUID messageUuid = UUID.fromString(message.getServerGuid()); final UUID messageUuid = UUID.fromString(message.getServerGuid());
final Item item = new Item().withBinary(KEY_PARTITION, partitionKey) final ImmutableMap.Builder<String, AttributeValue> item = ImmutableMap.<String, AttributeValue>builder()
.withBinary(KEY_SORT, convertSortKey(destinationDeviceId, message.getServerTimestamp(), messageUuid)) .put(KEY_PARTITION, partitionKey)
.withBinary(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid)) .put(KEY_SORT, convertSortKey(destinationDeviceId, message.getServerTimestamp(), messageUuid))
.withInt(KEY_TYPE, message.getType().getNumber()) .put(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid))
.withLong(KEY_TIMESTAMP, message.getTimestamp()) .put(KEY_TYPE, AttributeValues.fromInt(message.getType().getNumber()))
.withLong(KEY_TTL, getTtlForMessage(message)); .put(KEY_TIMESTAMP, AttributeValues.fromLong(message.getTimestamp()))
.put(KEY_TTL, AttributeValues.fromLong(getTtlForMessage(message)));
if (message.hasRelay() && message.getRelay().length() > 0) { if (message.hasRelay() && message.getRelay().length() > 0) {
item.withString(KEY_RELAY, message.getRelay()); item.put(KEY_RELAY, AttributeValues.fromString(message.getRelay()));
} }
if (message.hasSource()) { if (message.hasSource()) {
item.withString(KEY_SOURCE, message.getSource()); item.put(KEY_SOURCE, AttributeValues.fromString(message.getSource()));
} }
if (message.hasSourceUuid()) { if (message.hasSourceUuid()) {
item.withBinary(KEY_SOURCE_UUID, UUIDUtil.toBytes(UUID.fromString(message.getSourceUuid()))); item.put(KEY_SOURCE_UUID, AttributeValues.fromUUID(UUID.fromString(message.getSourceUuid())));
} }
if (message.hasSourceDevice()) { if (message.hasSourceDevice()) {
item.withInt(KEY_SOURCE_DEVICE, message.getSourceDevice()); item.put(KEY_SOURCE_DEVICE, AttributeValues.fromInt(message.getSourceDevice()));
} }
if (message.hasLegacyMessage()) { if (message.hasLegacyMessage()) {
item.withBinary(KEY_MESSAGE, message.getLegacyMessage().toByteArray()); item.put(KEY_MESSAGE, AttributeValues.fromByteArray(message.getLegacyMessage().toByteArray()));
} }
if (message.hasContent()) { if (message.hasContent()) {
item.withBinary(KEY_CONTENT, message.getContent().toByteArray()); item.put(KEY_CONTENT, AttributeValues.fromByteArray(message.getContent().toByteArray()));
} }
items.addItemToPut(item); writeItems.add(WriteRequest.builder().putRequest(PutRequest.builder()
.item(item.build())
.build()).build());
} }
executeTableWriteItemsUntilComplete(items); executeTableWriteItemsUntilComplete(Map.of(tableName, writeItems));
} }
public List<OutgoingMessageEntity> load(final UUID destinationAccountUuid, final long destinationDeviceId, final int requestedNumberOfMessagesToFetch) { public List<OutgoingMessageEntity> load(final UUID destinationAccountUuid, final long destinationDeviceId, final int requestedNumberOfMessagesToFetch) {
return loadTimer.record(() -> { return loadTimer.record(() -> {
final int numberOfMessagesToFetch = Math.min(requestedNumberOfMessagesToFetch, RESULT_SET_CHUNK_SIZE); final int numberOfMessagesToFetch = Math.min(requestedNumberOfMessagesToFetch, RESULT_SET_CHUNK_SIZE);
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withConsistentRead(true) final QueryRequest queryRequest = QueryRequest.builder()
.withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )") .tableName(tableName)
.withNameMap(Map.of("#part", KEY_PARTITION, .consistentRead(true)
"#sort", KEY_SORT)) .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.withValueMap(Map.of(":part", partitionKey, .expressionAttributeNames(Map.of(
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId))) "#part", KEY_PARTITION,
.withMaxResultSize(numberOfMessagesToFetch); "#sort", KEY_SORT))
final Table table = getDynamoDb().getTable(tableName); .expressionAttributeValues(Map.of(
":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.limit(numberOfMessagesToFetch)
.build();
List<OutgoingMessageEntity> messageEntities = new ArrayList<>(numberOfMessagesToFetch); List<OutgoingMessageEntity> messageEntities = new ArrayList<>(numberOfMessagesToFetch);
for (Item message : table.query(querySpec)) { for (Map<String, AttributeValue> message : db().query(queryRequest).items()) {
messageEntities.add(convertItemToOutgoingMessageEntity(message)); messageEntities.add(convertItemToOutgoingMessageEntity(message));
} }
return messageEntities; return messageEntities;
@ -136,53 +144,63 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
throw new IllegalArgumentException("must specify a source"); throw new IllegalArgumentException("must specify a source");
} }
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withProjectionExpression(KEY_SORT) final QueryRequest queryRequest = QueryRequest.builder()
.withConsistentRead(true) .tableName(tableName)
.withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )") .projectionExpression(KEY_SORT)
.withFilterExpression("#source = :source AND #timestamp = :timestamp") .consistentRead(true)
.withNameMap(Map.of("#part", KEY_PARTITION, .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
"#sort", KEY_SORT, .filterExpression("#source = :source AND #timestamp = :timestamp")
"#source", KEY_SOURCE, .expressionAttributeNames(Map.of(
"#timestamp", KEY_TIMESTAMP)) "#part", KEY_PARTITION,
.withValueMap(Map.of(":part", partitionKey, "#sort", KEY_SORT,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId), "#source", KEY_SOURCE,
":source", source, "#timestamp", KEY_TIMESTAMP))
":timestamp", timestamp)); .expressionAttributeValues(Map.of(
":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId),
":source", AttributeValues.fromString(source),
":timestamp", AttributeValues.fromLong(timestamp)))
.build();
final Table table = getDynamoDb().getTable(tableName); return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(partitionKey, queryRequest);
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(table, partitionKey, querySpec, table);
}); });
} }
public Optional<OutgoingMessageEntity> deleteMessageByDestinationAndGuid(final UUID destinationAccountUuid, final long destinationDeviceId, final UUID messageUuid) { public Optional<OutgoingMessageEntity> deleteMessageByDestinationAndGuid(final UUID destinationAccountUuid, final long destinationDeviceId, final UUID messageUuid) {
return deleteByGuid.record(() -> { return deleteByGuid.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withProjectionExpression(KEY_SORT) final QueryRequest queryRequest = QueryRequest.builder()
.withConsistentRead(true) .tableName(tableName)
.withKeyConditionExpression("#part = :part AND #uuid = :uuid") .indexName(LOCAL_INDEX_MESSAGE_UUID_NAME)
.withNameMap(Map.of("#part", KEY_PARTITION, .projectionExpression(KEY_SORT)
"#uuid", LOCAL_INDEX_MESSAGE_UUID_KEY_SORT)) .consistentRead(true)
.withValueMap(Map.of(":part", partitionKey, .keyConditionExpression("#part = :part AND #uuid = :uuid")
":uuid", convertLocalIndexMessageUuidSortKey(messageUuid))); .expressionAttributeNames(Map.of(
final Table table = getDynamoDb().getTable(tableName); "#part", KEY_PARTITION,
final Index index = table.getIndex(LOCAL_INDEX_MESSAGE_UUID_NAME); "#uuid", LOCAL_INDEX_MESSAGE_UUID_KEY_SORT))
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(table, partitionKey, querySpec, index); .expressionAttributeValues(Map.of(
":part", partitionKey,
":uuid", convertLocalIndexMessageUuidSortKey(messageUuid)))
.build();
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(partitionKey, queryRequest);
}); });
} }
@Nonnull @Nonnull
private Optional<OutgoingMessageEntity> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(Table table, byte[] partitionKey, QuerySpec querySpec, QueryApi queryApi) { private Optional<OutgoingMessageEntity> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(AttributeValue partitionKey, QueryRequest queryRequest) {
Optional<OutgoingMessageEntity> result = Optional.empty(); Optional<OutgoingMessageEntity> result = Optional.empty();
for (Item item : queryApi.query(querySpec)) { for (Map<String, AttributeValue> item : db().query(queryRequest).items()) {
final byte[] rangeKeyValue = item.getBinary(KEY_SORT); final byte[] rangeKeyValue = item.get(KEY_SORT).b().asByteArray();
DeleteItemSpec deleteItemSpec = new DeleteItemSpec().withPrimaryKey(KEY_PARTITION, partitionKey, KEY_SORT, rangeKeyValue); DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT, AttributeValues.fromByteArray(rangeKeyValue)));
if (result.isEmpty()) { if (result.isEmpty()) {
deleteItemSpec.withReturnValues(ReturnValue.ALL_OLD); deleteItemRequest.returnValues(ReturnValue.ALL_OLD);
} }
final DeleteItemOutcome deleteItemOutcome = table.deleteItem(deleteItemSpec); final DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest.build());
if (deleteItemOutcome.getItem() != null && deleteItemOutcome.getItem().hasAttribute(KEY_PARTITION)) { if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
result = Optional.of(convertItemToOutgoingMessageEntity(deleteItemOutcome.getItem())); result = Optional.of(convertItemToOutgoingMessageEntity(deleteItemResponse.attributes()));
} }
} }
return result; return result;
@ -190,74 +208,88 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
public void deleteAllMessagesForAccount(final UUID destinationAccountUuid) { public void deleteAllMessagesForAccount(final UUID destinationAccountUuid) {
deleteByAccount.record(() -> { deleteByAccount.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withHashKey(KEY_PARTITION, partitionKey) final QueryRequest queryRequest = QueryRequest.builder()
.withProjectionExpression(KEY_SORT) .tableName(tableName)
.withConsistentRead(true); .projectionExpression(KEY_SORT)
deleteRowsMatchingQuery(partitionKey, querySpec); .consistentRead(true)
.keyConditionExpression("#part = :part")
.expressionAttributeNames(Map.of("#part", KEY_PARTITION))
.expressionAttributeValues(Map.of(":part", partitionKey))
.build();
deleteRowsMatchingQuery(partitionKey, queryRequest);
}); });
} }
public void deleteAllMessagesForDevice(final UUID destinationAccountUuid, final long destinationDeviceId) { public void deleteAllMessagesForDevice(final UUID destinationAccountUuid, final long destinationDeviceId) {
deleteByDevice.record(() -> { deleteByDevice.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid); final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )") final QueryRequest queryRequest = QueryRequest.builder()
.withNameMap(Map.of("#part", KEY_PARTITION, .tableName(tableName)
"#sort", KEY_SORT)) .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.withValueMap(Map.of(":part", partitionKey, .expressionAttributeNames(Map.of(
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId))) "#part", KEY_PARTITION,
.withProjectionExpression(KEY_SORT) "#sort", KEY_SORT))
.withConsistentRead(true); .expressionAttributeValues(Map.of(
deleteRowsMatchingQuery(partitionKey, querySpec); ":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.projectionExpression(KEY_SORT)
.consistentRead(true)
.build();
deleteRowsMatchingQuery(partitionKey, queryRequest);
}); });
} }
private OutgoingMessageEntity convertItemToOutgoingMessageEntity(Item message) { private OutgoingMessageEntity convertItemToOutgoingMessageEntity(Map<String, AttributeValue> message) {
final SortKey sortKey = convertSortKey(message.getBinary(KEY_SORT)); final SortKey sortKey = convertSortKey(message.get(KEY_SORT).b().asByteArray());
final UUID messageUuid = convertLocalIndexMessageUuidSortKey(message.getBinary(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT)); final UUID messageUuid = convertLocalIndexMessageUuidSortKey(message.get(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT).b().asByteArray());
final int type = message.getInt(KEY_TYPE); final int type = AttributeValues.getInt(message, KEY_TYPE, 0);
final String relay = message.getString(KEY_RELAY); final String relay = AttributeValues.getString(message, KEY_RELAY, null);
final long timestamp = message.getLong(KEY_TIMESTAMP); final long timestamp = AttributeValues.getLong(message, KEY_TIMESTAMP, 0L);
final String source = message.getString(KEY_SOURCE); final String source = AttributeValues.getString(message, KEY_SOURCE, null);
final UUID sourceUuid = message.hasAttribute(KEY_SOURCE_UUID) ? convertUuidFromBytes(message.getBinary(KEY_SOURCE_UUID), "message source uuid") : null; final UUID sourceUuid = AttributeValues.getUUID(message, KEY_SOURCE_UUID, null);
final int sourceDevice = message.hasAttribute(KEY_SOURCE_DEVICE) ? message.getInt(KEY_SOURCE_DEVICE) : 0; final int sourceDevice = AttributeValues.getInt(message, KEY_SOURCE_DEVICE, 0);
final byte[] messageBytes = message.getBinary(KEY_MESSAGE); final byte[] messageBytes = AttributeValues.getByteArray(message, KEY_MESSAGE, null);
final byte[] content = message.getBinary(KEY_CONTENT); final byte[] content = AttributeValues.getByteArray(message, KEY_CONTENT, null);
return new OutgoingMessageEntity(-1L, false, messageUuid, type, relay, timestamp, source, sourceUuid, sourceDevice, messageBytes, content, sortKey.getServerTimestamp()); return new OutgoingMessageEntity(-1L, false, messageUuid, type, relay, timestamp, source, sourceUuid, sourceDevice, messageBytes, content, sortKey.getServerTimestamp());
} }
private void deleteRowsMatchingQuery(byte[] partitionKey, QuerySpec querySpec) { private void deleteRowsMatchingQuery(AttributeValue partitionKey, QueryRequest querySpec) {
final Table table = getDynamoDb().getTable(tableName); writeInBatches(db().query(querySpec).items(), (itemBatch) -> deleteItems(partitionKey, itemBatch));
writeInBatches(table.query(querySpec), (itemBatch) -> deleteItems(partitionKey, itemBatch));
} }
private void deleteItems(byte[] partitionKey, List<Item> items) { private void deleteItems(AttributeValue partitionKey, List<Map<String, AttributeValue>> items) {
final TableWriteItems tableWriteItems = new TableWriteItems(tableName); List<WriteRequest> deletes = items.stream()
items.stream().map(item -> new PrimaryKey(KEY_PARTITION, partitionKey, KEY_SORT, item.getBinary(KEY_SORT))).forEach(tableWriteItems::addPrimaryKeyToDelete); .map(item -> WriteRequest.builder()
executeTableWriteItemsUntilComplete(tableWriteItems); .deleteRequest(DeleteRequest.builder().key(Map.of(
KEY_PARTITION, partitionKey,
KEY_SORT, item.get(KEY_SORT))).build())
.build())
.collect(Collectors.toList());
executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
} }
private long getTtlForMessage(MessageProtos.Envelope message) { private long getTtlForMessage(MessageProtos.Envelope message) {
return message.getServerTimestamp() / 1000 + timeToLive.getSeconds(); return message.getServerTimestamp() / 1000 + timeToLive.getSeconds();
} }
private static byte[] convertPartitionKey(final UUID destinationAccountUuid) { private static AttributeValue convertPartitionKey(final UUID destinationAccountUuid) {
return UUIDUtil.toBytes(destinationAccountUuid); return AttributeValues.fromUUID(destinationAccountUuid);
} }
private static byte[] convertSortKey(final long destinationDeviceId, final long serverTimestamp, final UUID messageUuid) { private static AttributeValue convertSortKey(final long destinationDeviceId, final long serverTimestamp, final UUID messageUuid) {
ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[32]); ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[32]);
byteBuffer.putLong(destinationDeviceId); byteBuffer.putLong(destinationDeviceId);
byteBuffer.putLong(serverTimestamp); byteBuffer.putLong(serverTimestamp);
byteBuffer.putLong(messageUuid.getMostSignificantBits()); byteBuffer.putLong(messageUuid.getMostSignificantBits());
byteBuffer.putLong(messageUuid.getLeastSignificantBits()); byteBuffer.putLong(messageUuid.getLeastSignificantBits());
return byteBuffer.array(); return AttributeValues.fromByteBuffer(byteBuffer.flip());
} }
private static byte[] convertDestinationDeviceIdToSortKeyPrefix(final long destinationDeviceId) { private static AttributeValue convertDestinationDeviceIdToSortKeyPrefix(final long destinationDeviceId) {
ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]); ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(destinationDeviceId); byteBuffer.putLong(destinationDeviceId);
return byteBuffer.array(); return AttributeValues.fromByteBuffer(byteBuffer.flip());
} }
private static SortKey convertSortKey(final byte[] bytes) { private static SortKey convertSortKey(final byte[] bytes) {
@ -273,8 +305,8 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
return new SortKey(destinationDeviceId, serverTimestamp, new UUID(mostSigBits, leastSigBits)); return new SortKey(destinationDeviceId, serverTimestamp, new UUID(mostSigBits, leastSigBits));
} }
private static byte[] convertLocalIndexMessageUuidSortKey(final UUID messageUuid) { private static AttributeValue convertLocalIndexMessageUuidSortKey(final UUID messageUuid) {
return UUIDUtil.toBytes(messageUuid); return AttributeValues.fromUUID(messageUuid);
} }
private static UUID convertLocalIndexMessageUuidSortKey(final byte[] bytes) { private static UUID convertLocalIndexMessageUuidSortKey(final byte[] bytes) {

59
service/src/main/java/org/whispersystems/textsecuregcm/storage/MigrationDeletedAccounts.java
View File

@ -1,42 +1,52 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID; import java.util.UUID;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import java.util.stream.Collectors;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class MigrationDeletedAccounts extends AbstractDynamoDbStore { public class MigrationDeletedAccounts extends AbstractDynamoDbStore {
private final Table table; private final String tableName;
static final String KEY_UUID = "U"; static final String KEY_UUID = "U";
public MigrationDeletedAccounts(DynamoDB dynamoDb, String tableName) { public MigrationDeletedAccounts(DynamoDbClient dynamoDb, String tableName) {
super(dynamoDb); super(dynamoDb);
this.tableName = tableName;
table = dynamoDb.getTable(tableName);
} }
public void put(UUID uuid) { public void put(UUID uuid) {
table.putItem(new Item() db().putItem(PutItemRequest.builder()
.withPrimaryKey(primaryKey(uuid))); .tableName(tableName)
.item(primaryKey(uuid))
.build());
} }
public List<UUID> getRecentlyDeletedUuids() { public List<UUID> getRecentlyDeletedUuids() {
final List<UUID> uuids = new ArrayList<>(); final List<UUID> uuids = new ArrayList<>();
Optional<ScanResponse> firstPage = db().scanPaginator(ScanRequest.builder()
.tableName(tableName)
.build()).stream().findAny(); // get the first available response
for (Item item : table.scan(new ScanSpec()).firstPage()) { if (firstPage.isPresent()) {
// only process one page each time. If we have a significant backlog at the end of the migration for (Map<String, AttributeValue> item : firstPage.get().items()) {
// we can handle it separately // only process one page each time. If we have a significant backlog at the end of the migration
uuids.add(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_UUID))); // we can handle it separately
uuids.add(AttributeValues.getUUID(item, KEY_UUID, null));
}
} }
return uuids; return uuids;
@ -45,20 +55,17 @@ public class MigrationDeletedAccounts extends AbstractDynamoDbStore {
public void delete(List<UUID> uuids) { public void delete(List<UUID> uuids) {
writeInBatches(uuids, (batch) -> { writeInBatches(uuids, (batch) -> {
List<WriteRequest> deletes = batch.stream().map((uuid) -> WriteRequest.builder().deleteRequest(DeleteRequest.builder()
.key(primaryKey(uuid))
.build()).build()).collect(Collectors.toList());
final TableWriteItems deleteItems = new TableWriteItems(table.getTableName()); executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
for (UUID uuid : batch) {
deleteItems.addPrimaryKeyToDelete(primaryKey(uuid));
}
executeTableWriteItemsUntilComplete(deleteItems);
}); });
} }
@VisibleForTesting @VisibleForTesting
public static PrimaryKey primaryKey(UUID uuid) { public static Map<String, AttributeValue> primaryKey(UUID uuid) {
return new PrimaryKey(KEY_UUID, UUIDUtil.toBytes(uuid)); return Map.of(KEY_UUID, AttributeValues.fromUUID(uuid));
} }
} }

37
service/src/main/java/org/whispersystems/textsecuregcm/storage/MigrationRetryAccounts.java
View File

@ -1,43 +1,44 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.UUID; import java.util.UUID;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
public class MigrationRetryAccounts extends AbstractDynamoDbStore { public class MigrationRetryAccounts extends AbstractDynamoDbStore {
private final Table table; private final String tableName;
static final String KEY_UUID = "U"; static final String KEY_UUID = "U";
public MigrationRetryAccounts(DynamoDB dynamoDb, String tableName) { public MigrationRetryAccounts(DynamoDbClient dynamoDb, String tableName) {
super(dynamoDb); super(dynamoDb);
table = dynamoDb.getTable(tableName); this.tableName = tableName;
} }
public void put(UUID uuid) { public void put(UUID uuid) {
table.putItem(new Item() db().putItem(PutItemRequest.builder()
.withPrimaryKey(primaryKey(uuid))); .tableName(tableName)
.item(primaryKey(uuid))
.build());
} }
public List<UUID> getUuids(int max) { public List<UUID> getUuids(int max) {
final List<UUID> uuids = new ArrayList<>(); final List<UUID> uuids = new ArrayList<>();
for (Page<Item, ScanOutcome> page : table.scan(new ScanSpec()).pages()) { for (ScanResponse response : db().scanPaginator(ScanRequest.builder().tableName(tableName).build())) {
for (Item item : page) { for (Map<String, AttributeValue> item : response.items()) {
uuids.add(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_UUID))); uuids.add(AttributeValues.getUUID(item, KEY_UUID, null));
if (uuids.size() >= max) { if (uuids.size() >= max) {
break; break;
@ -53,8 +54,8 @@ public class MigrationRetryAccounts extends AbstractDynamoDbStore {
} }
@VisibleForTesting @VisibleForTesting
public static PrimaryKey primaryKey(UUID uuid) { public static Map<String, AttributeValue> primaryKey(UUID uuid) {
return new PrimaryKey(KEY_UUID, UUIDUtil.toBytes(uuid)); return Map.of(KEY_UUID, AttributeValues.fromUUID(uuid));
} }
} }

48
service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java
View File

@ -5,25 +5,23 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.PutItemSpec;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import java.time.Clock; import java.time.Clock;
import java.time.Duration; import java.time.Duration;
import java.util.Map; import java.util.Map;
import java.util.UUID; import java.util.UUID;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
/** /**
* Stores push challenge tokens. Users may have at most one outstanding push challenge token at a time. * Stores push challenge tokens. Users may have at most one outstanding push challenge token at a time.
*/ */
public class PushChallengeDynamoDb extends AbstractDynamoDbStore { public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
private final Table table; private final String tableName;
private final Clock clock; private final Clock clock;
static final String KEY_ACCOUNT_UUID = "U"; static final String KEY_ACCOUNT_UUID = "U";
@ -33,15 +31,15 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
private static final Map<String, String> UUID_NAME_MAP = Map.of("#uuid", KEY_ACCOUNT_UUID); private static final Map<String, String> UUID_NAME_MAP = Map.of("#uuid", KEY_ACCOUNT_UUID);
private static final Map<String, String> CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN); private static final Map<String, String> CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN);
public PushChallengeDynamoDb(final DynamoDB dynamoDB, final String tableName) { public PushChallengeDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
this(dynamoDB, tableName, Clock.systemUTC()); this(dynamoDB, tableName, Clock.systemUTC());
} }
@VisibleForTesting @VisibleForTesting
PushChallengeDynamoDb(final DynamoDB dynamoDB, final String tableName, final Clock clock) { PushChallengeDynamoDb(final DynamoDbClient dynamoDB, final String tableName, final Clock clock) {
super(dynamoDB); super(dynamoDB);
this.table = dynamoDB.getTable(tableName); this.tableName = tableName;
this.clock = clock; this.clock = clock;
} }
@ -57,13 +55,15 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
*/ */
public boolean add(final UUID accountUuid, final byte[] challengeToken, final Duration ttl) { public boolean add(final UUID accountUuid, final byte[] challengeToken, final Duration ttl) {
try { try {
table.putItem( new PutItemSpec() db().putItem(PutItemRequest.builder()
.withItem(new Item() .tableName(tableName)
.withBinary(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(accountUuid)) .item(Map.of(
.withBinary(ATTR_CHALLENGE_TOKEN, challengeToken) KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountUuid),
.withNumber(ATTR_TTL, getExpirationTimestamp(ttl))) ATTR_CHALLENGE_TOKEN, AttributeValues.fromByteArray(challengeToken),
.withConditionExpression("attribute_not_exists(#uuid)") ATTR_TTL, AttributeValues.fromLong(getExpirationTimestamp(ttl))))
.withNameMap(UUID_NAME_MAP)); .conditionExpression("attribute_not_exists(#uuid)")
.expressionAttributeNames(UUID_NAME_MAP)
.build());
return true; return true;
} catch (final ConditionalCheckFailedException e) { } catch (final ConditionalCheckFailedException e) {
return false; return false;
@ -84,11 +84,13 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
*/ */
public boolean remove(final UUID accountUuid, final byte[] challengeToken) { public boolean remove(final UUID accountUuid, final byte[] challengeToken) {
try { try {
table.deleteItem(new DeleteItemSpec() db().deleteItem(DeleteItemRequest.builder()
.withPrimaryKey(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(accountUuid)) .tableName(tableName)
.withConditionExpression("#challenge = :challenge") .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountUuid)))
.withNameMap(CHALLENGE_TOKEN_NAME_MAP) .conditionExpression("#challenge = :challenge")
.withValueMap(Map.of(":challenge", challengeToken))); .expressionAttributeNames(CHALLENGE_TOKEN_NAME_MAP)
.expressionAttributeValues(Map.of(":challenge", AttributeValues.fromByteArray(challengeToken)))
.build());
return true; return true;
} catch (final ConditionalCheckFailedException e) { } catch (final ConditionalCheckFailedException e) {
return false; return false;

52
service/src/main/java/org/whispersystems/textsecuregcm/storage/ReportMessageDynamoDb.java
View File

@ -1,13 +1,14 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome; import org.whispersystems.textsecuregcm.util.AttributeValues;
import com.amazonaws.services.dynamodbv2.document.DynamoDB; import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import com.amazonaws.services.dynamodbv2.document.Item; import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import com.amazonaws.services.dynamodbv2.document.Table; import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec; import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import com.amazonaws.services.dynamodbv2.model.ReturnValue; import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import java.time.Duration; import java.time.Duration;
import java.time.Instant; import java.time.Instant;
import java.util.Map;
public class ReportMessageDynamoDb { public class ReportMessageDynamoDb {
@ -16,33 +17,30 @@ public class ReportMessageDynamoDb {
static final Duration TIME_TO_LIVE = Duration.ofDays(7); static final Duration TIME_TO_LIVE = Duration.ofDays(7);
private final Table table; private final DynamoDbClient db;
private final String tableName;
public ReportMessageDynamoDb(final DynamoDB dynamoDB, final String tableName) { public ReportMessageDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
this.db = dynamoDB;
this.table = dynamoDB.getTable(tableName); this.tableName = tableName;
} }
public void store(byte[] hash) { public void store(byte[] hash) {
db.putItem(PutItemRequest.builder()
table.putItem(buildItemForHash(hash)); .tableName(tableName)
} .item(Map.of(
KEY_HASH, AttributeValues.fromByteArray(hash),
private Item buildItemForHash(byte[] hash) { ATTR_TTL, AttributeValues.fromLong(Instant.now().plus(TIME_TO_LIVE).getEpochSecond())
return new Item() ))
.withBinary(KEY_HASH, hash) .build());
.withLong(ATTR_TTL, Instant.now().plus(TIME_TO_LIVE).getEpochSecond());
} }
public boolean remove(byte[] hash) { public boolean remove(byte[] hash) {
final DeleteItemResponse deleteItemResponse = db.deleteItem(DeleteItemRequest.builder()
final DeleteItemSpec deleteItemSpec = new DeleteItemSpec() .tableName(tableName)
.withPrimaryKey(KEY_HASH, hash) .key(Map.of(KEY_HASH, AttributeValues.fromByteArray(hash)))
.withReturnValues(ReturnValue.ALL_OLD); .returnValues(ReturnValue.ALL_OLD)
.build());
final DeleteItemOutcome outcome = table.deleteItem(deleteItemSpec); return !deleteItemResponse.attributes().isEmpty();
return outcome.getItem() != null;
} }
} }

89
service/src/main/java/org/whispersystems/textsecuregcm/util/AttributeValues.java Normal file
View File

@ -0,0 +1,89 @@
/*
* Copyright 2013-2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.util;
import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
/** AwsAV provides static helper methods for working with AWS AttributeValues. */
public class AttributeValues {
public static AttributeValue fromString(String value) {
return AttributeValue.builder().s(value).build();
}
public static AttributeValue fromLong(long value) {
return AttributeValue.builder().n(Long.toString(value)).build();
}
public static AttributeValue fromInt(int value) {
return AttributeValue.builder().n(Integer.toString(value)).build();
}
public static AttributeValue fromByteArray(byte[] value) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteArray(value));
}
public static AttributeValue fromByteBuffer(ByteBuffer value) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteBuffer(value));
}
public static AttributeValue fromUUID(UUID uuid) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteArrayUnsafe(UUIDUtil.toBytes(uuid)));
}
public static AttributeValue fromSdkBytes(SdkBytes value) {
return AttributeValue.builder().b(value).build();
}
private static int toInt(AttributeValue av) {
return Integer.parseInt(av.n());
}
private static long toLong(AttributeValue av) {
return Long.parseLong(av.n());
}
private static UUID toUUID(AttributeValue av) {
return UUIDUtil.fromBytes(av.b().asByteArrayUnsafe()); // We're guaranteed not to modify the byte array
}
private static byte[] toByteArray(AttributeValue av) {
return av.b().asByteArray();
}
private static String toString(AttributeValue av) {
return av.s();
}
public static Optional<AttributeValue> get(Map<String, AttributeValue> item, String key) {
return Optional.ofNullable(item.get(key));
}
public static int getInt(Map<String, AttributeValue> item, String key, int defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toInt).orElse(defaultValue);
}
public static String getString(Map<String, AttributeValue> item, String key, String defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toString).orElse(defaultValue);
}
public static long getLong(Map<String, AttributeValue> item, String key, long defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toLong).orElse(defaultValue);
}
public static byte[] getByteArray(Map<String, AttributeValue> item, String key, byte[] defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toByteArray).orElse(defaultValue);
}
public static UUID getUUID(Map<String, AttributeValue> item, String key, UUID defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toUUID).orElse(defaultValue);
}
}

41
service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java Normal file
View File

@ -0,0 +1,41 @@
package org.whispersystems.textsecuregcm.util;
import org.whispersystems.textsecuregcm.configuration.DynamoDbConfiguration;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.core.client.config.ClientAsyncConfiguration;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClientBuilder;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import java.util.concurrent.Executor;
public class DynamoDbFromConfig {
private static ClientOverrideConfiguration clientOverrideConfiguration(DynamoDbConfiguration config) {
return ClientOverrideConfiguration.builder()
.apiCallTimeout(config.getClientExecutionTimeout())
.apiCallAttemptTimeout(config.getClientRequestTimeout())
.build();
}
public static DynamoDbClient client(DynamoDbConfiguration config, AwsCredentialsProvider credentialsProvider) {
return DynamoDbClient.builder()
.region(Region.of(config.getRegion()))
.credentialsProvider(credentialsProvider)
.overrideConfiguration(clientOverrideConfiguration(config))
.build();
}
public static DynamoDbAsyncClient asyncClient(DynamoDbConfiguration config, AwsCredentialsProvider credentialsProvider, Executor executor) {
DynamoDbAsyncClientBuilder builder = DynamoDbAsyncClient.builder()
.region(Region.of(config.getRegion()))
.credentialsProvider(credentialsProvider)
.overrideConfiguration(clientOverrideConfiguration(config));
if (executor != null) {
builder.asyncConfiguration(ClientAsyncConfiguration.builder()
.advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR,
executor)
.build());
}
return builder.build();
}
}

16
service/src/main/java/org/whispersystems/textsecuregcm/util/UUIDUtil.java
View File

@ -5,6 +5,7 @@
package org.whispersystems.textsecuregcm.util; package org.whispersystems.textsecuregcm.util;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.util.UUID; import java.util.UUID;
@ -26,12 +27,15 @@ public class UUIDUtil {
} }
public static UUID fromByteBuffer(final ByteBuffer byteBuffer) { public static UUID fromByteBuffer(final ByteBuffer byteBuffer) {
if (byteBuffer.array().length != 16) { try {
throw new IllegalArgumentException("unexpected byte array length; was " + byteBuffer.array().length + " but expected 16"); final long mostSigBits = byteBuffer.getLong();
final long leastSigBits = byteBuffer.getLong();
if (byteBuffer.hasRemaining()) {
throw new IllegalArgumentException("unexpected byte array length; was greater than 16");
}
return new UUID(mostSigBits, leastSigBits);
} catch (BufferUnderflowException e) {
throw new IllegalArgumentException("unexpected byte array length; was less than 16");
} }
final long mostSigBits = byteBuffer.getLong();
final long leastSigBits = byteBuffer.getLong();
return new UUID(mostSigBits, leastSigBits);
} }
} }

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

@ -9,11 +9,6 @@ import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.ClientConfiguration; import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.InstanceProfileCredentialsProvider; import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.DeserializationFeature;
import io.dropwizard.Application; import io.dropwizard.Application;
import io.dropwizard.cli.EnvironmentCommand; import io.dropwizard.cli.EnvironmentCommand;
@ -59,6 +54,9 @@ import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
import org.whispersystems.textsecuregcm.storage.ReservedUsernames; import org.whispersystems.textsecuregcm.storage.ReservedUsernames;
import org.whispersystems.textsecuregcm.storage.Usernames; import org.whispersystems.textsecuregcm.storage.Usernames;
import org.whispersystems.textsecuregcm.storage.UsernamesManager; import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfiguration> { public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfiguration> {
@ -100,64 +98,20 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("account_database_delete_user", accountJdbi, configuration.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration()); FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("account_database_delete_user", accountJdbi, configuration.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration());
ClientResources redisClusterClientResources = ClientResources.builder().build(); ClientResources redisClusterClientResources = ClientResources.builder().build();
AmazonDynamoDBClientBuilder clientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder keysDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getKeysDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getKeysDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getKeysDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder accountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS,
new LinkedBlockingDeque<>()); new LinkedBlockingDeque<>());
AmazonDynamoDBAsyncClientBuilder accountsDynamoDbAsyncClientBuilder = AmazonDynamoDBAsyncClientBuilder DynamoDbClient reportMessagesDynamoDb = DynamoDbFromConfig.client(configuration.getReportMessageDynamoDbConfiguration(),
.standard() software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withRegion(accountsDynamoDbClientBuilder.getRegion()) DynamoDbClient messageDynamoDb = DynamoDbFromConfig.client(configuration.getMessageDynamoDbConfiguration(),
.withClientConfiguration(accountsDynamoDbClientBuilder.getClientConfiguration()) software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.withCredentials(accountsDynamoDbClientBuilder.getCredentials()) DynamoDbClient preKeysDynamoDb = DynamoDbFromConfig.client(configuration.getKeysDynamoDbConfiguration(),
.withExecutorFactory(() -> accountsDynamoDbMigrationThreadPool); software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient accountsDynamoDbClient = DynamoDbFromConfig.client(configuration.getAccountsDynamoDbConfiguration(),
AmazonDynamoDBClientBuilder migrationDeletedAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
.standard() DynamoDbAsyncClient accountsDynamoDbAsyncClient = DynamoDbFromConfig.asyncClient(configuration.getAccountsDynamoDbConfiguration(),
.withRegion(configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getRegion()) software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create(),
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis())) accountsDynamoDbMigrationThreadPool);
.withRequestTimeout((int) configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder migrationRetryAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getMigrationRetryAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMigrationRetryAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getMigrationRetryAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder reportMessageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getReportMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getReportMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getReportMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDB messageDynamoDb = new DynamoDB(clientBuilder.build());
DynamoDB preKeysDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
DynamoDB reportMessagesDynamoDb = new DynamoDB(reportMessageDynamoDbClientBuilder.build());
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
AmazonDynamoDBAsync accountsDynamoDbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources); FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources);
@ -173,14 +127,16 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager); ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
DynamoDB migrationDeletedAccountsDynamoDb = new DynamoDB(migrationDeletedAccountsDynamoDbClientBuilder.build()); DynamoDbClient migrationDeletedAccountsDynamoDb = DynamoDbFromConfig.client(configuration.getMigrationDeletedAccountsDynamoDbConfiguration(),
DynamoDB migrationRetryAccountsDynamoDb = new DynamoDB(migrationRetryAccountsDynamoDbClientBuilder.build()); software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient migrationRetryAccountsDynamoDb = DynamoDbFromConfig.client(configuration.getMigrationRetryAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(migrationDeletedAccountsDynamoDb, configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName()); MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(migrationDeletedAccountsDynamoDb, configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, configuration.getMigrationRetryAccountsDynamoDbConfiguration().getTableName()); MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, configuration.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
Accounts accounts = new Accounts(accountDatabase); Accounts accounts = new Accounts(accountDatabase);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts); AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
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);

158
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDbTest.java
View File

@ -11,26 +11,12 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import io.github.resilience4j.circuitbreaker.CallNotPermittedException; import io.github.resilience4j.circuitbreaker.CallNotPermittedException;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.util.Arrays; import java.util.Arrays;
import java.util.Collections; import java.util.Collections;
import java.util.HashSet; import java.util.HashSet;
import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.Random; import java.util.Random;
import java.util.Set; import java.util.Set;
@ -46,7 +32,22 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration; import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.entities.SignedPreKey; import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
class AccountsDynamoDbTest { class AccountsDynamoDbTest {
@ -59,49 +60,75 @@ class AccountsDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder() static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(ACCOUNTS_TABLE_NAME) .tableName(ACCOUNTS_TABLE_NAME)
.hashKey(AccountsDynamoDb.KEY_ACCOUNT_UUID) .hashKey(AccountsDynamoDb.KEY_ACCOUNT_UUID)
.attributeDefinition(new AttributeDefinition(AccountsDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B)) .attributeDefinition(AttributeDefinition.builder()
.attributeName(AccountsDynamoDb.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build(); .build();
private AccountsDynamoDb accountsDynamoDb; private AccountsDynamoDb accountsDynamoDb;
private Table migrationDeletedAccountsTable;
private Table migrationRetryAccountsTable;
@BeforeEach @BeforeEach
void setupAccountsDao() { void setupAccountsDao() {
CreateTableRequest createNumbersTableRequest = CreateTableRequest.builder()
.tableName(NUMBERS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(AccountsDynamoDb.ATTR_ACCOUNT_E164)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(AccountsDynamoDb.ATTR_ACCOUNT_E164)
.attributeType(ScalarAttributeType.S)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
CreateTableRequest createNumbersTableRequest = new CreateTableRequest() dynamoDbExtension.getDynamoDbClient().createTable(createNumbersTableRequest);
.withTableName(NUMBERS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(AccountsDynamoDb.ATTR_ACCOUNT_E164, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(AccountsDynamoDb.ATTR_ACCOUNT_E164, ScalarAttributeType.S))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
final Table numbersTable = dynamoDbExtension.getDynamoDB().createTable(createNumbersTableRequest); final CreateTableRequest createMigrationDeletedAccountsTableRequest = CreateTableRequest.builder()
.tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
final CreateTableRequest createMigrationDeletedAccountsTableRequest = new CreateTableRequest() dynamoDbExtension.getDynamoDbClient().createTable(createMigrationDeletedAccountsTableRequest);
.withTableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(MigrationDeletedAccounts.KEY_UUID, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(MigrationDeletedAccounts.KEY_UUID, ScalarAttributeType.B))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
migrationDeletedAccountsTable = dynamoDbExtension.getDynamoDB().createTable(createMigrationDeletedAccountsTableRequest); MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(
dynamoDbExtension.getDynamoDbClient(), MIGRATION_DELETED_ACCOUNTS_TABLE_NAME);
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDB(), final CreateTableRequest createMigrationRetryAccountsTableRequest = CreateTableRequest.builder()
migrationDeletedAccountsTable.getTableName()); .tableName(MIGRATION_RETRY_ACCOUNTS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
final CreateTableRequest createMigrationRetryAccountsTableRequest = new CreateTableRequest() dynamoDbExtension.getDynamoDbClient().createTable(createMigrationRetryAccountsTableRequest);
.withTableName(MIGRATION_RETRY_ACCOUNTS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(MigrationRetryAccounts.KEY_UUID, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(MigrationRetryAccounts.KEY_UUID, ScalarAttributeType.B))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
migrationRetryAccountsTable = dynamoDbExtension.getDynamoDB().createTable(createMigrationRetryAccountsTableRequest); MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts((dynamoDbExtension.getDynamoDbClient()),
MIGRATION_RETRY_ACCOUNTS_TABLE_NAME);
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts((dynamoDbExtension.getDynamoDB()), this.accountsDynamoDb = new AccountsDynamoDb(
migrationRetryAccountsTable.getTableName()); dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getDynamoDbAsyncClient(),
this.accountsDynamoDb = new AccountsDynamoDb(dynamoDbExtension.getClient(), dynamoDbExtension.getAsyncClient(), new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingDeque<>()), dynamoDbExtension.getDynamoDB(), dynamoDbExtension.getTableName(), numbersTable.getTableName(), new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingDeque<>()),
migrationDeletedAccounts, migrationRetryAccounts); dynamoDbExtension.getTableName(),
NUMBERS_TABLE_NAME,
migrationDeletedAccounts,
migrationRetryAccounts);
} }
@Test @Test
@ -262,8 +289,12 @@ class AccountsDynamoDbTest {
verifyRecentlyDeletedAccountsTableItemCount(1); verifyRecentlyDeletedAccountsTableItemCount(1);
assertThat(migrationDeletedAccountsTable Map<String, AttributeValue> primaryKey = MigrationDeletedAccounts.primaryKey(deletedAccount.getUuid());
.getItem(MigrationDeletedAccounts.primaryKey(deletedAccount.getUuid()))).isNotNull(); assertThat(dynamoDbExtension.getDynamoDbClient().getItem(GetItemRequest.builder()
.tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.key(Map.of(MigrationDeletedAccounts.KEY_UUID, primaryKey.get(MigrationDeletedAccounts.KEY_UUID)))
.build()))
.isNotNull();
accountsDynamoDb.deleteRecentlyDeletedUuids(); accountsDynamoDb.deleteRecentlyDeletedUuids();
@ -273,8 +304,10 @@ class AccountsDynamoDbTest {
private void verifyRecentlyDeletedAccountsTableItemCount(int expectedItemCount) { private void verifyRecentlyDeletedAccountsTableItemCount(int expectedItemCount) {
int totalItems = 0; int totalItems = 0;
for (Page<Item, ScanOutcome> page : migrationDeletedAccountsTable.scan(new ScanSpec()).pages()) { for (ScanResponse page : dynamoDbExtension.getDynamoDbClient().scanPaginator(ScanRequest.builder()
for (Item ignored : page) { .tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.build())) {
for (Map<String, AttributeValue> item : page.items()) {
totalItems++; totalItems++;
} }
} }
@ -306,16 +339,15 @@ class AccountsDynamoDbTest {
configuration.setRingBufferSizeInClosedState(2); configuration.setRingBufferSizeInClosedState(2);
configuration.setFailureRateThreshold(50); configuration.setFailureRateThreshold(50);
final AmazonDynamoDB client = mock(AmazonDynamoDB.class); final DynamoDbClient client = mock(DynamoDbClient.class);
final DynamoDB dynamoDB = new DynamoDB(client);
when(client.transactWriteItems(any())) when(client.transactWriteItems(any(TransactWriteItemsRequest.class)))
.thenThrow(RuntimeException.class); .thenThrow(RuntimeException.class);
when(client.updateItem(any())) when(client.updateItem(any(UpdateItemRequest.class)))
.thenThrow(RuntimeException.class); .thenThrow(RuntimeException.class);
AccountsDynamoDb accounts = new AccountsDynamoDb(client, mock(AmazonDynamoDBAsync.class), mock(ThreadPoolExecutor.class), dynamoDB, ACCOUNTS_TABLE_NAME, NUMBERS_TABLE_NAME, mock( AccountsDynamoDb accounts = new AccountsDynamoDb(client, mock(DynamoDbAsyncClient.class), mock(ThreadPoolExecutor.class), ACCOUNTS_TABLE_NAME, NUMBERS_TABLE_NAME, mock(
MigrationDeletedAccounts.class), mock(MigrationRetryAccounts.class)); MigrationDeletedAccounts.class), mock(MigrationRetryAccounts.class));
Account account = generateAccount("+14151112222", UUID.randomUUID()); Account account = generateAccount("+14151112222", UUID.randomUUID());
@ -408,20 +440,22 @@ class AccountsDynamoDbTest {
} }
private void verifyStoredState(String number, UUID uuid, Account expecting) { private void verifyStoredState(String number, UUID uuid, Account expecting) {
final Table accounts = dynamoDbExtension.getDynamoDB().getTable(dynamoDbExtension.getTableName()); final DynamoDbClient db = dynamoDbExtension.getDynamoDbClient();
Item item = accounts.getItem(new GetItemSpec() final GetItemResponse get = db.getItem(GetItemRequest.builder()
.withPrimaryKey(AccountsDynamoDb.KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(uuid)) .tableName(dynamoDbExtension.getTableName())
.withConsistentRead(true)); .key(Map.of(AccountsDynamoDb.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.consistentRead(true)
.build());
if (item != null) { if (get.hasItem()) {
String data = new String(item.getBinary(AccountsDynamoDb.ATTR_ACCOUNT_DATA), StandardCharsets.UTF_8); String data = new String(get.item().get(AccountsDynamoDb.ATTR_ACCOUNT_DATA).b().asByteArray(), StandardCharsets.UTF_8);
assertThat(data).isNotEmpty(); assertThat(data).isNotEmpty();
assertThat(item.getNumber(AccountsDynamoDb.ATTR_MIGRATION_VERSION).intValue()) assertThat(AttributeValues.getInt(get.item(), AccountsDynamoDb.ATTR_MIGRATION_VERSION, -1))
.isEqualTo(expecting.getDynamoDbMigrationVersion()); .isEqualTo(expecting.getDynamoDbMigrationVersion());
Account result = AccountsDynamoDb.fromItem(item); Account result = AccountsDynamoDb.fromItem(get.item());
verifyStoredState(number, uuid, result, expecting); verifyStoredState(number, uuid, result, expecting);
} else { } else {
throw new AssertionError("No data"); throw new AssertionError("No data");

102
service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtension.java
View File

@ -1,33 +1,35 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.almworks.sqlite4java.SQLite; import com.almworks.sqlite4java.SQLite;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.client.builder.AwsClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner; import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer; import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.GlobalSecondaryIndex;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
import java.net.ServerSocket; import java.net.ServerSocket;
import java.net.URI;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import org.junit.jupiter.api.extension.AfterEachCallback; import org.junit.jupiter.api.extension.AfterEachCallback;
import org.junit.jupiter.api.extension.BeforeEachCallback; import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtensionContext; import org.junit.jupiter.api.extension.ExtensionContext;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndex;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback { public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback {
static final String DEFAULT_TABLE_NAME = "test_table"; static final String DEFAULT_TABLE_NAME = "test_table";
static final ProvisionedThroughput DEFAULT_PROVISIONED_THROUGHPUT = new ProvisionedThroughput(20L, 20L); static final ProvisionedThroughput DEFAULT_PROVISIONED_THROUGHPUT = ProvisionedThroughput.builder()
.readCapacityUnits(20L)
.writeCapacityUnits(20L)
.build();
private DynamoDBProxyServer server; private DynamoDBProxyServer server;
private int port; private int port;
@ -42,9 +44,8 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
private final long readCapacityUnits; private final long readCapacityUnits;
private final long writeCapacityUnits; private final long writeCapacityUnits;
private AmazonDynamoDB client; private DynamoDbClient dynamoDB2;
private AmazonDynamoDBAsync asyncClient; private DynamoDbAsyncClient dynamoAsyncDB2;
private DynamoDB dynamoDB;
private DynamoDbExtension(String tableName, String hashKey, String rangeKey, List<AttributeDefinition> attributeDefinitions, List<GlobalSecondaryIndex> globalSecondaryIndexes, long readCapacityUnits, private DynamoDbExtension(String tableName, String hashKey, String rangeKey, List<AttributeDefinition> attributeDefinitions, List<GlobalSecondaryIndex> globalSecondaryIndexes, long readCapacityUnits,
long writeCapacityUnits) { long writeCapacityUnits) {
@ -87,26 +88,33 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
KeySchemaElement[] keySchemaElements; KeySchemaElement[] keySchemaElements;
if (rangeKeyName == null) { if (rangeKeyName == null) {
keySchemaElements = new KeySchemaElement[] { keySchemaElements = new KeySchemaElement[] {
new KeySchemaElement(hashKeyName, "HASH"), KeySchemaElement.builder().attributeName(hashKeyName).keyType(KeyType.HASH).build(),
}; };
} else { } else {
keySchemaElements = new KeySchemaElement[] { keySchemaElements = new KeySchemaElement[] {
new KeySchemaElement(hashKeyName, "HASH"), KeySchemaElement.builder().attributeName(hashKeyName).keyType(KeyType.HASH).build(),
new KeySchemaElement(rangeKeyName, "RANGE") KeySchemaElement.builder().attributeName(rangeKeyName).keyType(KeyType.RANGE).build(),
}; };
} }
final CreateTableRequest createTableRequest = new CreateTableRequest() final CreateTableRequest createTableRequest = CreateTableRequest.builder()
.withTableName(tableName) .tableName(tableName)
.withKeySchema(keySchemaElements) .keySchema(keySchemaElements)
.withAttributeDefinitions(attributeDefinitions.isEmpty() ? null : attributeDefinitions) .attributeDefinitions(attributeDefinitions.isEmpty() ? null : attributeDefinitions)
.withGlobalSecondaryIndexes(globalSecondaryIndexes.isEmpty() ? null : globalSecondaryIndexes) .globalSecondaryIndexes(globalSecondaryIndexes.isEmpty() ? null : globalSecondaryIndexes)
.withProvisionedThroughput(new ProvisionedThroughput(readCapacityUnits, writeCapacityUnits)); .provisionedThroughput(ProvisionedThroughput.builder()
.readCapacityUnits(readCapacityUnits)
.writeCapacityUnits(writeCapacityUnits)
.build())
.build();
getDynamoDB().createTable(createTableRequest); getDynamoDbClient().createTable(createTableRequest);
} }
private void startServer() throws Exception { private void startServer() throws Exception {
// Even though we're using AWS SDK v2, Dynamo's local implementation's canonical location
// is within v1 (https://github.com/aws/aws-sdk-java-v2/issues/982). This does support
// v2 clients, though.
SQLite.setLibraryPath("target/lib"); // if you see a library failed to load error, you need to run mvn test-compile at least once first SQLite.setLibraryPath("target/lib"); // if you see a library failed to load error, you need to run mvn test-compile at least once first
ServerSocket serverSocket = new ServerSocket(0); ServerSocket serverSocket = new ServerSocket(0);
serverSocket.setReuseAddress(false); serverSocket.setReuseAddress(false);
@ -117,18 +125,18 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
} }
private void initializeClient() { private void initializeClient() {
AmazonDynamoDBClientBuilder clientBuilder = AmazonDynamoDBClientBuilder.standard() dynamoDB2 = DynamoDbClient.builder()
.withEndpointConfiguration( .endpointOverride(URI.create("http://localhost:" + port))
new AwsClientBuilder.EndpointConfiguration("http://localhost:" + port, "local-test-region")) .region(Region.of("local-test-region"))
.withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("accessKey", "secretKey"))); .credentialsProvider(StaticCredentialsProvider.create(
client = clientBuilder.build(); AwsBasicCredentials.create("accessKey", "secretKey")))
.build();
asyncClient = AmazonDynamoDBAsyncClientBuilder.standard() dynamoAsyncDB2 = DynamoDbAsyncClient.builder()
.withEndpointConfiguration(clientBuilder.getEndpoint()) .endpointOverride(URI.create("http://localhost:" + port))
.withCredentials(clientBuilder.getCredentials()) .region(Region.of("local-test-region"))
.build(); .credentialsProvider(StaticCredentialsProvider.create(
AwsBasicCredentials.create("accessKey", "secretKey")))
dynamoDB = new DynamoDB(client); .build();
} }
static class DynamoDbExtensionBuilder { static class DynamoDbExtensionBuilder {
@ -140,8 +148,8 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
private List<AttributeDefinition> attributeDefinitions = new ArrayList<>(); private List<AttributeDefinition> attributeDefinitions = new ArrayList<>();
private List<GlobalSecondaryIndex> globalSecondaryIndexes = new ArrayList<>(); private List<GlobalSecondaryIndex> globalSecondaryIndexes = new ArrayList<>();
private long readCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.getReadCapacityUnits(); private long readCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.readCapacityUnits();
private long writeCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.getWriteCapacityUnits(); private long writeCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.writeCapacityUnits();
private DynamoDbExtensionBuilder() { private DynamoDbExtensionBuilder() {
@ -178,16 +186,12 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
} }
} }
public AmazonDynamoDB getClient() { public DynamoDbClient getDynamoDbClient() {
return client; return dynamoDB2;
} }
public AmazonDynamoDBAsync getAsyncClient() { public DynamoDbAsyncClient getDynamoDbAsyncClient() {
return asyncClient; return dynamoAsyncDB2;
}
public DynamoDB getDynamoDB() {
return dynamoDB;
} }
public String getTableName() { public String getTableName() {

40
service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbRule.java
View File

@ -5,13 +5,13 @@
package org.whispersystems.textsecuregcm.storage; package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import org.whispersystems.textsecuregcm.tests.util.LocalDynamoDbRule; import org.whispersystems.textsecuregcm.tests.util.LocalDynamoDbRule;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
public class KeysDynamoDbRule extends LocalDynamoDbRule { public class KeysDynamoDbRule extends LocalDynamoDbRule {
public static final String TABLE_NAME = "Signal_Keys_Test"; public static final String TABLE_NAME = "Signal_Keys_Test";
@ -19,18 +19,22 @@ public class KeysDynamoDbRule extends LocalDynamoDbRule {
@Override @Override
protected void before() throws Throwable { protected void before() throws Throwable {
super.before(); super.before();
getDynamoDbClient().createTable(CreateTableRequest.builder()
final DynamoDB dynamoDB = getDynamoDB(); .tableName(TABLE_NAME)
.keySchema(
final CreateTableRequest createTableRequest = new CreateTableRequest() KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(),
.withTableName(TABLE_NAME) KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE)
.withKeySchema(new KeySchemaElement(KeysDynamoDb.KEY_ACCOUNT_UUID, "HASH"), .build())
new KeySchemaElement(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID, "RANGE")) .attributeDefinitions(AttributeDefinition.builder()
.withAttributeDefinitions(new AttributeDefinition(KeysDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B), .attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID)
new AttributeDefinition(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID, ScalarAttributeType.B)) .attributeType(ScalarAttributeType.B)
.withProvisionedThroughput(new ProvisionedThroughput(20L, 20L)); .build(),
AttributeDefinition.builder()
dynamoDB.createTable(createTableRequest); .attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build())
.build());
} }
@Override @Override

10
service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDbTest.java
View File

@ -9,6 +9,7 @@ import org.junit.Before;
import org.junit.ClassRule; import org.junit.ClassRule;
import org.junit.Test; import org.junit.Test;
import org.whispersystems.textsecuregcm.entities.PreKey; import org.whispersystems.textsecuregcm.entities.PreKey;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.util.Collections; import java.util.Collections;
import java.util.List; import java.util.List;
@ -17,6 +18,7 @@ import java.util.Optional;
import java.util.Set; import java.util.Set;
import java.util.UUID; import java.util.UUID;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
@ -34,7 +36,7 @@ public class KeysDynamoDbTest {
@Before @Before
public void setup() { public void setup() {
keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDB(), KeysDynamoDbRule.TABLE_NAME); keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME);
account = mock(Account.class); account = mock(Account.class);
when(account.getNumber()).thenReturn(ACCOUNT_NUMBER); when(account.getNumber()).thenReturn(ACCOUNT_NUMBER);
@ -133,4 +135,10 @@ public class KeysDynamoDbTest {
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID)); assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1)); assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
} }
@Test
public void testSortKeyPrefix() {
AttributeValue got = KeysDynamoDb.getSortKeyPrefix(123);
assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, got.b().asByteArray());
}
} }

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

@ -9,12 +9,6 @@ import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.ItemCollection;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.protobuf.ByteString; import com.google.protobuf.ByteString;
import io.lettuce.core.cluster.SlotHash; import io.lettuce.core.cluster.SlotHash;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
@ -22,6 +16,7 @@ import java.time.Duration;
import java.time.Instant; import java.time.Instant;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorService;
@ -38,6 +33,10 @@ import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.metrics.PushLatencyManager; import org.whispersystems.textsecuregcm.metrics.PushLatencyManager;
import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest; import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule; import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest { public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
@ -62,7 +61,7 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
connection.sync().upstream().commands().configSet("notify-keyspace-events", "K$glz"); connection.sync().upstream().commands().configSet("notify-keyspace-events", "K$glz");
}); });
final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7)); final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
final AccountsManager accountsManager = mock(AccountsManager.class); final AccountsManager accountsManager = mock(AccountsManager.class);
final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class); final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
@ -142,17 +141,16 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
final List<MessageProtos.Envelope> persistedMessages = new ArrayList<>(messageCount); final List<MessageProtos.Envelope> persistedMessages = new ArrayList<>(messageCount);
DynamoDB dynamoDB = messagesDynamoDbRule.getDynamoDB(); DynamoDbClient dynamoDB = messagesDynamoDbRule.getDynamoDbClient();
Table table = dynamoDB.getTable(MessagesDynamoDbRule.TABLE_NAME); for (Map<String, AttributeValue> item : dynamoDB
final ItemCollection<ScanOutcome> scan = table.scan(new ScanSpec()); .scan(ScanRequest.builder().tableName(MessagesDynamoDbRule.TABLE_NAME).build()).items()) {
for (Item item : scan) { persistedMessages.add(MessageProtos.Envelope.newBuilder()
persistedMessages.add(MessageProtos.Envelope.newBuilder() .setServerGuid(AttributeValues.getUUID(item, "U", null).toString())
.setServerGuid(convertBinaryToUuid(item.getBinary("U")).toString()) .setType(MessageProtos.Envelope.Type.valueOf(AttributeValues.getInt(item, "T", -1)))
.setType(MessageProtos.Envelope.Type.valueOf(item.getInt("T"))) .setTimestamp(AttributeValues.getLong(item, "TS", -1))
.setTimestamp(item.getLong("TS")) .setServerTimestamp(extractServerTimestamp(AttributeValues.getByteArray(item, "S", null)))
.setServerTimestamp(extractServerTimestamp(item.getBinary("S"))) .setContent(ByteString.copyFrom(AttributeValues.getByteArray(item, "C", null)))
.setContent(ByteString.copyFrom(item.getBinary("C"))) .build());
.build());
} }
assertEquals(expectedMessages, persistedMessages); assertEquals(expectedMessages, persistedMessages);

11
service/src/test/java/org/whispersystems/textsecuregcm/storage/MigrationDeletedAccountsTest.java
View File

@ -4,10 +4,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.List; import java.util.List;
import java.util.UUID; import java.util.UUID;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class MigrationDeletedAccountsTest { class MigrationDeletedAccountsTest {
@ -15,13 +15,16 @@ class MigrationDeletedAccountsTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder() static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName("deleted_accounts_test") .tableName("deleted_accounts_test")
.hashKey(MigrationDeletedAccounts.KEY_UUID) .hashKey(MigrationDeletedAccounts.KEY_UUID)
.attributeDefinition(new AttributeDefinition(MigrationDeletedAccounts.KEY_UUID, ScalarAttributeType.B)) .attributeDefinition(AttributeDefinition.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build(); .build();
@Test @Test
void test() { void test() {
final MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDB(), final MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName()); dynamoDbExtension.getTableName());
UUID firstUuid = UUID.randomUUID(); UUID firstUuid = UUID.randomUUID();

11
service/src/test/java/org/whispersystems/textsecuregcm/storage/MigrationRetryAccountsTest.java
View File

@ -2,12 +2,12 @@ package org.whispersystems.textsecuregcm.storage;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.util.List; import java.util.List;
import java.util.UUID; import java.util.UUID;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class MigrationRetryAccountsTest { class MigrationRetryAccountsTest {
@ -15,13 +15,16 @@ class MigrationRetryAccountsTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder() static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName("account_migration_errors_test") .tableName("account_migration_errors_test")
.hashKey(MigrationRetryAccounts.KEY_UUID) .hashKey(MigrationRetryAccounts.KEY_UUID)
.attributeDefinition(new AttributeDefinition(MigrationRetryAccounts.KEY_UUID, ScalarAttributeType.B)) .attributeDefinition(AttributeDefinition.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build(); .build();
@Test @Test
void test() { void test() {
final MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(dynamoDbExtension.getDynamoDB(), final MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName()); dynamoDbExtension.getTableName());
UUID firstUuid = UUID.randomUUID(); UUID firstUuid = UUID.randomUUID();

11
service/src/test/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDbTest.java
View File

@ -9,8 +9,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.time.Clock; import java.time.Clock;
import java.time.Duration; import java.time.Duration;
import java.time.Instant; import java.time.Instant;
@ -20,6 +18,8 @@ import java.util.UUID;
import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class PushChallengeDynamoDbTest { class PushChallengeDynamoDbTest {
@ -34,12 +34,15 @@ class PushChallengeDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder() static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(TABLE_NAME) .tableName(TABLE_NAME)
.hashKey(PushChallengeDynamoDb.KEY_ACCOUNT_UUID) .hashKey(PushChallengeDynamoDb.KEY_ACCOUNT_UUID)
.attributeDefinition(new AttributeDefinition(PushChallengeDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B)) .attributeDefinition(AttributeDefinition.builder()
.attributeName(PushChallengeDynamoDb.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build(); .build();
@BeforeEach @BeforeEach
void setUp() { void setUp() {
this.pushChallengeDynamoDb = new PushChallengeDynamoDb(dynamoDbExtension.getDynamoDB(), TABLE_NAME, Clock.fixed( this.pushChallengeDynamoDb = new PushChallengeDynamoDb(dynamoDbExtension.getDynamoDbClient(), TABLE_NAME, Clock.fixed(
Instant.ofEpochMilli(CURRENT_TIME_MILLIS), ZoneId.systemDefault())); Instant.ofEpochMilli(CURRENT_TIME_MILLIS), ZoneId.systemDefault()));
} }

11
service/src/test/java/org/whispersystems/textsecuregcm/storage/ReportMessageDynamoDbTest.java
View File

@ -4,13 +4,13 @@ import static org.junit.jupiter.api.Assertions.assertAll;
import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.util.UUID; import java.util.UUID;
import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class ReportMessageDynamoDbTest { class ReportMessageDynamoDbTest {
@ -22,13 +22,16 @@ class ReportMessageDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder() static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(TABLE_NAME) .tableName(TABLE_NAME)
.hashKey(ReportMessageDynamoDb.KEY_HASH) .hashKey(ReportMessageDynamoDb.KEY_HASH)
.attributeDefinition(new AttributeDefinition(ReportMessageDynamoDb.KEY_HASH, ScalarAttributeType.B)) .attributeDefinition(AttributeDefinition.builder()
.attributeName(ReportMessageDynamoDb.KEY_HASH)
.attributeType(ScalarAttributeType.B)
.build())
.build(); .build();
@BeforeEach @BeforeEach
void setUp() { void setUp() {
this.reportMessageDynamoDb = new ReportMessageDynamoDb(dynamoDbExtension.getDynamoDB(), TABLE_NAME); this.reportMessageDynamoDb = new ReportMessageDynamoDb(dynamoDbExtension.getDynamoDbClient(), TABLE_NAME);
} }
@Test @Test

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

@ -20,7 +20,6 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.verifyZeroInteractions; import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import io.lettuce.core.RedisException; import io.lettuce.core.RedisException;
import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
import java.util.HashSet; import java.util.HashSet;
@ -49,6 +48,7 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.ProfilesManager; import org.whispersystems.textsecuregcm.storage.ProfilesManager;
import org.whispersystems.textsecuregcm.storage.UsernamesManager; import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper; import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
class AccountsManagerTest { class AccountsManagerTest {

2
service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/MessagesDynamoDbTest.java
View File

@ -67,7 +67,7 @@ public class MessagesDynamoDbTest {
@Before @Before
public void setup() { public void setup() {
messagesDynamoDb = new MessagesDynamoDb(dynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7)); messagesDynamoDb = new MessagesDynamoDb(dynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
} }
@Test @Test

23
service/src/test/java/org/whispersystems/textsecuregcm/tests/util/LocalDynamoDbRule.java
View File

@ -6,16 +6,16 @@
package org.whispersystems.textsecuregcm.tests.util; package org.whispersystems.textsecuregcm.tests.util;
import com.almworks.sqlite4java.SQLite; import com.almworks.sqlite4java.SQLite;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.client.builder.AwsClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner; import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer; import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
import org.junit.rules.ExternalResource; import org.junit.rules.ExternalResource;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import java.net.ServerSocket; import java.net.ServerSocket;
import java.net.URI;
public class LocalDynamoDbRule extends ExternalResource { public class LocalDynamoDbRule extends ExternalResource {
private DynamoDBProxyServer server; private DynamoDBProxyServer server;
@ -43,11 +43,12 @@ public class LocalDynamoDbRule extends ExternalResource {
super.after(); super.after();
} }
public DynamoDB getDynamoDB() { public DynamoDbClient getDynamoDbClient() {
AmazonDynamoDBClientBuilder clientBuilder = return DynamoDbClient.builder()
AmazonDynamoDBClientBuilder.standard() .endpointOverride(URI.create("http://localhost:" + port))
.withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration("http://localhost:" + port, "local-test-region")) .region(Region.of("local-test-region"))
.withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("accessKey", "secretKey"))); .credentialsProvider(StaticCredentialsProvider.create(
return new DynamoDB(clientBuilder.build()); AwsBasicCredentials.create("accessKey", "secretKey")))
.build();
} }
} }

47
service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MessagesDynamoDbRule.java
View File

@ -5,15 +5,15 @@
package org.whispersystems.textsecuregcm.tests.util; package org.whispersystems.textsecuregcm.tests.util;
import com.amazonaws.services.dynamodbv2.document.DynamoDB; import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest; import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement; import software.amazon.awssdk.services.dynamodb.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.LocalSecondaryIndex; import software.amazon.awssdk.services.dynamodb.model.LocalSecondaryIndex;
import com.amazonaws.services.dynamodbv2.model.Projection; import software.amazon.awssdk.services.dynamodb.model.Projection;
import com.amazonaws.services.dynamodbv2.model.ProjectionType; import software.amazon.awssdk.services.dynamodb.model.ProjectionType;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput; import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType; import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
public class MessagesDynamoDbRule extends LocalDynamoDbRule { public class MessagesDynamoDbRule extends LocalDynamoDbRule {
@ -22,20 +22,21 @@ public class MessagesDynamoDbRule extends LocalDynamoDbRule {
@Override @Override
protected void before() throws Throwable { protected void before() throws Throwable {
super.before(); super.before();
DynamoDB dynamoDB = getDynamoDB(); getDynamoDbClient().createTable(CreateTableRequest.builder()
CreateTableRequest createTableRequest = new CreateTableRequest() .tableName(TABLE_NAME)
.withTableName(TABLE_NAME) .keySchema(KeySchemaElement.builder().attributeName("H").keyType(KeyType.HASH).build(),
.withKeySchema(new KeySchemaElement("H", "HASH"), KeySchemaElement.builder().attributeName("S").keyType(KeyType.RANGE).build())
new KeySchemaElement("S", "RANGE")) .attributeDefinitions(
.withAttributeDefinitions(new AttributeDefinition("H", ScalarAttributeType.B), AttributeDefinition.builder().attributeName("H").attributeType(ScalarAttributeType.B).build(),
new AttributeDefinition("S", ScalarAttributeType.B), AttributeDefinition.builder().attributeName("S").attributeType(ScalarAttributeType.B).build(),
new AttributeDefinition("U", ScalarAttributeType.B)) AttributeDefinition.builder().attributeName("U").attributeType(ScalarAttributeType.B).build())
.withProvisionedThroughput(new ProvisionedThroughput(20L, 20L)) .provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build())
.withLocalSecondaryIndexes(new LocalSecondaryIndex().withIndexName("Message_UUID_Index") .localSecondaryIndexes(LocalSecondaryIndex.builder().indexName("Message_UUID_Index")
.withKeySchema(new KeySchemaElement("H", "HASH"), .keySchema(KeySchemaElement.builder().attributeName("H").keyType(KeyType.HASH).build(),
new KeySchemaElement("U", "RANGE")) KeySchemaElement.builder().attributeName("U").keyType(KeyType.RANGE).build())
.withProjection(new Projection().withProjectionType(ProjectionType.KEYS_ONLY))); .projection(Projection.builder().projectionType(ProjectionType.KEYS_ONLY).build())
dynamoDB.createTable(createTableRequest); .build())
.build());
} }
@Override @Override

1
service/src/test/java/org/whispersystems/textsecuregcm/util/AsnTableTest.java
View File

@ -12,7 +12,6 @@ import java.io.InputStreamReader;
import java.net.Inet4Address; import java.net.Inet4Address;
import java.net.UnknownHostException; import java.net.UnknownHostException;
import java.util.Optional; import java.util.Optional;
import java.util.zip.GZIPInputStream;
import static org.junit.jupiter.api.Assertions.*; import static org.junit.jupiter.api.Assertions.*;

60
service/src/test/java/org/whispersystems/textsecuregcm/util/AttributeValuesTest.java Normal file
View File

@ -0,0 +1,60 @@
/*
* Copyright 2013-2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.util;
import org.junit.jupiter.api.Test;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.UUID;
import static org.junit.jupiter.api.Assertions.*;
public class AttributeValuesTest {
@Test
void testUUIDRoundTrip() {
UUID orig = UUID.randomUUID();
AttributeValue av = AttributeValues.fromUUID(orig);
UUID returned = AttributeValues.getUUID(Map.of("foo", av), "foo", null);
assertEquals(orig, returned);
}
@Test
void testLongRoundTrip() {
long orig = 12345;
AttributeValue av = AttributeValues.fromLong(orig);
long returned = AttributeValues.getLong(Map.of("foo", av), "foo", -1);
assertEquals(orig, returned);
}
@Test
void testIntRoundTrip() {
int orig = 12345;
AttributeValue av = AttributeValues.fromInt(orig);
int returned = AttributeValues.getInt(Map.of("foo", av), "foo", -1);
assertEquals(orig, returned);
}
@Test
void testByteBuffer() {
byte[] bytes = {1, 2, 3};
ByteBuffer bb = ByteBuffer.wrap(bytes);
AttributeValue av = AttributeValues.fromByteBuffer(bb);
byte[] returned = av.b().asByteArray();
assertArrayEquals(bytes, returned);
returned = AttributeValues.getByteArray(Map.of("foo", av), "foo", null);
assertArrayEquals(bytes, returned);
}
@Test
void testByteBuffer2() {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(123);
assertEquals(byteBuffer.remaining(), 0);
AttributeValue av = AttributeValues.fromByteBuffer(byteBuffer.flip());
assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, AttributeValues.getByteArray(Map.of("foo", av), "foo", null));
}
}

2
service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java
View File

@ -79,7 +79,7 @@ public class WebSocketConnectionIntegrationTest extends AbstractRedisClusterTest
executorService = Executors.newSingleThreadExecutor(); executorService = Executors.newSingleThreadExecutor();
messagesCache = new MessagesCache(getRedisCluster(), getRedisCluster(), executorService); messagesCache = new MessagesCache(getRedisCluster(), getRedisCluster(), executorService);
messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7)); messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
reportMessageManager = mock(ReportMessageManager.class); reportMessageManager = mock(ReportMessageManager.class);
account = mock(Account.class); account = mock(Account.class);
device = mock(Device.class); device = mock(Device.class);