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

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),

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),

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())

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());
}
}

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());
}
}

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 );

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);

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);