Clarify behavioral contract of the pre-key store

This commit is contained in:
Jon Chambers
2021-02-08 11:45:57 -05:00
committed by GitHub
parent 04728ea4bc
commit aa99e202b4
6 changed files with 39 additions and 132 deletions

View File

@@ -11,8 +11,9 @@ import org.junit.Test;
import org.whispersystems.textsecuregcm.entities.PreKey;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
@@ -70,7 +71,7 @@ public class KeysDynamoDbTest {
when(secondDevice.getId()).thenReturn(DEVICE_ID + 1);
when(account.getDevices()).thenReturn(Set.of(firstDevice, secondDevice));
assertEquals(Collections.emptyList(), keysDynamoDb.take(account));
assertEquals(Collections.emptyMap(), keysDynamoDb.take(account));
final PreKey firstDevicePreKey = new PreKey(1, "public-key");
final PreKey secondDevicePreKey = new PreKey(2, "second-key");
@@ -78,23 +79,22 @@ public class KeysDynamoDbTest {
keysDynamoDb.store(account, DEVICE_ID, List.of(firstDevicePreKey));
keysDynamoDb.store(account, DEVICE_ID + 1, List.of(secondDevicePreKey));
final Set<KeyRecord> expectedKeys = Set.of(
new KeyRecord(-1, ACCOUNT_NUMBER, DEVICE_ID, firstDevicePreKey.getKeyId(), firstDevicePreKey.getPublicKey()),
new KeyRecord(-1, ACCOUNT_NUMBER, DEVICE_ID + 1, secondDevicePreKey.getKeyId(), secondDevicePreKey.getPublicKey()));
final Map<Long, PreKey> expectedKeys = Map.of(DEVICE_ID, firstDevicePreKey,
DEVICE_ID + 1, secondDevicePreKey);
assertEquals(expectedKeys, new HashSet<>(keysDynamoDb.take(account)));
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(Collections.emptyList(), keysDynamoDb.take(account, DEVICE_ID));
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(List.of(new KeyRecord(-1, ACCOUNT_NUMBER, DEVICE_ID, preKey.getKeyId(), preKey.getPublicKey())), keysDynamoDb.take(account, DEVICE_ID));
assertEquals(Optional.of(preKey), keysDynamoDb.take(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID));
}

View File

@@ -28,7 +28,6 @@ import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
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.KeysDynamoDb;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
@@ -38,6 +37,7 @@ import javax.ws.rs.core.Response;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
@@ -59,10 +59,10 @@ public class KeysControllerTest {
private static int SAMPLE_REGISTRATION_ID2 = 1002;
private static int SAMPLE_REGISTRATION_ID4 = 1555;
private final KeyRecord SAMPLE_KEY = new KeyRecord(1, EXISTS_NUMBER, Device.MASTER_ID, 1234, "test1");
private final KeyRecord SAMPLE_KEY2 = new KeyRecord(2, EXISTS_NUMBER, 2, 5667, "test3");
private final KeyRecord SAMPLE_KEY3 = new KeyRecord(3, EXISTS_NUMBER, 3, 334, "test5");
private final KeyRecord SAMPLE_KEY4 = new KeyRecord(4, EXISTS_NUMBER, 4, 336, "test6");
private final PreKey SAMPLE_KEY = new PreKey(1234, "test1");
private final PreKey SAMPLE_KEY2 = new PreKey(5667, "test3");
private final PreKey SAMPLE_KEY3 = new PreKey(334, "test5");
private final PreKey SAMPLE_KEY4 = new PreKey(336, "test6");
private final SignedPreKey SAMPLE_SIGNED_KEY = new SignedPreKey( 1111, "foofoo", "sig11" );
@@ -140,16 +140,12 @@ public class KeysControllerTest {
when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter);
List<KeyRecord> singleDevice = new LinkedList<>();
singleDevice.add(SAMPLE_KEY);
when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(singleDevice);
when(keysDynamoDb.take(eq(existsAccount), eq(1L))).thenReturn(Optional.of(SAMPLE_KEY));
List<KeyRecord> multiDevice = new LinkedList<>();
multiDevice.add(SAMPLE_KEY);
multiDevice.add(SAMPLE_KEY2);
multiDevice.add(SAMPLE_KEY3);
multiDevice.add(SAMPLE_KEY4);
when(keysDynamoDb.take(existsAccount)).thenReturn(multiDevice);
when(keysDynamoDb.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);