From 04728ea4bcbc162072fd9b22da3e2714b9015b1e Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Fri, 5 Feb 2021 16:16:27 -0500 Subject: [PATCH] Drop the old Postgres-based pre-key store. --- .../textsecuregcm/WhisperServerService.java | 6 +- .../controllers/KeysController.java | 28 +- .../storage/AccountsManager.java | 5 +- .../textsecuregcm/storage/Keys.java | 164 ------- .../textsecuregcm/storage/KeysDynamoDb.java | 7 +- .../textsecuregcm/storage/PreKeyStore.java | 23 - .../workers/DeleteUserCommand.java | 9 +- .../textsecuregcm/workers/VacuumCommand.java | 5 - .../tests/controllers/KeysControllerTest.java | 47 +- .../tests/storage/AccountsManagerTest.java | 19 +- .../textsecuregcm/tests/storage/KeysTest.java | 462 ------------------ 11 files changed, 40 insertions(+), 735 deletions(-) delete mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java delete mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/PreKeyStore.java delete mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index 2af1c748e..82a96089f 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -125,7 +125,6 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.FeatureFlags; import org.whispersystems.textsecuregcm.storage.FeatureFlagsManager; -import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.MessagePersister; import org.whispersystems.textsecuregcm.storage.Messages; @@ -293,7 +292,6 @@ public class WhisperServerService extends Application 0) { count = count - 1; @@ -104,7 +94,7 @@ public class KeysController { } } - getPreKeyStoreForProducer(account).store(account, device.getId(), preKeys.getPreKeys()); + keysDynamoDb.store(account, device.getId(), preKeys.getPreKeys()); } @Timed @@ -185,22 +175,14 @@ public class KeysController { private List getLocalKeys(Account destination, String deviceIdSelector) { try { if (deviceIdSelector.equals("*")) { - return getPreKeyStoreForConsumer(destination).take(destination); + return keysDynamoDb.take(destination); } long deviceId = Long.parseLong(deviceIdSelector); - return getPreKeyStoreForConsumer(destination).take(destination, deviceId); + return keysDynamoDb.take(destination, deviceId); } catch (NumberFormatException e) { throw new WebApplicationException(Response.status(422).build()); } } - - private PreKeyStore getPreKeyStoreForProducer(final Account account) { - return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_PRODUCER_EXPERIMENT) ? keysDynamoDb : keys; - } - - private PreKeyStore getPreKeyStoreForConsumer(final Account account) { - return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_CONSUMER_EXPERIMENT) ? keysDynamoDb : keys; - } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index 3d939df0f..eb4aeb1c6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -55,7 +55,6 @@ public class AccountsManager { private final FaultTolerantRedisCluster cacheCluster; private final DirectoryManager directory; private final DirectoryQueue directoryQueue; - private final Keys keys; private final KeysDynamoDb keysDynamoDb; private final MessagesManager messagesManager; private final UsernamesManager usernamesManager; @@ -74,12 +73,11 @@ public class AccountsManager { } } - public AccountsManager(Accounts accounts, DirectoryManager directory, FaultTolerantRedisCluster cacheCluster, final DirectoryQueue directoryQueue, final Keys keys, final KeysDynamoDb keysDynamoDb, final MessagesManager messagesManager, final UsernamesManager usernamesManager, final ProfilesManager profilesManager) { + public AccountsManager(Accounts accounts, DirectoryManager directory, FaultTolerantRedisCluster cacheCluster, final DirectoryQueue directoryQueue, final KeysDynamoDb keysDynamoDb, final MessagesManager messagesManager, final UsernamesManager usernamesManager, final ProfilesManager profilesManager) { this.accounts = accounts; this.directory = directory; this.cacheCluster = cacheCluster; this.directoryQueue = directoryQueue; - this.keys = keys; this.keysDynamoDb = keysDynamoDb; this.messagesManager = messagesManager; this.usernamesManager = usernamesManager; @@ -152,7 +150,6 @@ public class AccountsManager { directoryQueue.deleteAccount(account); directory.remove(account.getNumber()); profilesManager.deleteAll(account.getUuid()); - keys.delete(account); keysDynamoDb.delete(account); messagesManager.clear(account.getNumber(), account.getUuid()); redisDelete(account); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java deleted file mode 100644 index 156276d8f..000000000 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java +++ /dev/null @@ -1,164 +0,0 @@ -/* - * Copyright 2013-2020 Signal Messenger, LLC - * SPDX-License-Identifier: AGPL-3.0-only - */ -package org.whispersystems.textsecuregcm.storage; - -import com.codahale.metrics.Counter; -import com.codahale.metrics.Meter; -import com.codahale.metrics.MetricRegistry; -import com.codahale.metrics.SharedMetricRegistries; -import com.codahale.metrics.Timer; -import io.github.resilience4j.retry.Retry; -import org.jdbi.v3.core.JdbiException; -import org.jdbi.v3.core.statement.PreparedBatch; -import org.jdbi.v3.core.statement.UnableToExecuteStatementException; -import org.jdbi.v3.core.transaction.SerializableTransactionRunner; -import org.jdbi.v3.core.transaction.TransactionIsolationLevel; -import org.whispersystems.textsecuregcm.configuration.RetryConfiguration; -import org.whispersystems.textsecuregcm.entities.PreKey; -import org.whispersystems.textsecuregcm.storage.mappers.KeyRecordRowMapper; -import org.whispersystems.textsecuregcm.util.Constants; - -import java.util.Collections; -import java.util.LinkedList; -import java.util.List; -import java.util.function.Supplier; - -import static com.codahale.metrics.MetricRegistry.name; - -public class Keys implements PreKeyStore { - - private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME); - private final Meter fallbackMeter = metricRegistry.meter(name(Keys.class, "fallback")); - private final Timer storeTimer = metricRegistry.timer(name(Keys.class, "store" )); - private final Timer getDevicetTimer = metricRegistry.timer(name(Keys.class, "getDevice")); - private final Timer getTimer = metricRegistry.timer(name(Keys.class, "get" )); - private final Timer getCountTimer = metricRegistry.timer(name(Keys.class, "getCount" )); - private final Timer vacuumTimer = metricRegistry.timer(name(Keys.class, "vacuum" )); - - private final FaultTolerantDatabase database; - private final Retry retry; - - public Keys(FaultTolerantDatabase database, RetryConfiguration retryConfiguration) { - this.database = database; - this.database.getDatabase().registerRowMapper(new KeyRecordRowMapper()); - this.database.getDatabase().setTransactionHandler(new SerializableTransactionRunner()); - this.database.getDatabase().getConfig(SerializableTransactionRunner.Configuration.class).setMaxRetries(10); - - this.retry = Retry.of("keys", retryConfiguration.toRetryConfigBuilder().build()); - } - - @Override - public void store(Account account, long deviceId, List keys) { - final String number = account.getNumber(); - - retry.executeRunnable(() -> { - database.use(jdbi -> jdbi.useTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> { - try (Timer.Context ignored = storeTimer.time()) { - PreparedBatch preparedBatch = handle.prepareBatch("INSERT INTO keys (number, device_id, key_id, public_key) VALUES (:number, :device_id, :key_id, :public_key)"); - - for (PreKey key : keys) { - preparedBatch.bind("number", number) - .bind("device_id", deviceId) - .bind("key_id", key.getKeyId()) - .bind("public_key", key.getPublicKey()) - .add(); - } - - handle.createUpdate("DELETE FROM keys WHERE number = :number AND device_id = :device_id") - .bind("number", number) - .bind("device_id", deviceId) - .execute(); - - preparedBatch.execute(); - } - })); - }); - } - - @Override - public List take(Account account, long deviceId) { - /* - final String number = account.getNumber(); - - try { - return database.with(jdbi -> jdbi.inTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> { - try (Timer.Context ignored = getDevicetTimer.time()) { - return handle.createQuery("DELETE FROM keys WHERE id IN (SELECT id FROM keys WHERE number = :number AND device_id = :device_id ORDER BY key_id ASC LIMIT 1) RETURNING *") - .bind("number", number) - .bind("device_id", deviceId) - .mapTo(KeyRecord.class) - .list(); - } - })); - } catch (JdbiException e) { - // TODO 2021-01-13 Replace this with a retry once desktop clients better handle HTTP/500 responses - fallbackMeter.mark(); - return Collections.emptyList(); - } */ - - // 2021-01-15 Emergency service recovery measure - return new LinkedList<>(); - } - - @Override - public List take(Account account) { - /* - final String number = account.getNumber(); - - try { - return database.with(jdbi -> jdbi.inTransaction(TransactionIsolationLevel.SERIALIZABLE, handle -> { - try (Timer.Context ignored = getTimer.time()) { - return handle.createQuery("DELETE FROM keys WHERE id IN (SELECT DISTINCT ON (number, device_id) id FROM keys WHERE number = :number ORDER BY number, device_id, key_id ASC) RETURNING *") - .bind("number", number) - .mapTo(KeyRecord.class) - .list(); - } - })); - } catch (JdbiException e) { - // TODO 2021-01-13 Replace this with a retry once desktop clients better handle HTTP/500 responses - fallbackMeter.mark(); - return Collections.emptyList(); - } */ - - // 2021-01-15 Emergency service recovery measure - return new LinkedList<>(); - } - - @Override - public int getCount(Account account, long deviceId) { - final String number = account.getNumber(); - - return database.with(jdbi -> jdbi.withHandle(handle -> { - try (Timer.Context ignored = getCountTimer.time()) { - return handle.createQuery("SELECT COUNT(*) FROM keys WHERE number = :number AND device_id = :device_id") - .bind("number", number) - .bind("device_id", deviceId) - .mapTo(Integer.class) - .findOnly(); - } - })); - } - - public void delete(final Account account) { - final String number = account.getNumber(); - - database.use(jdbi -> jdbi.useHandle(handle -> { - try (Timer.Context ignored = getCountTimer.time()) { - handle.createUpdate("DELETE FROM keys WHERE number = :number") - .bind("number", number) - .execute(); - } - })); - } - - public void vacuum() { - database.use(jdbi -> jdbi.useHandle(handle -> { - try (Timer.Context ignored = vacuumTimer.time()) { - handle.execute("VACUUM keys"); - } - })); - } - -} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java index c95cbb71f..009afaf60 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java @@ -31,7 +31,7 @@ import java.util.UUID; import static com.codahale.metrics.MetricRegistry.name; -public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { +public class KeysDynamoDb extends AbstractDynamoDbStore { private final Table table; @@ -53,7 +53,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { this.table = dynamoDB.getTable(tableName); } - @Override public void store(final Account account, final long deviceId, final List keys) { STORE_KEYS_TIMER.record(() -> { delete(account, deviceId); @@ -70,7 +69,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { }); } - @Override public List take(final Account account, final long deviceId) { return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> { final byte[] partitionKey = getPartitionKey(account.getUuid()); @@ -106,7 +104,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { }); } - @Override public List take(final Account account) { return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { final List keyRecords = new ArrayList<>(); @@ -119,7 +116,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { }); } - @Override public int getCount(final Account account, final long deviceId) { return GET_KEY_COUNT_TIMER.record(() -> { final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)") @@ -133,7 +129,6 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { }); } - @Override public void delete(final Account account) { DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> { final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid") diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PreKeyStore.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/PreKeyStore.java deleted file mode 100644 index 35551789c..000000000 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PreKeyStore.java +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Copyright 2021 Signal Messenger, LLC - * SPDX-License-Identifier: AGPL-3.0-only - */ - -package org.whispersystems.textsecuregcm.storage; - -import org.whispersystems.textsecuregcm.entities.PreKey; - -import java.util.List; - -public interface PreKeyStore { - - void store(Account account, long deviceId, List keys); - - int getCount(Account account, long deviceId); - - List take(Account account, long deviceId); - - List take(Account account); - - void delete(Account account); -} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java index 880b91acb..9022d0629 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java @@ -33,7 +33,6 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.DirectoryManager; import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; -import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.Messages; import org.whispersystems.textsecuregcm.storage.MessagesCache; @@ -107,21 +106,19 @@ public class DeleteUserCommand extends EnvironmentCommand account = accountsManager.get(user); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/VacuumCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/VacuumCommand.java index 1356cd695..78a3bb8ab 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/VacuumCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/VacuumCommand.java @@ -14,7 +14,6 @@ import org.whispersystems.textsecuregcm.configuration.DatabaseConfiguration; import org.whispersystems.textsecuregcm.storage.Accounts; import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; import org.whispersystems.textsecuregcm.storage.FeatureFlags; -import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.Messages; import org.whispersystems.textsecuregcm.storage.PendingAccounts; @@ -46,7 +45,6 @@ public class VacuumCommand extends ConfiguredCommand FaultTolerantDatabase messageDatabase = new FaultTolerantDatabase("message_database_vacuum", messageJdbi, messageDbConfig.getCircuitBreakerConfiguration()); Accounts accounts = new Accounts(accountDatabase); - Keys keys = new Keys(accountDatabase, config.getAccountsDatabaseConfiguration().getKeyOperationRetryConfiguration()); PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase); Messages messages = new Messages(messageDatabase); FeatureFlags featureFlags = new FeatureFlags(accountDatabase); @@ -57,9 +55,6 @@ public class VacuumCommand extends ConfiguredCommand logger.info("Vacuuming pending_accounts..."); pendingAccounts.vacuum(); - logger.info("Vacuuming keys..."); - keys.vacuum(); - logger.info("Vacuuming messages..."); messages.vacuum(); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java index af2c64812..954504331 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java @@ -29,7 +29,6 @@ import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.KeyRecord; -import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.tests.util.AuthHelper; @@ -71,12 +70,10 @@ public class KeysControllerTest { private final SignedPreKey SAMPLE_SIGNED_KEY3 = new SignedPreKey( 3333, "barfoo", "sig33" ); private final SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid"); - private final Keys keys = mock(Keys.class ); private final KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class ); private final AccountsManager accounts = mock(AccountsManager.class ); private final DirectoryQueue directoryQueue = mock(DirectoryQueue.class ); private final Account existsAccount = mock(Account.class ); - private final ExperimentEnrollmentManager experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class); private RateLimiters rateLimiters = mock(RateLimiters.class); private RateLimiter rateLimiter = mock(RateLimiter.class ); @@ -86,7 +83,7 @@ public class KeysControllerTest { .addProvider(AuthHelper.getAuthFilter()) .addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(Account.class, DisabledPermittedAccount.class))) .setTestContainerFactory(new GrizzlyWebTestContainerFactory()) - .addResource(new KeysController(rateLimiters, keys, keysDynamoDb, accounts, directoryQueue, experimentEnrollmentManager)) + .addResource(new KeysController(rateLimiters, keysDynamoDb, accounts, directoryQueue)) .build(); @Before @@ -145,16 +142,16 @@ public class KeysControllerTest { List singleDevice = new LinkedList<>(); singleDevice.add(SAMPLE_KEY); - when(keys.take(eq(existsAccount), eq(1L))).thenReturn(singleDevice); + when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(singleDevice); List multiDevice = new LinkedList<>(); multiDevice.add(SAMPLE_KEY); multiDevice.add(SAMPLE_KEY2); multiDevice.add(SAMPLE_KEY3); multiDevice.add(SAMPLE_KEY4); - when(keys.take(existsAccount)).thenReturn(multiDevice); + when(keysDynamoDb.take(existsAccount)).thenReturn(multiDevice); - when(keys.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5); + when(keysDynamoDb.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5); when(AuthHelper.VALID_DEVICE.getSignedPreKey()).thenReturn(VALID_DEVICE_SIGNED_KEY); when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null); @@ -171,7 +168,7 @@ public class KeysControllerTest { assertThat(result.getCount()).isEqualTo(4); - verify(keys).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); + verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); } @Test @@ -185,7 +182,7 @@ public class KeysControllerTest { assertThat(result.getCount()).isEqualTo(4); - verify(keys).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); + verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L)); } @@ -285,8 +282,8 @@ public class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keys).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(keysDynamoDb); } @Test @@ -303,8 +300,8 @@ public class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keys).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(keysDynamoDb); } @@ -322,8 +319,8 @@ public class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keys).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(keysDynamoDb); } @Test @@ -340,8 +337,8 @@ public class KeysControllerTest { assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey()); - verify(keys).take(eq(existsAccount), eq(1L)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount), eq(1L)); + verifyNoMoreInteractions(keysDynamoDb); } @@ -354,7 +351,7 @@ public class KeysControllerTest { .get(); assertThat(response.getStatus()).isEqualTo(401); - verifyNoMoreInteractions(keys); + verifyNoMoreInteractions(keysDynamoDb); } @Test @@ -366,7 +363,7 @@ public class KeysControllerTest { .get(); assertThat(response.getStatus()).isEqualTo(401); - verifyNoMoreInteractions(keys); + verifyNoMoreInteractions(keysDynamoDb); } @@ -416,8 +413,8 @@ public class KeysControllerTest { assertThat(signedPreKey).isNull(); assertThat(deviceId).isEqualTo(4); - verify(keys).take(eq(existsAccount)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount)); + verifyNoMoreInteractions(keysDynamoDb); } @Test @@ -466,8 +463,8 @@ public class KeysControllerTest { assertThat(signedPreKey).isNull(); assertThat(deviceId).isEqualTo(4); - verify(keys).take(eq(existsAccount)); - verifyNoMoreInteractions(keys); + verify(keysDynamoDb).take(eq(existsAccount)); + verifyNoMoreInteractions(keysDynamoDb); } @@ -535,7 +532,7 @@ public class KeysControllerTest { assertThat(response.getStatus()).isEqualTo(204); ArgumentCaptor listCaptor = ArgumentCaptor.forClass(List.class); - verify(keys).store(eq(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture()); + verify(keysDynamoDb).store(eq(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture()); List capturedList = listCaptor.getValue(); assertThat(capturedList.size()).isEqualTo(1); @@ -569,7 +566,7 @@ public class KeysControllerTest { assertThat(response.getStatus()).isEqualTo(204); ArgumentCaptor listCaptor = ArgumentCaptor.forClass(List.class); - verify(keys).store(eq(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture()); + verify(keysDynamoDb).store(eq(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture()); List capturedList = listCaptor.getValue(); assertThat(capturedList.size()).isEqualTo(1); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java index 72dd3b989..55b8ad100 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java @@ -14,7 +14,6 @@ import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.Accounts; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.DirectoryManager; -import org.whispersystems.textsecuregcm.storage.Keys; import org.whispersystems.textsecuregcm.storage.KeysDynamoDb; import org.whispersystems.textsecuregcm.storage.MessagesManager; import org.whispersystems.textsecuregcm.storage.ProfilesManager; @@ -45,7 +44,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -56,7 +54,7 @@ public class AccountsManagerTest { when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(uuid.toString()); when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}"); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional account = accountsManager.get("+14152222222"); assertTrue(account.isPresent()); @@ -76,7 +74,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -86,7 +83,7 @@ public class AccountsManagerTest { when(commands.get(eq("Account3::" + uuid.toString()))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}"); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional account = accountsManager.get(uuid); assertTrue(account.isPresent()); @@ -107,7 +104,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -118,7 +114,7 @@ public class AccountsManagerTest { when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(null); when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account)); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional retrieved = accountsManager.get("+14152222222"); assertTrue(retrieved.isPresent()); @@ -140,7 +136,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -151,7 +146,7 @@ public class AccountsManagerTest { when(commands.get(eq("Account3::" + uuid))).thenReturn(null); when(accounts.get(eq(uuid))).thenReturn(Optional.of(account)); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional retrieved = accountsManager.get(uuid); assertTrue(retrieved.isPresent()); @@ -173,7 +168,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -184,7 +178,7 @@ public class AccountsManagerTest { when(commands.get(eq("AccountMap::+14152222222"))).thenThrow(new RedisException("Connection lost!")); when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account)); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional retrieved = accountsManager.get("+14152222222"); assertTrue(retrieved.isPresent()); @@ -206,7 +200,6 @@ public class AccountsManagerTest { Accounts accounts = mock(Accounts.class); DirectoryManager directoryManager = mock(DirectoryManager.class); DirectoryQueue directoryQueue = mock(DirectoryQueue.class); - Keys keys = mock(Keys.class); KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); MessagesManager messagesManager = mock(MessagesManager.class); UsernamesManager usernamesManager = mock(UsernamesManager.class); @@ -217,7 +210,7 @@ public class AccountsManagerTest { when(commands.get(eq("Account3::" + uuid))).thenThrow(new RedisException("Connection lost!")); when(accounts.get(eq(uuid))).thenReturn(Optional.of(account)); - AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keys, keysDynamoDb, messagesManager, usernamesManager, profilesManager); + AccountsManager accountsManager = new AccountsManager(accounts, directoryManager, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager); Optional retrieved = accountsManager.get(uuid); assertTrue(retrieved.isPresent()); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java deleted file mode 100644 index 437188628..000000000 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java +++ /dev/null @@ -1,462 +0,0 @@ -/* - * Copyright 2013-2020 Signal Messenger, LLC - * SPDX-License-Identifier: AGPL-3.0-only - */ - -package org.whispersystems.textsecuregcm.tests.storage; - -import com.opentable.db.postgres.embedded.LiquibasePreparer; -import com.opentable.db.postgres.junit.EmbeddedPostgresRules; -import com.opentable.db.postgres.junit.PreparedDbRule; -import io.github.resilience4j.circuitbreaker.CallNotPermittedException; -import org.jdbi.v3.core.HandleCallback; -import org.jdbi.v3.core.HandleConsumer; -import org.jdbi.v3.core.Jdbi; -import org.jdbi.v3.core.statement.UnableToExecuteStatementException; -import org.jdbi.v3.core.transaction.SerializableTransactionRunner; -import org.jdbi.v3.core.transaction.TransactionException; -import org.jdbi.v3.core.transaction.TransactionIsolationLevel; -import org.junit.Before; -import org.junit.Ignore; -import org.junit.Rule; -import org.junit.Test; -import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration; -import org.whispersystems.textsecuregcm.configuration.AccountsDatabaseConfiguration; -import org.whispersystems.textsecuregcm.configuration.RetryConfiguration; -import org.whispersystems.textsecuregcm.entities.PreKey; -import org.whispersystems.textsecuregcm.storage.Account; -import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase; -import org.whispersystems.textsecuregcm.storage.KeyRecord; -import org.whispersystems.textsecuregcm.storage.Keys; - -import java.sql.PreparedStatement; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.Collections; -import java.util.LinkedList; -import java.util.List; - -import static org.assertj.core.api.AssertionsForClassTypes.assertThat; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class KeysTest { - - @Rule - public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(LiquibasePreparer.forClasspathLocation("accountsdb.xml")); - - private Account firstAccount; - private Account secondAccount; - private Keys keys; - - @Before - public void setup() { - FaultTolerantDatabase faultTolerantDatabase = new FaultTolerantDatabase("keysTest", - Jdbi.create(db.getTestDatabase()), - new CircuitBreakerConfiguration()); - - this.keys = new Keys(faultTolerantDatabase, new RetryConfiguration()); - - this.firstAccount = mock(Account.class); - this.secondAccount = mock(Account.class); - - when(firstAccount.getNumber()).thenReturn("+14152222222"); - when(secondAccount.getNumber()).thenReturn("+14151111111"); - } - - - @Test - public void testPopulateKeys() throws SQLException { - List deviceOnePreKeys = new LinkedList<>(); - List deviceTwoPreKeys = new LinkedList<>(); - - List oldAnotherDeviceOnePrKeys = new LinkedList<>(); - List anotherDeviceOnePreKeys = new LinkedList<>(); - List anotherDeviceTwoPreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i)); - } - - for (int i=1;i<=100;i++) { - oldAnotherDeviceOnePrKeys.add(new PreKey(i, "OldPublicKey" + i)); - anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i)); - anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i)); - } - - keys.store(firstAccount, 1, deviceOnePreKeys); - keys.store(firstAccount, 2, deviceTwoPreKeys); - - keys.store(secondAccount, 1, oldAnotherDeviceOnePrKeys); - keys.store(secondAccount, 1, anotherDeviceOnePreKeys); - keys.store(secondAccount, 2, anotherDeviceTwoPreKeys); - - PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("SELECT * FROM keys WHERE number = ? AND device_id = ? ORDER BY key_id"); - verifyStoredState(statement, firstAccount, 1); - verifyStoredState(statement, firstAccount, 2); - verifyStoredState(statement, secondAccount, 1); - verifyStoredState(statement, secondAccount, 2); - } - - @Test - public void testKeyCount() { - List deviceOnePreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - } - - - keys.store(firstAccount, 1, deviceOnePreKeys); - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100); - } - - @Ignore - @Test - public void testGetForDevice() { - List deviceOnePreKeys = new LinkedList<>(); - List deviceTwoPreKeys = new LinkedList<>(); - - List anotherDeviceOnePreKeys = new LinkedList<>(); - List anotherDeviceTwoPreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i)); - } - - for (int i=1;i<=100;i++) { - anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i)); - anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i)); - } - - keys.store(firstAccount, 1, deviceOnePreKeys); - keys.store(firstAccount, 2, deviceTwoPreKeys); - - keys.store(secondAccount, 1, anotherDeviceOnePreKeys); - keys.store(secondAccount, 2, anotherDeviceTwoPreKeys); - - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100); - List records = keys.take(firstAccount, 1); - - assertThat(records.size()).isEqualTo(1); - assertThat(records.get(0).getKeyId()).isEqualTo(1); - assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device1PublicKey1"); - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(99); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100); - - records = keys.take(firstAccount, 1); - - assertThat(records.size()).isEqualTo(1); - assertThat(records.get(0).getKeyId()).isEqualTo(2); - assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device1PublicKey2"); - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100); - - records = keys.take(firstAccount, 2); - - assertThat(records.size()).isEqualTo(1); - assertThat(records.get(0).getKeyId()).isEqualTo(1); - assertThat(records.get(0).getPublicKey()).isEqualTo("+14152222222Device2PublicKey1"); - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(99); - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100); - } - - @Ignore - @Test - public void testGetForAllDevices() { - List deviceOnePreKeys = new LinkedList<>(); - List deviceTwoPreKeys = new LinkedList<>(); - - List anotherDeviceOnePreKeys = new LinkedList<>(); - List anotherDeviceTwoPreKeys = new LinkedList<>(); - List anotherDeviceThreePreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i)); - } - - for (int i=1;i<=100;i++) { - anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i)); - anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i)); - anotherDeviceThreePreKeys.add(new PreKey(i, "+14151111111Device3PublicKey" + i)); - } - - keys.store(firstAccount, 1, deviceOnePreKeys); - keys.store(firstAccount, 2, deviceTwoPreKeys); - - keys.store(secondAccount, 1, anotherDeviceOnePreKeys); - keys.store(secondAccount, 2, anotherDeviceTwoPreKeys); - keys.store(secondAccount, 3, anotherDeviceThreePreKeys); - - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100); - - List records = keys.take(firstAccount); - - assertThat(records.size()).isEqualTo(2); - assertThat(records.get(0).getKeyId()).isEqualTo(1); - assertThat(records.get(1).getKeyId()).isEqualTo(1); - - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device1PublicKey1"))).isTrue(); - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device2PublicKey1"))).isTrue(); - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(99); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(99); - - records = keys.take(firstAccount); - - assertThat(records.size()).isEqualTo(2); - assertThat(records.get(0).getKeyId()).isEqualTo(2); - assertThat(records.get(1).getKeyId()).isEqualTo(2); - - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device1PublicKey2"))).isTrue(); - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14152222222Device2PublicKey2"))).isTrue(); - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(98); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(98); - - - records = keys.take(secondAccount); - - assertThat(records.size()).isEqualTo(3); - assertThat(records.get(0).getKeyId()).isEqualTo(1); - assertThat(records.get(1).getKeyId()).isEqualTo(1); - assertThat(records.get(2).getKeyId()).isEqualTo(1); - - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device1PublicKey1"))).isTrue(); - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device2PublicKey1"))).isTrue(); - assertThat(records.stream().anyMatch(record -> record.getPublicKey().equals("+14151111111Device3PublicKey1"))).isTrue(); - - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(99); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(99); - assertThat(keys.getCount(secondAccount, 3)).isEqualTo(99); - } - - @Ignore - @Test - public void testGetForAllDevicesParallel() throws InterruptedException { - List deviceOnePreKeys = new LinkedList<>(); - List deviceTwoPreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i)); - } - - keys.store(firstAccount, 1, deviceOnePreKeys); - keys.store(firstAccount, 2, deviceTwoPreKeys); - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100); - - List threads = new LinkedList<>(); - - for (int i=0;i<20;i++) { - Thread thread = new Thread(() -> { - List results = null; - final int MAX_RETRIES = 5; - for (int retryAttempt = 0; results == null && retryAttempt < MAX_RETRIES; ++retryAttempt) { - try { - results = keys.take(firstAccount); - } catch (UnableToExecuteStatementException e) { - if (retryAttempt == MAX_RETRIES - 1) { - throw e; - } - } - } - assertThat(results).isNotNull(); - assertThat(results.size()).isEqualTo(2); - }); - thread.start(); - threads.add(thread); - } - - for (Thread thread : threads) { - thread.join(); - } - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(80); - assertThat(keys.getCount(firstAccount,2)).isEqualTo(80); - } - - @Test - public void testDelete() { - List deviceOnePreKeys = new LinkedList<>(); - List deviceTwoPreKeys = new LinkedList<>(); - - List anotherDeviceOnePreKeys = new LinkedList<>(); - List anotherDeviceTwoPreKeys = new LinkedList<>(); - List anotherDeviceThreePreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - deviceTwoPreKeys.add(new PreKey(i, "+14152222222Device2PublicKey" + i)); - } - - for (int i=1;i<=100;i++) { - anotherDeviceOnePreKeys.add(new PreKey(i, "+14151111111Device1PublicKey" + i)); - anotherDeviceTwoPreKeys.add(new PreKey(i, "+14151111111Device2PublicKey" + i)); - anotherDeviceThreePreKeys.add(new PreKey(i, "+14151111111Device3PublicKey" + i)); - } - - keys.store(firstAccount, 1, deviceOnePreKeys); - keys.store(firstAccount, 2, deviceTwoPreKeys); - - keys.store(secondAccount, 1, anotherDeviceOnePreKeys); - keys.store(secondAccount, 2, anotherDeviceTwoPreKeys); - keys.store(secondAccount, 3, anotherDeviceThreePreKeys); - - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 3)).isEqualTo(100); - - keys.delete(firstAccount); - - assertThat(keys.getCount(firstAccount, 1)).isEqualTo(0); - assertThat(keys.getCount(firstAccount, 2)).isEqualTo(0); - assertThat(keys.getCount(secondAccount, 1)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 2)).isEqualTo(100); - assertThat(keys.getCount(secondAccount, 3)).isEqualTo(100); - } - - @Test - public void testEmptyKeyGet() { - List records = keys.take(firstAccount); - - assertThat(records.isEmpty()).isTrue(); - } - - @Test - public void testVacuum() { - keys.vacuum(); - } - - @Test - public void testBreaker() throws InterruptedException { - Jdbi jdbi = mock(Jdbi.class); - doThrow(new TransactionException("Database error!")).when(jdbi).useTransaction(any(TransactionIsolationLevel.class), any(HandleConsumer.class)); - when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class)); - - CircuitBreakerConfiguration configuration = new CircuitBreakerConfiguration(); - configuration.setWaitDurationInOpenStateInSeconds(1); - configuration.setRingBufferSizeInHalfOpenState(1); - configuration.setRingBufferSizeInClosedState(2); - configuration.setFailureRateThreshold(50); - - RetryConfiguration retryConfiguration = new RetryConfiguration(); - retryConfiguration.setMaxAttempts(1); - - Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, configuration), retryConfiguration); - - List deviceOnePreKeys = new LinkedList<>(); - - for (int i=1;i<=100;i++) { - deviceOnePreKeys.add(new PreKey(i, "+14152222222Device1PublicKey" + i)); - } - - try { - keys.store(firstAccount, 1, deviceOnePreKeys); - throw new AssertionError(); - } catch (TransactionException e) { - // good - } - - try { - keys.store(firstAccount, 1, deviceOnePreKeys); - throw new AssertionError(); - } catch (TransactionException e) { - // good - } - - try { - keys.store(firstAccount, 1, deviceOnePreKeys); - throw new AssertionError(); - } catch (CallNotPermittedException e) { - // good - } - - Thread.sleep(1100); - - try { - keys.store(firstAccount, 1, deviceOnePreKeys); - throw new AssertionError(); - } catch (TransactionException e) { - // good - } - - } - - @Test - public void testRetry() { - Jdbi jdbi = mock(Jdbi.class); - doThrow(new TransactionException("Database error!")).doNothing().when(jdbi).useTransaction(any(TransactionIsolationLevel.class), any(HandleConsumer.class)); - when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class)); - - Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, new CircuitBreakerConfiguration()), new RetryConfiguration()); - - // We're happy as long as nothing throws an exception - Account account = mock(Account.class); - when(account.getNumber()).thenReturn("+18005551234"); - - keys.store(account, 1, Collections.emptyList()); - } - - @Test - public void testGetKeysWithException() { - Jdbi jdbi = mock(Jdbi.class); - when(jdbi.getConfig(any())).thenReturn(mock(SerializableTransactionRunner.Configuration.class)); - - when(jdbi.inTransaction(any(TransactionIsolationLevel.class), any(HandleCallback.class))) - .thenThrow(new TransactionException("Database error!")); - - Keys keys = new Keys(new FaultTolerantDatabase("testBreaker", jdbi, new CircuitBreakerConfiguration()), new RetryConfiguration()); - - Account account = mock(Account.class); - when(account.getNumber()).thenReturn("+18005551234"); - - assertThat(keys.take(account)).isEqualTo(Collections.emptyList()); - assertThat(keys.take(account, 1)).isEqualTo(Collections.emptyList()); - } - - private void verifyStoredState(PreparedStatement statement, Account account, int deviceId) throws SQLException { - statement.setString(1, account.getNumber()); - statement.setInt(2, deviceId); - - ResultSet resultSet = statement.executeQuery(); - int rowCount = 1; - - while (resultSet.next()) { - long keyId = resultSet.getLong("key_id"); - String publicKey = resultSet.getString("public_key"); - - - assertThat(keyId).isEqualTo(rowCount); - assertThat(publicKey).isEqualTo(account.getNumber() + "Device" + deviceId + "PublicKey" + rowCount); - - rowCount++; - } - - resultSet.close(); - - assertThat(rowCount).isEqualTo(101); - - } - -}