mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-21 18:38:05 +01:00
Build Dynamo DB backed Message Store (#358)
* Work in progress... * Finish first pass draft of MessagesDynamoDb * Use begins_with everywhere for destination device id * Remove now unused methods * First basic test built * Add another test case * Remove comment * Verify more of the message contents * Ensure all methods are tested * Integrate MessagesDynamoDb into the MessagesManager This change plugs the MessagesDynamoDb class into the live serving flow in MessagesManager. Tests are not yet as comprehensive for this big a change as they should be, but they now compile and pass so checkpointing here with a commit. * Put DynamoDB before RDBS when deleting specific messages * Extract method * Make aws sdk version into a property * Rename clientBuilder * Discard messages with no GUID * Unify batching logic into one function * Comment on the source of the value in this constant * Inline method * Variable name swizzle * Add timers to all public methods * Add missing return statements * Reject messages that are too large with response code 413 * Add configuration to control dynamo DB timeouts * Set server timestamp from the ReceiptSender * Change to shorter key names to optimize IOPS * Fix tests broken by changing column names * Fix broken copyright template output * Remove copyright template error text * Add experiments to control use of dynamo and rds in message storage * Specify instance profile credentials for the dynamic configuration manager * Use property for aws sdk version * Switch dynamo to instance profile credentials * Add metrics to the batch write loop * Use placeholders in logging
This commit is contained in:
Ehren Kret
GitHub
@@ -13,7 +13,6 @@ import org.whispersystems.textsecuregcm.storage.Account;
|
||||
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.Set;
|
||||
import java.util.UUID;
|
||||
@@ -50,18 +49,18 @@ public class ExperimentEnrollmentManagerTest {
|
||||
|
||||
@Test
|
||||
public void testIsEnrolled() {
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account, EXPERIMENT_NAME));
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account, EXPERIMENT_NAME + "-unrelated-experiment"));
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account.getUuid(), EXPERIMENT_NAME));
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account.getUuid(), EXPERIMENT_NAME + "-unrelated-experiment"));
|
||||
|
||||
when(experimentEnrollmentConfiguration.getEnrolledUuids()).thenReturn(Set.of(ACCOUNT_UUID));
|
||||
assertTrue(experimentEnrollmentManager.isEnrolled(account, EXPERIMENT_NAME));
|
||||
assertTrue(experimentEnrollmentManager.isEnrolled(account.getUuid(), EXPERIMENT_NAME));
|
||||
|
||||
when(experimentEnrollmentConfiguration.getEnrolledUuids()).thenReturn(Collections.emptySet());
|
||||
when(experimentEnrollmentConfiguration.getEnrollmentPercentage()).thenReturn(0);
|
||||
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account, EXPERIMENT_NAME));
|
||||
assertFalse(experimentEnrollmentManager.isEnrolled(account.getUuid(), EXPERIMENT_NAME));
|
||||
|
||||
when(experimentEnrollmentConfiguration.getEnrollmentPercentage()).thenReturn(100);
|
||||
assertTrue(experimentEnrollmentManager.isEnrolled(account, EXPERIMENT_NAME));
|
||||
assertTrue(experimentEnrollmentManager.isEnrolled(account.getUuid(), EXPERIMENT_NAME));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,12 @@
|
||||
|
||||
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.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.opentable.db.postgres.embedded.LiquibasePreparer;
|
||||
import com.opentable.db.postgres.junit.EmbeddedPostgresRules;
|
||||
@@ -18,11 +24,12 @@ import org.junit.Rule;
|
||||
import org.junit.Test;
|
||||
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
|
||||
import org.whispersystems.textsecuregcm.metrics.PushLatencyManager;
|
||||
import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
|
||||
import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule;
|
||||
|
||||
import java.sql.PreparedStatement;
|
||||
import java.sql.ResultSet;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.time.Duration;
|
||||
import java.time.Instant;
|
||||
import java.util.ArrayList;
|
||||
@@ -35,6 +42,8 @@ import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
import static org.junit.Assert.assertEquals;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
@@ -43,11 +52,15 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
|
||||
@Rule
|
||||
public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(LiquibasePreparer.forClasspathLocation("messagedb.xml"));
|
||||
|
||||
@Rule
|
||||
public MessagesDynamoDbRule messagesDynamoDbRule = new MessagesDynamoDbRule();
|
||||
|
||||
private ExecutorService notificationExecutorService;
|
||||
private MessagesCache messagesCache;
|
||||
private MessagesManager messagesManager;
|
||||
private MessagePersister messagePersister;
|
||||
private Account account;
|
||||
private ExperimentEnrollmentManager experimentEnrollmentManager;
|
||||
|
||||
private static final Duration PERSIST_DELAY = Duration.ofMinutes(10);
|
||||
|
||||
@@ -62,11 +75,15 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
|
||||
});
|
||||
|
||||
final Messages messages = new Messages(new FaultTolerantDatabase("messages-test", Jdbi.create(db.getTestDatabase()), new CircuitBreakerConfiguration()));
|
||||
final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
|
||||
final AccountsManager accountsManager = mock(AccountsManager.class);
|
||||
|
||||
experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class);
|
||||
when(experimentEnrollmentManager.isEnrolled(any(UUID.class), anyString())).thenReturn(Boolean.TRUE);
|
||||
|
||||
notificationExecutorService = Executors.newSingleThreadExecutor();
|
||||
messagesCache = new MessagesCache(getRedisCluster(), getRedisCluster(), notificationExecutorService);
|
||||
messagesManager = new MessagesManager(messages, messagesCache, mock(PushLatencyManager.class));
|
||||
messagesManager = new MessagesManager(messages, messagesDynamoDb, messagesCache, mock(PushLatencyManager.class), experimentEnrollmentManager);
|
||||
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, mock(FeatureFlagsManager.class), PERSIST_DELAY);
|
||||
|
||||
account = mock(Account.class);
|
||||
@@ -139,25 +156,35 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
|
||||
|
||||
final List<MessageProtos.Envelope> persistedMessages = new ArrayList<>(messageCount);
|
||||
|
||||
try (final PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("SELECT * FROM messages WHERE destination = ? ORDER BY timestamp ASC")) {
|
||||
statement.setString(1, account.getNumber());
|
||||
|
||||
try (final ResultSet resultSet = statement.executeQuery()) {
|
||||
while (resultSet.next()) {
|
||||
persistedMessages.add(MessageProtos.Envelope.newBuilder()
|
||||
.setServerGuid(resultSet.getString("guid"))
|
||||
.setType(MessageProtos.Envelope.Type.valueOf(resultSet.getInt("type")))
|
||||
.setTimestamp(resultSet.getLong("timestamp"))
|
||||
.setServerTimestamp(resultSet.getLong("server_timestamp"))
|
||||
.setContent(ByteString.copyFrom(resultSet.getBytes("content")))
|
||||
.build());
|
||||
}
|
||||
}
|
||||
DynamoDB dynamoDB = messagesDynamoDbRule.getDynamoDB();
|
||||
Table table = dynamoDB.getTable(MessagesDynamoDbRule.TABLE_NAME);
|
||||
final ItemCollection<ScanOutcome> scan = table.scan(new ScanSpec());
|
||||
for (Item item : scan) {
|
||||
persistedMessages.add(MessageProtos.Envelope.newBuilder()
|
||||
.setServerGuid(convertBinaryToUuid(item.getBinary("U")).toString())
|
||||
.setType(MessageProtos.Envelope.Type.valueOf(item.getInt("T")))
|
||||
.setTimestamp(item.getLong("TS"))
|
||||
.setServerTimestamp(extractServerTimestamp(item.getBinary("S")))
|
||||
.setContent(ByteString.copyFrom(item.getBinary("C")))
|
||||
.build());
|
||||
}
|
||||
|
||||
assertEquals(expectedMessages, persistedMessages);
|
||||
}
|
||||
|
||||
private static UUID convertBinaryToUuid(byte[] bytes) {
|
||||
ByteBuffer bb = ByteBuffer.wrap(bytes);
|
||||
long msb = bb.getLong();
|
||||
long lsb = bb.getLong();
|
||||
return new UUID(msb, lsb);
|
||||
}
|
||||
|
||||
private static long extractServerTimestamp(byte[] bytes) {
|
||||
ByteBuffer bb = ByteBuffer.wrap(bytes);
|
||||
bb.getLong();
|
||||
return bb.getLong();
|
||||
}
|
||||
|
||||
private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final long timestamp) {
|
||||
return MessageProtos.Envelope.newBuilder()
|
||||
.setTimestamp(timestamp)
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
/*
|
||||
* Copyright 2021 Signal Messenger, LLC
|
||||
* SPDX-License-Identifier: AGPL-3.0-only
|
||||
*/
|
||||
|
||||
package org.whispersystems.textsecuregcm.tests.storage;
|
||||
|
||||
import com.google.protobuf.ByteString;
|
||||
import org.junit.Before;
|
||||
import org.junit.ClassRule;
|
||||
import org.junit.Test;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
|
||||
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
|
||||
import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Random;
|
||||
import java.util.UUID;
|
||||
import java.util.function.Consumer;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
public class MessagesDynamoDbTest {
|
||||
private static final Random random = new Random();
|
||||
private static final MessageProtos.Envelope MESSAGE1;
|
||||
private static final MessageProtos.Envelope MESSAGE2;
|
||||
private static final MessageProtos.Envelope MESSAGE3;
|
||||
|
||||
static {
|
||||
final long serverTimestamp = System.currentTimeMillis();
|
||||
MessageProtos.Envelope.Builder builder = MessageProtos.Envelope.newBuilder();
|
||||
builder.setType(MessageProtos.Envelope.Type.UNIDENTIFIED_SENDER);
|
||||
builder.setTimestamp(123456789L);
|
||||
builder.setContent(ByteString.copyFrom(new byte[]{(byte) 0xDE, (byte) 0xAD, (byte) 0xBE, (byte) 0xEF}));
|
||||
builder.setServerGuid(UUID.randomUUID().toString());
|
||||
builder.setServerTimestamp(serverTimestamp);
|
||||
|
||||
MESSAGE1 = builder.build();
|
||||
|
||||
builder.setType(MessageProtos.Envelope.Type.CIPHERTEXT);
|
||||
builder.setSource("12348675309");
|
||||
builder.setSourceUuid(UUID.randomUUID().toString());
|
||||
builder.setSourceDevice(1);
|
||||
builder.setContent(ByteString.copyFromUtf8("MOO"));
|
||||
builder.setServerGuid(UUID.randomUUID().toString());
|
||||
builder.setServerTimestamp(serverTimestamp + 1);
|
||||
|
||||
MESSAGE2 = builder.build();
|
||||
|
||||
builder.setType(MessageProtos.Envelope.Type.UNIDENTIFIED_SENDER);
|
||||
builder.clearSource();
|
||||
builder.clearSourceUuid();
|
||||
builder.clearSourceDevice();
|
||||
builder.setContent(ByteString.copyFromUtf8("COW"));
|
||||
builder.setServerGuid(UUID.randomUUID().toString());
|
||||
builder.setServerTimestamp(serverTimestamp); // Test same millisecond arrival for two different messages
|
||||
|
||||
MESSAGE3 = builder.build();
|
||||
}
|
||||
|
||||
private MessagesDynamoDb messagesDynamoDb;
|
||||
|
||||
@ClassRule
|
||||
public static MessagesDynamoDbRule dynamoDbRule = new MessagesDynamoDbRule();
|
||||
|
||||
@Before
|
||||
public void setup() {
|
||||
messagesDynamoDb = new MessagesDynamoDb(dynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testServerStart() {
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testSimpleFetchAfterInsert() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final int destinationDeviceId = random.nextInt(255) + 1;
|
||||
messagesDynamoDb.store(List.of(MESSAGE1, MESSAGE2, MESSAGE3), destinationUuid, destinationDeviceId);
|
||||
|
||||
final List<OutgoingMessageEntity> messagesStored = messagesDynamoDb.load(destinationUuid, destinationDeviceId, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE);
|
||||
assertThat(messagesStored).isNotNull().hasSize(3);
|
||||
final MessageProtos.Envelope firstMessage = MESSAGE1.getServerGuid().compareTo(MESSAGE3.getServerGuid()) < 0 ? MESSAGE1 : MESSAGE3;
|
||||
final MessageProtos.Envelope secondMessage = firstMessage == MESSAGE1 ? MESSAGE3 : MESSAGE1;
|
||||
assertThat(messagesStored).element(0).satisfies(verify(firstMessage));
|
||||
assertThat(messagesStored).element(1).satisfies(verify(secondMessage));
|
||||
assertThat(messagesStored).element(2).satisfies(verify(MESSAGE2));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDeleteForDestination() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final UUID secondDestinationUuid = UUID.randomUUID();
|
||||
messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
|
||||
messagesDynamoDb.deleteAllMessagesForAccount(destinationUuid);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDeleteForDestinationDevice() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final UUID secondDestinationUuid = UUID.randomUUID();
|
||||
messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
|
||||
messagesDynamoDb.deleteAllMessagesForDevice(destinationUuid, 2);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDeleteMessageByDestinationAndSourceAndTimestamp() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final UUID secondDestinationUuid = UUID.randomUUID();
|
||||
messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
|
||||
messagesDynamoDb.deleteMessageByDestinationAndSourceAndTimestamp(secondDestinationUuid, 1, MESSAGE2.getSource(), MESSAGE2.getTimestamp());
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDeleteMessageByDestinationAndGuid() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final UUID secondDestinationUuid = UUID.randomUUID();
|
||||
messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
|
||||
messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE2));
|
||||
|
||||
messagesDynamoDb.deleteMessageByDestinationAndGuid(secondDestinationUuid, 1, UUID.fromString(MESSAGE2.getServerGuid()));
|
||||
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE1));
|
||||
assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1).element(0).satisfies(verify(MESSAGE3));
|
||||
assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
|
||||
}
|
||||
|
||||
private static void verify(OutgoingMessageEntity retrieved, MessageProtos.Envelope inserted) {
|
||||
assertThat(retrieved.getTimestamp()).isEqualTo(inserted.getTimestamp());
|
||||
assertThat(retrieved.getSource()).isEqualTo(inserted.hasSource() ? inserted.getSource() : null);
|
||||
assertThat(retrieved.getSourceUuid()).isEqualTo(inserted.hasSourceUuid() ? UUID.fromString(inserted.getSourceUuid()) : null);
|
||||
assertThat(retrieved.getSourceDevice()).isEqualTo(inserted.getSourceDevice());
|
||||
assertThat(retrieved.getRelay()).isEqualTo(inserted.hasRelay() ? inserted.getRelay() : null);
|
||||
assertThat(retrieved.getType()).isEqualTo(inserted.getType().getNumber());
|
||||
assertThat(retrieved.getContent()).isEqualTo(inserted.hasContent() ? inserted.getContent().toByteArray() : null);
|
||||
assertThat(retrieved.getMessage()).isEqualTo(inserted.hasLegacyMessage() ? inserted.getLegacyMessage().toByteArray() : null);
|
||||
assertThat(retrieved.getServerTimestamp()).isEqualTo(inserted.getServerTimestamp());
|
||||
assertThat(retrieved.getGuid()).isEqualTo(UUID.fromString(inserted.getServerGuid()));
|
||||
}
|
||||
|
||||
private static VerifyMessage verify(MessageProtos.Envelope expected) {
|
||||
return new VerifyMessage(expected);
|
||||
}
|
||||
|
||||
private static final class VerifyMessage implements Consumer<OutgoingMessageEntity> {
|
||||
private final MessageProtos.Envelope expected;
|
||||
|
||||
public VerifyMessage(MessageProtos.Envelope expected) {
|
||||
this.expected = expected;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void accept(OutgoingMessageEntity outgoingMessageEntity) {
|
||||
verify(outgoingMessageEntity, expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
/*
|
||||
* Copyright 2021 Signal Messenger, LLC
|
||||
* SPDX-License-Identifier: AGPL-3.0-only
|
||||
*/
|
||||
|
||||
package org.whispersystems.textsecuregcm.tests.util;
|
||||
|
||||
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.server.DynamoDBProxyServer;
|
||||
import org.junit.rules.ExternalResource;
|
||||
|
||||
import java.net.ServerSocket;
|
||||
|
||||
public class LocalDynamoDbRule extends ExternalResource {
|
||||
private DynamoDBProxyServer server;
|
||||
private int port;
|
||||
|
||||
@Override
|
||||
protected void before() throws Throwable {
|
||||
super.before();
|
||||
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.setReuseAddress(false);
|
||||
port = serverSocket.getLocalPort();
|
||||
serverSocket.close();
|
||||
server = ServerRunner.createServerFromCommandLineArgs(new String[]{"-inMemory", "-port", String.valueOf(port)});
|
||||
server.start();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void after() {
|
||||
try {
|
||||
server.stop();
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
super.after();
|
||||
}
|
||||
|
||||
public DynamoDB getDynamoDB() {
|
||||
AmazonDynamoDBClientBuilder clientBuilder =
|
||||
AmazonDynamoDBClientBuilder.standard()
|
||||
.withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration("http://localhost:" + port, "local-test-region"))
|
||||
.withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("accessKey", "secretKey")));
|
||||
return new DynamoDB(clientBuilder.build());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
/*
|
||||
* Copyright 2021 Signal Messenger, LLC
|
||||
* SPDX-License-Identifier: AGPL-3.0-only
|
||||
*/
|
||||
|
||||
package org.whispersystems.textsecuregcm.tests.util;
|
||||
|
||||
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.LocalSecondaryIndex;
|
||||
import com.amazonaws.services.dynamodbv2.model.Projection;
|
||||
import com.amazonaws.services.dynamodbv2.model.ProjectionType;
|
||||
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
|
||||
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
|
||||
|
||||
public class MessagesDynamoDbRule extends LocalDynamoDbRule {
|
||||
|
||||
public static final String TABLE_NAME = "Signal_Messages_UnitTest";
|
||||
|
||||
@Override
|
||||
protected void before() throws Throwable {
|
||||
super.before();
|
||||
DynamoDB dynamoDB = getDynamoDB();
|
||||
CreateTableRequest createTableRequest = new CreateTableRequest()
|
||||
.withTableName(TABLE_NAME)
|
||||
.withKeySchema(new KeySchemaElement("H", "HASH"),
|
||||
new KeySchemaElement("S", "RANGE"))
|
||||
.withAttributeDefinitions(new AttributeDefinition("H", ScalarAttributeType.B),
|
||||
new AttributeDefinition("S", ScalarAttributeType.B),
|
||||
new AttributeDefinition("U", ScalarAttributeType.B))
|
||||
.withProvisionedThroughput(new ProvisionedThroughput(20L, 20L))
|
||||
.withLocalSecondaryIndexes(new LocalSecondaryIndex().withIndexName("Message_UUID_Index")
|
||||
.withKeySchema(new KeySchemaElement("H", "HASH"),
|
||||
new KeySchemaElement("U", "RANGE"))
|
||||
.withProjection(new Projection().withProjectionType(ProjectionType.KEYS_ONLY)));
|
||||
dynamoDB.createTable(createTableRequest);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void after() {
|
||||
super.after();
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@ import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.stubbing.Answer;
|
||||
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
|
||||
import org.whispersystems.textsecuregcm.metrics.PushLatencyManager;
|
||||
import org.whispersystems.textsecuregcm.push.ReceiptSender;
|
||||
import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
|
||||
@@ -28,11 +29,14 @@ import org.whispersystems.textsecuregcm.storage.Device;
|
||||
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
|
||||
import org.whispersystems.textsecuregcm.storage.Messages;
|
||||
import org.whispersystems.textsecuregcm.storage.MessagesCache;
|
||||
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
|
||||
import org.whispersystems.textsecuregcm.storage.MessagesManager;
|
||||
import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule;
|
||||
import org.whispersystems.websocket.WebSocketClient;
|
||||
import org.whispersystems.websocket.messages.WebSocketResponseMessage;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
@@ -47,6 +51,7 @@ import static org.junit.Assert.assertEquals;
|
||||
import static org.junit.Assert.fail;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyList;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.atMost;
|
||||
import static org.mockito.Mockito.mock;
|
||||
@@ -60,13 +65,18 @@ public class WebSocketConnectionIntegrationTest extends AbstractRedisClusterTest
|
||||
@Rule
|
||||
public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(LiquibasePreparer.forClasspathLocation("messagedb.xml"));
|
||||
|
||||
@Rule
|
||||
public MessagesDynamoDbRule messagesDynamoDbRule = new MessagesDynamoDbRule();
|
||||
|
||||
private ExecutorService executorService;
|
||||
private Messages messages;
|
||||
private MessagesDynamoDb messagesDynamoDb;
|
||||
private MessagesCache messagesCache;
|
||||
private Account account;
|
||||
private Device device;
|
||||
private WebSocketClient webSocketClient;
|
||||
private WebSocketConnection webSocketConnection;
|
||||
private ExperimentEnrollmentManager experimentEnrollmentManager;
|
||||
|
||||
private long serialTimestamp = System.currentTimeMillis();
|
||||
|
||||
@@ -82,17 +92,20 @@ public class WebSocketConnectionIntegrationTest extends AbstractRedisClusterTest
|
||||
executorService = Executors.newSingleThreadExecutor();
|
||||
messages = new Messages(new FaultTolerantDatabase("messages-test", Jdbi.create(db.getTestDatabase()), new CircuitBreakerConfiguration()));
|
||||
messagesCache = new MessagesCache(getRedisCluster(), getRedisCluster(), executorService);
|
||||
messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
|
||||
account = mock(Account.class);
|
||||
device = mock(Device.class);
|
||||
webSocketClient = mock(WebSocketClient.class);
|
||||
experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class);
|
||||
|
||||
when(account.getNumber()).thenReturn("+18005551234");
|
||||
when(account.getUuid()).thenReturn(UUID.randomUUID());
|
||||
when(device.getId()).thenReturn(1L);
|
||||
when(experimentEnrollmentManager.isEnrolled(any(UUID.class), anyString())).thenReturn(Boolean.FALSE);
|
||||
|
||||
webSocketConnection = new WebSocketConnection(
|
||||
mock(ReceiptSender.class),
|
||||
new MessagesManager(messages, messagesCache, mock(PushLatencyManager.class)),
|
||||
new MessagesManager(messages, messagesDynamoDb, messagesCache, mock(PushLatencyManager.class), experimentEnrollmentManager),
|
||||
account,
|
||||
device,
|
||||
webSocketClient);
|
||||
|
||||
@@ -195,7 +195,7 @@ public class WebSocketConnectionTest {
|
||||
futures.get(0).completeExceptionally(new IOException());
|
||||
futures.get(2).completeExceptionally(new IOException());
|
||||
|
||||
verify(storedMessages, times(1)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), eq(2L), eq(false));
|
||||
verify(storedMessages, times(1)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), eq(outgoingMessages.get(1).getGuid()));
|
||||
verify(receiptSender, times(1)).sendReceipt(eq(account), eq("sender1"), eq(2222L));
|
||||
|
||||
connection.stop();
|
||||
@@ -712,7 +712,7 @@ public class WebSocketConnectionTest {
|
||||
|
||||
// We should delete all three messages even though we only sent two; one got discarded because it was too big for
|
||||
// desktop clients.
|
||||
verify(storedMessages, times(3)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), anyLong(), anyBoolean());
|
||||
verify(storedMessages, times(3)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), any(UUID.class));
|
||||
|
||||
connection.stop();
|
||||
verify(client).close(anyInt(), anyString());
|
||||
@@ -785,7 +785,7 @@ public class WebSocketConnectionTest {
|
||||
futures.get(1).complete(response);
|
||||
futures.get(2).complete(response);
|
||||
|
||||
verify(storedMessages, times(3)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), anyLong(), anyBoolean());
|
||||
verify(storedMessages, times(3)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), any(UUID.class));
|
||||
|
||||
connection.stop();
|
||||
verify(client).close(anyInt(), anyString());
|
||||
|
||||
Reference in New Issue
Block a user