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

s/KeysDynamoDb/Keys/

This commit is contained in:
Jon Chambers 2021-11-03 17:50:57 -04:00 committed by Jon Chambers
parent fa6e3d3690
commit 5e1334e8de
15 changed files with 212 additions and 213 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -182,7 +182,7 @@ import org.whispersystems.textsecuregcm.storage.DirectoryReconciliationClient;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagePersister;
import org.whispersystems.textsecuregcm.storage.MessagesCache;
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
@ -379,7 +379,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
Usernames usernames = new Usernames(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase);
KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName());
Keys keys = new Keys(preKeyDynamoDb, config.getKeysDynamoDbConfiguration().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
config.getMessageDynamoDbConfiguration().getTableName(),
config.getMessageDynamoDbConfiguration().getTimeToLive());
@ -467,7 +467,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
deletedAccountsLockDynamoDbClient, config.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
DeadLetterHandler deadLetterHandler = new DeadLetterHandler(accountsManager, messagesManager);
@ -649,7 +649,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
smsSender, dynamicConfigurationManager, turnTokenGenerator, config.getTestDevices(),
transitionalRecaptchaClient, gcmSender, apnSender, backupCredentialsGenerator,
verifyExperimentEnrollmentManager));
environment.jersey().register(new KeysController(rateLimiters, keysDynamoDb, accountsManager, preKeyRateLimiter, rateLimitChallengeManager));
environment.jersey().register(new KeysController(rateLimiters, keys, accountsManager, preKeyRateLimiter, rateLimitChallengeManager));
final List<Object> commonControllers = Lists.newArrayList(
new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket()),
@ -657,7 +657,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
new AttachmentControllerV3(rateLimiters, config.getGcpAttachmentsConfiguration().getDomain(), config.getGcpAttachmentsConfiguration().getEmail(), config.getGcpAttachmentsConfiguration().getMaxSizeInBytes(), config.getGcpAttachmentsConfiguration().getPathPrefix(), config.getGcpAttachmentsConfiguration().getRsaSigningKey()),
new CertificateController(new CertificateGenerator(config.getDeliveryCertificate().getCertificate(), config.getDeliveryCertificate().getPrivateKey(), config.getDeliveryCertificate().getExpiresDays()), zkAuthOperations),
new ChallengeController(rateLimitChallengeManager),
new DeviceController(pendingDevicesManager, accountsManager, messagesManager, keysDynamoDb, rateLimiters, config.getMaxDevices()),
new DeviceController(pendingDevicesManager, accountsManager, messagesManager, keys, rateLimiters, config.getMaxDevices()),
new DirectoryController(directoryCredentialsGenerator),
new DirectoryV2Controller(directoryV2CredentialsGenerator),
new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),

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

@ -41,7 +41,7 @@ import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
import org.whispersystems.textsecuregcm.util.Util;
@ -57,14 +57,14 @@ public class DeviceController {
private final StoredVerificationCodeManager pendingDevices;
private final AccountsManager accounts;
private final MessagesManager messages;
private final KeysDynamoDb keys;
private final Keys keys;
private final RateLimiters rateLimiters;
private final Map<String, Integer> maxDeviceConfiguration;
public DeviceController(StoredVerificationCodeManager pendingDevices,
AccountsManager accounts,
MessagesManager messages,
KeysDynamoDb keys,
Keys keys,
RateLimiters rateLimiters,
Map<String, Integer> maxDeviceConfiguration)
{

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

@ -44,7 +44,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.util.Util;
@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
@ -52,7 +52,7 @@ import org.whispersystems.textsecuregcm.util.Util;
public class KeysController {
private final RateLimiters rateLimiters;
private final KeysDynamoDb keysDynamoDb;
private final Keys keys;
private final AccountsManager accounts;
private final PreKeyRateLimiter preKeyRateLimiter;
@ -64,11 +64,11 @@ public class KeysController {
private static final String SOURCE_COUNTRY_TAG_NAME = "sourceCountry";
private static final String INTERNATIONAL_TAG_NAME = "international";
public KeysController(RateLimiters rateLimiters, KeysDynamoDb keysDynamoDb, AccountsManager accounts,
public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts,
PreKeyRateLimiter preKeyRateLimiter,
RateLimitChallengeManager rateLimitChallengeManager) {
this.rateLimiters = rateLimiters;
this.keysDynamoDb = keysDynamoDb;
this.keys = keys;
this.accounts = accounts;
this.preKeyRateLimiter = preKeyRateLimiter;
this.rateLimitChallengeManager = rateLimitChallengeManager;
@ -77,7 +77,7 @@ public class KeysController {
@GET
@Produces(MediaType.APPLICATION_JSON)
public PreKeyCount getStatus(@Auth AuthenticatedAccount auth) {
int count = keysDynamoDb.getCount(auth.getAccount(), auth.getAuthenticatedDevice().getId());
int count = keys.getCount(auth.getAccount(), auth.getAuthenticatedDevice().getId());
if (count > 0) {
count = count - 1;
@ -109,7 +109,7 @@ public class KeysController {
});
}
keysDynamoDb.store(account, device.getId(), preKeys.getPreKeys());
keys.store(account, device.getId(), preKeys.getPreKeys());
}
@Timed
@ -213,12 +213,12 @@ public class KeysController {
private Map<Long, PreKey> getLocalKeys(Account destination, String deviceIdSelector) {
try {
if (deviceIdSelector.equals("*")) {
return keysDynamoDb.take(destination);
return keys.take(destination);
}
long deviceId = Long.parseLong(deviceIdSelector);
return keysDynamoDb.take(destination, deviceId)
return keys.take(destination, deviceId)
.map(preKey -> Map.of(deviceId, preKey))
.orElse(Collections.emptyMap());
} catch (NumberFormatException e) {

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

@ -68,7 +68,7 @@ public class AccountsManager {
private final FaultTolerantRedisCluster cacheCluster;
private final DeletedAccountsManager deletedAccountsManager;
private final DirectoryQueue directoryQueue;
private final KeysDynamoDb keysDynamoDb;
private final Keys keys;
private final MessagesManager messagesManager;
private final UsernamesManager usernamesManager;
private final ProfilesManager profilesManager;
@ -96,7 +96,7 @@ public class AccountsManager {
final FaultTolerantRedisCluster cacheCluster,
final DeletedAccountsManager deletedAccountsManager,
final DirectoryQueue directoryQueue,
final KeysDynamoDb keysDynamoDb,
final Keys keys,
final MessagesManager messagesManager,
final UsernamesManager usernamesManager,
final ProfilesManager profilesManager,
@ -110,7 +110,7 @@ public class AccountsManager {
this.cacheCluster = cacheCluster;
this.deletedAccountsManager = deletedAccountsManager;
this.directoryQueue = directoryQueue;
this.keysDynamoDb = keysDynamoDb;
this.keys = keys;
this.messagesManager = messagesManager;
this.usernamesManager = usernamesManager;
this.profilesManager = profilesManager;
@ -178,7 +178,7 @@ public class AccountsManager {
// account and need to clear out messages and keys that may have been stored for the old account.
if (!originalUuid.equals(actualUuid)) {
messagesManager.clear(actualUuid);
keysDynamoDb.delete(actualUuid);
keys.delete(actualUuid);
profilesManager.deleteAll(actualUuid);
}
@ -437,11 +437,11 @@ public class AccountsManager {
usernamesManager.delete(account.getUuid());
profilesManager.deleteAll(account.getUuid());
keysDynamoDb.delete(account.getUuid());
keys.delete(account.getUuid());
messagesManager.clear(account.getUuid());
account.getPhoneNumberIdentifier().ifPresent(pni -> {
keysDynamoDb.delete(pni);
keys.delete(pni);
messagesManager.clear(pni);
});

20
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java → service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java
View File

@ -32,7 +32,7 @@ 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 Keys extends AbstractDynamoDbStore {
private final String tableName;
@ -40,16 +40,16 @@ public class KeysDynamoDb extends AbstractDynamoDbStore {
static final String KEY_DEVICE_ID_KEY_ID = "DK";
static final String KEY_PUBLIC_KEY = "P";
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_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 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 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 Timer STORE_KEYS_TIMER = Metrics.timer(name(Keys.class, "storeKeys"));
private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "takeKeyForDevice"));
private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "takeKeyForAccount"));
private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(Keys.class, "getKeyCount"));
private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForDevice"));
private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForAccount"));
private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(Keys.class, "contestedKeys"));
private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(Keys.class, "keyCount"));
public KeysDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
public Keys(final DynamoDbClient dynamoDB, final String tableName) {
super(dynamoDB);
this.tableName = tableName;
}

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

@ -44,7 +44,7 @@ import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagesCache;
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
@ -167,7 +167,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
Usernames usernames = new Usernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeysDynamoDb,
Keys keys = new Keys(preKeysDynamoDb,
configuration.getKeysDynamoDbConfiguration().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
configuration.getMessageDynamoDbConfiguration().getTableName(),
@ -207,7 +207,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
for (String user : users) {

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

@ -42,7 +42,7 @@ import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagesCache;
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
@ -172,7 +172,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
Usernames usernames = new Usernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
KeysDynamoDb keysDynamoDb = new KeysDynamoDb(preKeysDynamoDb,
Keys keys = new Keys(preKeysDynamoDb,
configuration.getKeysDynamoDbConfiguration().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb,
configuration.getMessageDynamoDbConfiguration().getTableName(),
@ -210,7 +210,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
deletedAccountsManager, directoryQueue, keys, messagesManager, usernamesManager, profilesManager,
pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
Optional<Account> maybeAccount;

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

@ -189,7 +189,7 @@ class AccountsManagerChangeNumberIntegrationTest {
CACHE_CLUSTER_EXTENSION.getRedisCluster(),
deletedAccountsManager,
mock(DirectoryQueue.class),
mock(KeysDynamoDb.class),
mock(Keys.class),
mock(MessagesManager.class),
mock(UsernamesManager.class),
mock(ProfilesManager.class),

3
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java
View File

@ -36,7 +36,6 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.ArgumentCaptor;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
import org.whispersystems.textsecuregcm.entities.AccountAttributes;
@ -146,7 +145,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
RedisClusterHelper.buildMockRedisCluster(commands),
deletedAccountsManager,
mock(DirectoryQueue.class),
mock(KeysDynamoDb.class),
mock(Keys.class),
mock(MessagesManager.class),
mock(UsernamesManager.class),
mock(ProfilesManager.class),

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

@ -22,15 +22,15 @@ public class KeysDynamoDbRule extends LocalDynamoDbRule {
getDynamoDbClient().createTable(CreateTableRequest.builder()
.tableName(TABLE_NAME)
.keySchema(
KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE)
KeySchemaElement.builder().attributeName(Keys.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName(Keys.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID)
.attributeName(Keys.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build(),
AttributeDefinition.builder()
.attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID)
.attributeName(Keys.KEY_DEVICE_ID_KEY_ID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build())

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

@ -1,144 +0,0 @@
/*
* Copyright 2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.whispersystems.textsecuregcm.entities.PreKey;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
public class KeysDynamoDbTest {
private Account account;
private KeysDynamoDb keysDynamoDb;
@ClassRule
public static KeysDynamoDbRule dynamoDbRule = new KeysDynamoDbRule();
private static final String ACCOUNT_NUMBER = "+18005551234";
private static final long DEVICE_ID = 1L;
@Before
public void setup() {
keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME);
account = mock(Account.class);
when(account.getNumber()).thenReturn(ACCOUNT_NUMBER);
when(account.getUuid()).thenReturn(UUID.randomUUID());
}
@Test
public void testStore() {
assertEquals("Initial pre-key count for an account should be zero",
0, keysDynamoDb.getCount(account, DEVICE_ID));
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID));
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals("Repeatedly storing same key should have no effect",
1, keysDynamoDb.getCount(account, DEVICE_ID));
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(2, "different-public-key")));
assertEquals("Inserting a new key should overwrite all prior keys for the given account/device",
1, keysDynamoDb.getCount(account, DEVICE_ID));
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(3, "third-public-key"), new PreKey(4, "fourth-public-key")));
assertEquals("Inserting multiple new keys should overwrite all prior keys for the given account/device",
2, keysDynamoDb.getCount(account, DEVICE_ID));
}
@Test
public void testTakeAccount() {
final Device firstDevice = mock(Device.class);
final Device secondDevice = mock(Device.class);
when(firstDevice.getId()).thenReturn(DEVICE_ID);
when(secondDevice.getId()).thenReturn(DEVICE_ID + 1);
when(account.getDevices()).thenReturn(Set.of(firstDevice, secondDevice));
assertEquals(Collections.emptyMap(), keysDynamoDb.take(account));
final PreKey firstDevicePreKey = new PreKey(1, "public-key");
final PreKey secondDevicePreKey = new PreKey(2, "second-key");
keysDynamoDb.store(account, DEVICE_ID, List.of(firstDevicePreKey));
keysDynamoDb.store(account, DEVICE_ID + 1, List.of(secondDevicePreKey));
final Map<Long, PreKey> expectedKeys = Map.of(DEVICE_ID, firstDevicePreKey,
DEVICE_ID + 1, secondDevicePreKey);
assertEquals(expectedKeys, keysDynamoDb.take(account));
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1));
}
@Test
public void testTakeAccountAndDeviceId() {
assertEquals(Optional.empty(), keysDynamoDb.take(account, DEVICE_ID));
final PreKey preKey = new PreKey(1, "public-key");
keysDynamoDb.store(account, DEVICE_ID, List.of(preKey, new PreKey(2, "different-pre-key")));
assertEquals(Optional.of(preKey), keysDynamoDb.take(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID));
}
@Test
public void testGetCount() {
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID));
}
@Test
public void testDeleteByAccount() {
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key")));
keysDynamoDb.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device")));
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
keysDynamoDb.delete(account.getUuid());
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID + 1));
}
@Test
public void testDeleteByAccountAndDevice() {
keysDynamoDb.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key")));
keysDynamoDb.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device")));
assertEquals(2, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
keysDynamoDb.delete(account.getUuid(), DEVICE_ID);
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
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());
}
}

144
service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysTest.java Normal file
View File

@ -0,0 +1,144 @@
/*
* Copyright 2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.whispersystems.textsecuregcm.entities.PreKey;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
public class KeysTest {
private Account account;
private Keys keys;
@ClassRule
public static KeysDynamoDbRule dynamoDbRule = new KeysDynamoDbRule();
private static final String ACCOUNT_NUMBER = "+18005551234";
private static final long DEVICE_ID = 1L;
@Before
public void setup() {
keys = new Keys(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME);
account = mock(Account.class);
when(account.getNumber()).thenReturn(ACCOUNT_NUMBER);
when(account.getUuid()).thenReturn(UUID.randomUUID());
}
@Test
public void testStore() {
assertEquals("Initial pre-key count for an account should be zero",
0, keys.getCount(account, DEVICE_ID));
keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals(1, keys.getCount(account, DEVICE_ID));
keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals("Repeatedly storing same key should have no effect",
1, keys.getCount(account, DEVICE_ID));
keys.store(account, DEVICE_ID, List.of(new PreKey(2, "different-public-key")));
assertEquals("Inserting a new key should overwrite all prior keys for the given account/device",
1, keys.getCount(account, DEVICE_ID));
keys.store(account, DEVICE_ID, List.of(new PreKey(3, "third-public-key"), new PreKey(4, "fourth-public-key")));
assertEquals("Inserting multiple new keys should overwrite all prior keys for the given account/device",
2, keys.getCount(account, DEVICE_ID));
}
@Test
public void testTakeAccount() {
final Device firstDevice = mock(Device.class);
final Device secondDevice = mock(Device.class);
when(firstDevice.getId()).thenReturn(DEVICE_ID);
when(secondDevice.getId()).thenReturn(DEVICE_ID + 1);
when(account.getDevices()).thenReturn(Set.of(firstDevice, secondDevice));
assertEquals(Collections.emptyMap(), keys.take(account));
final PreKey firstDevicePreKey = new PreKey(1, "public-key");
final PreKey secondDevicePreKey = new PreKey(2, "second-key");
keys.store(account, DEVICE_ID, List.of(firstDevicePreKey));
keys.store(account, DEVICE_ID + 1, List.of(secondDevicePreKey));
final Map<Long, PreKey> expectedKeys = Map.of(DEVICE_ID, firstDevicePreKey,
DEVICE_ID + 1, secondDevicePreKey);
assertEquals(expectedKeys, keys.take(account));
assertEquals(0, keys.getCount(account, DEVICE_ID));
assertEquals(0, keys.getCount(account, DEVICE_ID + 1));
}
@Test
public void testTakeAccountAndDeviceId() {
assertEquals(Optional.empty(), keys.take(account, DEVICE_ID));
final PreKey preKey = new PreKey(1, "public-key");
keys.store(account, DEVICE_ID, List.of(preKey, new PreKey(2, "different-pre-key")));
assertEquals(Optional.of(preKey), keys.take(account, DEVICE_ID));
assertEquals(1, keys.getCount(account, DEVICE_ID));
}
@Test
public void testGetCount() {
assertEquals(0, keys.getCount(account, DEVICE_ID));
keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key")));
assertEquals(1, keys.getCount(account, DEVICE_ID));
}
@Test
public void testDeleteByAccount() {
keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key")));
keys.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device")));
assertEquals(2, keys.getCount(account, DEVICE_ID));
assertEquals(1, keys.getCount(account, DEVICE_ID + 1));
keys.delete(account.getUuid());
assertEquals(0, keys.getCount(account, DEVICE_ID));
assertEquals(0, keys.getCount(account, DEVICE_ID + 1));
}
@Test
public void testDeleteByAccountAndDevice() {
keys.store(account, DEVICE_ID, List.of(new PreKey(1, "public-key"), new PreKey(2, "different-public-key")));
keys.store(account, DEVICE_ID + 1, List.of(new PreKey(3, "public-key-for-different-device")));
assertEquals(2, keys.getCount(account, DEVICE_ID));
assertEquals(1, keys.getCount(account, DEVICE_ID + 1));
keys.delete(account.getUuid(), DEVICE_ID);
assertEquals(0, keys.getCount(account, DEVICE_ID));
assertEquals(1, keys.getCount(account, DEVICE_ID + 1));
}
@Test
public void testSortKeyPrefix() {
AttributeValue got = Keys.getSortKeyPrefix(123);
assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, got.b().asByteArray());
}
}

6
service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java
View File

@ -52,7 +52,7 @@ import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
@ -66,7 +66,7 @@ class DeviceControllerTest {
public DumbVerificationDeviceController(StoredVerificationCodeManager pendingDevices,
AccountsManager accounts,
MessagesManager messages,
KeysDynamoDb keys,
Keys keys,
RateLimiters rateLimiters,
Map<String, Integer> deviceConfiguration)
{
@ -82,7 +82,7 @@ class DeviceControllerTest {
private static StoredVerificationCodeManager pendingDevicesManager = mock(StoredVerificationCodeManager.class);
private static AccountsManager accountsManager = mock(AccountsManager.class );
private static MessagesManager messagesManager = mock(MessagesManager.class);
private static KeysDynamoDb keys = mock(KeysDynamoDb.class);
private static Keys keys = mock(Keys.class);
private static RateLimiters rateLimiters = mock(RateLimiters.class );
private static RateLimiter rateLimiter = mock(RateLimiter.class );
private static Account account = mock(Account.class );

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

@ -61,7 +61,7 @@ import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
@ -89,7 +89,7 @@ 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 static KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class );
private final static Keys KEYS = mock(Keys.class );
private final static AccountsManager accounts = mock(AccountsManager.class );
private final static PreKeyRateLimiter preKeyRateLimiter = mock(PreKeyRateLimiter.class );
private final static RateLimitChallengeManager rateLimitChallengeManager = mock(RateLimitChallengeManager.class );
@ -106,7 +106,7 @@ class KeysControllerTest {
.addResource(new RateLimitChallengeExceptionMapper(rateLimitChallengeManager))
.addResource(new ServerRejectedExceptionMapper())
.addResource(
new KeysController(rateLimiters, keysDynamoDb, accounts, preKeyRateLimiter, rateLimitChallengeManager))
new KeysController(rateLimiters, KEYS, accounts, preKeyRateLimiter, rateLimitChallengeManager))
.build();
@BeforeEach
@ -161,14 +161,14 @@ class KeysControllerTest {
when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter);
when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(Optional.of(SAMPLE_KEY));
when(KEYS.take(eq(existsAccount), eq(1L))).thenReturn(Optional.of(SAMPLE_KEY));
when(keysDynamoDb.take(existsAccount)).thenReturn(Map.of(1L, SAMPLE_KEY,
when(KEYS.take(existsAccount)).thenReturn(Map.of(1L, SAMPLE_KEY,
2L, SAMPLE_KEY2,
3L, SAMPLE_KEY3,
4L, SAMPLE_KEY4));
when(keysDynamoDb.getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L))).thenReturn(5);
when(KEYS.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);
@ -177,7 +177,7 @@ class KeysControllerTest {
@AfterEach
void teardown() {
reset(
keysDynamoDb,
KEYS,
accounts,
preKeyRateLimiter,
existsAccount,
@ -198,7 +198,7 @@ class KeysControllerTest {
assertThat(result.getCount()).isEqualTo(4);
verify(keysDynamoDb).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L));
verify(KEYS).getCount(eq(AuthHelper.VALID_ACCOUNT), eq(1L));
}
@ -257,8 +257,8 @@ class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keysDynamoDb);
verify(KEYS).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(KEYS);
}
@Test
@ -275,8 +275,8 @@ class KeysControllerTest {
assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey());
assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getSignedPreKey());
verify(keysDynamoDb).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(keysDynamoDb);
verify(KEYS).take(eq(existsAccount), eq(1L));
verifyNoMoreInteractions(KEYS);
}
@Test
@ -303,7 +303,7 @@ class KeysControllerTest {
.get();
assertThat(response.getStatus()).isEqualTo(401);
verifyNoMoreInteractions(keysDynamoDb);
verifyNoMoreInteractions(KEYS);
}
@Test
@ -315,7 +315,7 @@ class KeysControllerTest {
.get();
assertThat(response.getStatus()).isEqualTo(401);
verifyNoMoreInteractions(keysDynamoDb);
verifyNoMoreInteractions(KEYS);
}
@ -365,8 +365,8 @@ class KeysControllerTest {
assertThat(signedPreKey).isNull();
assertThat(deviceId).isEqualTo(4);
verify(keysDynamoDb).take(eq(existsAccount));
verifyNoMoreInteractions(keysDynamoDb);
verify(KEYS).take(eq(existsAccount));
verifyNoMoreInteractions(KEYS);
}
@ -434,7 +434,7 @@ class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class);
verify(keysDynamoDb).store(eqUuid(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture());
verify(KEYS).store(eqUuid(AuthHelper.VALID_ACCOUNT), eq(1L), listCaptor.capture());
List<PreKey> capturedList = listCaptor.getValue();
assertThat(capturedList.size()).isEqualTo(1);
@ -468,7 +468,7 @@ class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List> listCaptor = ArgumentCaptor.forClass(List.class);
verify(keysDynamoDb).store(eqUuid(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture());
verify(KEYS).store(eqUuid(AuthHelper.DISABLED_ACCOUNT), eq(1L), listCaptor.capture());
List<PreKey> capturedList = listCaptor.getValue();
assertThat(capturedList.size()).isEqualTo(1);

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

@ -56,7 +56,7 @@ import org.whispersystems.textsecuregcm.storage.ContestedOptimisticLockException
import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
import org.whispersystems.textsecuregcm.storage.ProfilesManager;
@ -69,7 +69,7 @@ class AccountsManagerTest {
private Accounts accounts;
private DeletedAccountsManager deletedAccountsManager;
private DirectoryQueue directoryQueue;
private KeysDynamoDb keys;
private Keys keys;
private MessagesManager messagesManager;
private ProfilesManager profilesManager;
@ -89,7 +89,7 @@ class AccountsManagerTest {
accounts = mock(Accounts.class);
deletedAccountsManager = mock(DeletedAccountsManager.class);
directoryQueue = mock(DirectoryQueue.class);
keys = mock(KeysDynamoDb.class);
keys = mock(Keys.class);
messagesManager = mock(MessagesManager.class);
profilesManager = mock(ProfilesManager.class);