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

Store signed EC pre-keys in a dedicated table

This commit is contained in:
Jon Chambers 2023-06-20 10:29:03 -04:00 committed by Jon Chambers
parent 93c78b6e40
commit 1a1defb055
18 changed files with 114 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
service/config/sample.yml
View File

@ -98,6 +98,8 @@ dynamoDbTables:
generator: abcdefg12345678= # random base64-encoded binary sequence
ecKeys:
tableName: Example_Keys
ecSignedPreKeys:
tableName: Example_EC_Signed_Pre_Keys
pqKeys:
tableName: Example_PQ_Keys
pqLastResortKeys:

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

@ -316,6 +316,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
KeysManager keys = new KeysManager(
dynamoDbAsyncClient,
config.getDynamoDbTables().getEcKeys().getTableName(),
config.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
config.getDynamoDbTables().getKemKeys().getTableName(),
config.getDynamoDbTables().getKemLastResortKeys().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

9
service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbTables.java
View File

@ -51,6 +51,7 @@ public class DynamoDbTables {
private final Table deletedAccountsLock;
private final IssuedReceiptsTableConfiguration issuedReceipts;
private final Table ecKeys;
private final Table ecSignedPreKeys;
private final Table kemKeys;
private final Table kemLastResortKeys;
private final TableWithExpiration messages;
@ -72,6 +73,7 @@ public class DynamoDbTables {
@JsonProperty("deletedAccountsLock") final Table deletedAccountsLock,
@JsonProperty("issuedReceipts") final IssuedReceiptsTableConfiguration issuedReceipts,
@JsonProperty("ecKeys") final Table ecKeys,
@JsonProperty("ecSignedPreKeys") final Table ecSignedPreKeys,
@JsonProperty("pqKeys") final Table kemKeys,
@JsonProperty("pqLastResortKeys") final Table kemLastResortKeys,
@JsonProperty("messages") final TableWithExpiration messages,
@ -92,6 +94,7 @@ public class DynamoDbTables {
this.deletedAccountsLock = deletedAccountsLock;
this.issuedReceipts = issuedReceipts;
this.ecKeys = ecKeys;
this.ecSignedPreKeys = ecSignedPreKeys;
this.kemKeys = kemKeys;
this.kemLastResortKeys = kemLastResortKeys;
this.messages = messages;
@ -138,6 +141,12 @@ public class DynamoDbTables {
return ecKeys;
}
@NotNull
@Valid
public Table getEcSignedPreKeys() {
return ecSignedPreKeys;
}
@NotNull
@Valid
public Table getKemKeys() {

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

@ -341,7 +341,9 @@ public class DeviceController {
keys.delete(a.getPhoneNumberIdentifier(), device.getId());
maybeDeviceActivationRequest.ifPresent(deviceActivationRequest -> {
keys.storeEcSignedPreKeys(a.getUuid(), Map.of(device.getId(), deviceActivationRequest.aciSignedPreKey().get()));
keys.storePqLastResort(a.getUuid(), Map.of(device.getId(), deviceActivationRequest.aciPqLastResortPreKey().get()));
keys.storeEcSignedPreKeys(a.getPhoneNumberIdentifier(), Map.of(device.getId(), deviceActivationRequest.pniSignedPreKey().get()));
keys.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(device.getId(), deviceActivationRequest.pniPqLastResortPreKey().get()));
});

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

@ -167,7 +167,7 @@ public class KeysController {
keys.store(
getIdentifier(account, identityType), device.getId(),
preKeys.getPreKeys(), preKeys.getPqPreKeys(), preKeys.getPqLastResortPreKey());
preKeys.getPreKeys(), preKeys.getPqPreKeys(), preKeys.getSignedPreKey(), preKeys.getPqLastResortPreKey());
}
@Timed

2
service/src/main/java/org/whispersystems/textsecuregcm/controllers/RegistrationController.java
View File

@ -176,7 +176,9 @@ public class RegistrationController {
registrationRequest.deviceActivationRequest().gcmToken().ifPresent(gcmRegistrationId ->
device.setGcmId(gcmRegistrationId.gcmRegistrationId()));
keysManager.storeEcSignedPreKeys(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciSignedPreKey().get()));
keysManager.storePqLastResort(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciPqLastResortPreKey().get()));
keysManager.storeEcSignedPreKeys(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniSignedPreKey().get()));
keysManager.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniPqLastResortPreKey().get()));
});
}

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

@ -312,7 +312,10 @@ public class AccountsManager {
numberChangedAccount = updateWithRetries(
account,
a -> { setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds); return true; },
a -> {
setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds);
return true;
},
a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier),
() -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
@ -322,6 +325,8 @@ public class AccountsManager {
keysManager.delete(phoneNumberIdentifier);
keysManager.delete(originalPhoneNumberIdentifier);
keysManager.storeEcSignedPreKeys(phoneNumberIdentifier, pniSignedPreKeys);
if (pniPqLastResortPreKeys != null) {
keysManager.storePqLastResort(
phoneNumberIdentifier,
@ -362,6 +367,7 @@ public class AccountsManager {
final List<Long> pqEnabledDeviceIDs = keysManager.getPqEnabledDevices(pni);
keysManager.delete(pni);
keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys);
if (pniPqLastResortPreKeys != null) {
keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
}

21
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java
View File

@ -14,6 +14,7 @@ import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import javax.annotation.Nullable;
import org.whispersystems.textsecuregcm.entities.ECPreKey;
import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
@ -21,26 +22,30 @@ public class KeysManager {
private final SingleUseECPreKeyStore ecPreKeys;
private final SingleUseKEMPreKeyStore pqPreKeys;
private final RepeatedUseECSignedPreKeyStore ecSignedPreKeys;
private final RepeatedUseKEMSignedPreKeyStore pqLastResortKeys;
public KeysManager(
final DynamoDbAsyncClient dynamoDbAsyncClient,
final String ecTableName,
final String pqTableName,
final String ecSignedPreKeysTableName,
final String pqLastResortTableName) {
this.ecPreKeys = new SingleUseECPreKeyStore(dynamoDbAsyncClient, ecTableName);
this.pqPreKeys = new SingleUseKEMPreKeyStore(dynamoDbAsyncClient, pqTableName);
this.ecSignedPreKeys = new RepeatedUseECSignedPreKeyStore(dynamoDbAsyncClient, ecSignedPreKeysTableName);
this.pqLastResortKeys = new RepeatedUseKEMSignedPreKeyStore(dynamoDbAsyncClient, pqLastResortTableName);
}
public void store(final UUID identifier, final long deviceId, final List<ECPreKey> keys) {
store(identifier, deviceId, keys, null, null);
store(identifier, deviceId, keys, null, null, null);
}
public void store(
final UUID identifier, final long deviceId,
@Nullable final List<ECPreKey> ecKeys,
@Nullable final List<KEMSignedPreKey> pqKeys,
@Nullable final ECSignedPreKey ecSignedPreKey,
@Nullable final KEMSignedPreKey pqLastResortKey) {
final List<CompletableFuture<Void>> storeFutures = new ArrayList<>();
@ -53,6 +58,10 @@ public class KeysManager {
storeFutures.add(pqPreKeys.store(identifier, deviceId, pqKeys));
}
if (ecSignedPreKey != null) {
storeFutures.add(ecSignedPreKeys.store(identifier, deviceId, ecSignedPreKey));
}
if (pqLastResortKey != null) {
storeFutures.add(pqLastResortKeys.store(identifier, deviceId, pqLastResortKey));
}
@ -60,6 +69,10 @@ public class KeysManager {
CompletableFuture.allOf(storeFutures.toArray(new CompletableFuture[0])).join();
}
public void storeEcSignedPreKeys(final UUID identifier, final Map<Long, ECSignedPreKey> keys) {
ecSignedPreKeys.store(identifier, keys).join();
}
public void storePqLastResort(final UUID identifier, final Map<Long, KEMSignedPreKey> keys) {
pqLastResortKeys.store(identifier, keys).join();
}
@ -80,6 +93,10 @@ public class KeysManager {
return pqLastResortKeys.find(identifier, deviceId).join();
}
public CompletableFuture<Optional<ECSignedPreKey>> getEcSignedPreKey(final UUID identifier, final long deviceId) {
return ecSignedPreKeys.find(identifier, deviceId);
}
public List<Long> getPqEnabledDevices(final UUID identifier) {
return pqLastResortKeys.getDeviceIdsWithKeys(identifier).collectList().block();
}
@ -96,6 +113,7 @@ public class KeysManager {
CompletableFuture.allOf(
ecPreKeys.delete(accountUuid),
pqPreKeys.delete(accountUuid),
ecSignedPreKeys.delete(accountUuid),
pqLastResortKeys.delete(accountUuid))
.join();
}
@ -104,6 +122,7 @@ public class KeysManager {
CompletableFuture.allOf(
ecPreKeys.delete(accountUuid, deviceId),
pqPreKeys.delete(accountUuid, deviceId),
ecSignedPreKeys.delete(accountUuid, deviceId),
pqLastResortKeys.delete(accountUuid, deviceId))
.join();
}

1
service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java
View File

@ -165,6 +165,7 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
KeysManager keys = new KeysManager(
dynamoDbAsyncClient,
configuration.getDynamoDbTables().getEcKeys().getTableName(),
configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
configuration.getDynamoDbTables().getKemKeys().getTableName(),
configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

1
service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
View File

@ -149,6 +149,7 @@ record CommandDependencies(
KeysManager keys = new KeysManager(
dynamoDbAsyncClient,
configuration.getDynamoDbTables().getEcKeys().getTableName(),
configuration.getDynamoDbTables().getEcSignedPreKeys().getTableName(),
configuration.getDynamoDbTables().getKemKeys().getTableName(),
configuration.getDynamoDbTables().getKemLastResortKeys().getTableName());
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,

2
service/src/test/java/org/whispersystems/textsecuregcm/controllers/RegistrationControllerTest.java
View File

@ -671,6 +671,8 @@ class RegistrationControllerTest {
verify(device).setSignedPreKey(expectedAciSignedPreKey);
verify(device).setPhoneNumberIdentitySignedPreKey(expectedPniSignedPreKey);
verify(keysManager).storeEcSignedPreKeys(accountIdentifier, Map.of(Device.MASTER_ID, expectedAciSignedPreKey));
verify(keysManager).storeEcSignedPreKeys(phoneNumberIdentifier, Map.of(Device.MASTER_ID, expectedPniSignedPreKey));
verify(keysManager).storePqLastResort(accountIdentifier, Map.of(Device.MASTER_ID, expectedAciPqLastResortPreKey));
verify(keysManager).storePqLastResort(phoneNumberIdentifier, Map.of(Device.MASTER_ID, expectedPniPqLastResortPreKey));

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

@ -709,6 +709,7 @@ class AccountsManagerTest {
verify(keysManager).delete(originalPni);
verify(keysManager, atLeastOnce()).delete(targetPni);
verify(keysManager).delete(newPni);
verify(keysManager).storeEcSignedPreKeys(eq(newPni), any());
verifyNoMoreInteractions(keysManager);
}
@ -748,6 +749,7 @@ class AccountsManagerTest {
verify(keysManager).delete(newPni);
verify(keysManager).delete(originalPni);
verify(keysManager).getPqEnabledDevices(uuid);
verify(keysManager).storeEcSignedPreKeys(newPni, newSignedKeys);
verify(keysManager).storePqLastResort(eq(newPni), eq(Map.of(1L, newSignedPqKeys.get(1L))));
verifyNoMoreInteractions(keysManager);
}
@ -852,6 +854,7 @@ class AccountsManagerTest {
verifyNoInteractions(deletedAccounts);
verify(keysManager).delete(oldPni);
verify(keysManager).storeEcSignedPreKeys(oldPni, newSignedKeys);
// only the pq key for the already-pq-enabled device should be saved
verify(keysManager).storePqLastResort(eq(oldPni), eq(Map.of(1L, newSignedPqKeys.get(1L))));

16
service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java
View File

@ -115,7 +115,21 @@ public final class DynamoDbExtensionSchema {
.build()),
List.of(), List.of()),
REPEATED_USE_SIGNED_PRE_KEYS("repeated_use_signed_pre_keys_test",
REPEATED_USE_EC_SIGNED_PRE_KEYS("repeated_use_signed_ec_pre_keys_test",
RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID,
RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID,
List.of(
AttributeDefinition.builder()
.attributeName(RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build(),
AttributeDefinition.builder()
.attributeName(RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID)
.attributeType(ScalarAttributeType.N)
.build()),
List.of(), List.of()),
REPEATED_USE_KEM_SIGNED_PRE_KEYS("repeated_use_signed_kem_pre_keys_test",
RepeatedUseSignedPreKeyStore.KEY_ACCOUNT_UUID,
RepeatedUseSignedPreKeyStore.KEY_DEVICE_ID,
List.of(

57
service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java
View File

@ -32,7 +32,7 @@ class KeysManagerTest {
@RegisterExtension
static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
Tables.EC_KEYS, Tables.PQ_KEYS, Tables.REPEATED_USE_SIGNED_PRE_KEYS);
Tables.EC_KEYS, Tables.PQ_KEYS, Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS, Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
private static final UUID ACCOUNT_UUID = UUID.randomUUID();
private static final long DEVICE_ID = 1L;
@ -45,7 +45,8 @@ class KeysManagerTest {
DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
Tables.EC_KEYS.tableName(),
Tables.PQ_KEYS.tableName(),
Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName(),
Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName());
}
@Test
@ -64,25 +65,25 @@ class KeysManagerTest {
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Repeatedly storing same key should have no effect");
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(generateTestKEMSignedPreKey(1)), null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(generateTestKEMSignedPreKey(1)), null, null);
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Uploading new PQ prekeys should have no effect on EC prekeys");
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, null, generateTestKEMSignedPreKey(1001));
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, null, null, generateTestKEMSignedPreKey(1001));
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Uploading new PQ last-resort prekey should have no effect on EC prekeys");
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
"Uploading new PQ last-resort prekey should have no effect on one-time PQ prekeys");
assertEquals(1001, keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).get().keyId());
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(2)), null, null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(2)), null, null, null);
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Inserting a new key should overwrite all prior keys of the same type for the given account/device");
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
"Uploading new EC prekeys should have no effect on PQ prekeys");
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(3)), List.of(generateTestKEMSignedPreKey(2)), null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(3)), List.of(generateTestKEMSignedPreKey(2)), null, null);
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Inserting a new key should overwrite all prior keys of the same type for the given account/device");
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
@ -90,7 +91,7 @@ class KeysManagerTest {
keysManager.store(ACCOUNT_UUID, DEVICE_ID,
List.of(generateTestPreKey(4), generateTestPreKey(5)),
List.of(generateTestKEMSignedPreKey(6), generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(1002));
List.of(generateTestKEMSignedPreKey(6), generateTestKEMSignedPreKey(7)), null, generateTestKEMSignedPreKey(1002));
assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID),
"Inserting multiple new keys should overwrite all prior keys for the given account/device");
assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID),
@ -119,7 +120,7 @@ class KeysManagerTest {
final KEMSignedPreKey preKey2 = generateTestKEMSignedPreKey(2);
final KEMSignedPreKey preKeyLast = generateTestKEMSignedPreKey(1001);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(preKey1, preKey2), preKeyLast);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(preKey1, preKey2), null, preKeyLast);
assertEquals(Optional.of(preKey1), keysManager.takePQ(ACCOUNT_UUID, DEVICE_ID));
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
@ -139,7 +140,7 @@ class KeysManagerTest {
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(1)), List.of(generateTestKEMSignedPreKey(1)), null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, List.of(generateTestPreKey(1)), List.of(generateTestKEMSignedPreKey(1)), null, null);
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
}
@ -148,53 +149,69 @@ class KeysManagerTest {
void testDeleteByAccount() {
keysManager.store(ACCOUNT_UUID, DEVICE_ID,
List.of(generateTestPreKey(1), generateTestPreKey(2)),
List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)), generateTestKEMSignedPreKey(5));
List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)),
generateTestECSignedPreKey(5),
generateTestKEMSignedPreKey(6));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1,
List.of(generateTestPreKey(6)),
List.of(generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(8));
List.of(generateTestPreKey(7)),
List.of(generateTestKEMSignedPreKey(8)),
generateTestECSignedPreKey(9),
generateTestKEMSignedPreKey(10));
assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
keysManager.delete(ACCOUNT_UUID);
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
}
@Test
void testDeleteByAccountAndDevice() {
keysManager.store(ACCOUNT_UUID, DEVICE_ID,
List.of(generateTestPreKey(1), generateTestPreKey(2)),
List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)), generateTestKEMSignedPreKey(5));
List.of(generateTestPreKey(1),generateTestPreKey(2)),
List.of(generateTestKEMSignedPreKey(3), generateTestKEMSignedPreKey(4)),
generateTestECSignedPreKey(5),
generateTestKEMSignedPreKey(6));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1,
List.of(generateTestPreKey(6)),
List.of(generateTestKEMSignedPreKey(7)), generateTestKEMSignedPreKey(8));
List.of(generateTestPreKey(7)),
List.of(generateTestKEMSignedPreKey(8)),
generateTestECSignedPreKey(9),
generateTestKEMSignedPreKey(10));
assertEquals(2, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(2, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
keysManager.delete(ACCOUNT_UUID, DEVICE_ID);
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID));
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID));
assertFalse(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID).join().isPresent());
assertFalse(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID).isPresent());
assertEquals(1, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertEquals(1, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID + 1));
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).isPresent());
}
@ -225,10 +242,10 @@ class KeysManagerTest {
void testGetPqEnabledDevices() {
final ECKeyPair identityKeyPair = Curve.generateKeyPair();
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(KeysHelper.signedKEMPreKey(1, identityKeyPair)), null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1, null, null, KeysHelper.signedKEMPreKey(2, identityKeyPair));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 2, null, List.of(KeysHelper.signedKEMPreKey(3, identityKeyPair)), KeysHelper.signedKEMPreKey(4, identityKeyPair));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 3, null, null, null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID, null, List.of(KeysHelper.signedKEMPreKey(1, identityKeyPair)), null, null);
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 1, null, null, null, KeysHelper.signedKEMPreKey(2, identityKeyPair));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 2, null, List.of(KeysHelper.signedKEMPreKey(3, identityKeyPair)), null, KeysHelper.signedKEMPreKey(4, identityKeyPair));
keysManager.store(ACCOUNT_UUID, DEVICE_ID + 3, null, null, null, null);
assertIterableEquals(
Set.of(DEVICE_ID + 1, DEVICE_ID + 2),
Set.copyOf(keysManager.getPqEnabledDevices(ACCOUNT_UUID)));

4
service/src/test/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStoreTest.java
View File

@ -20,14 +20,14 @@ class RepeatedUseECSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTes
@RegisterExtension
static final DynamoDbExtension DYNAMO_DB_EXTENSION =
new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS);
new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS);
private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
@BeforeEach
void setUp() {
keyStore = new RepeatedUseECSignedPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
DynamoDbExtensionSchema.Tables.REPEATED_USE_EC_SIGNED_PRE_KEYS.tableName());
}
@Override

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

@ -12,8 +12,6 @@ import org.signal.libsignal.protocol.ecc.ECKeyPair;
import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
import static org.junit.jupiter.api.Assertions.*;
class RepeatedUseKEMSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTest<KEMSignedPreKey> {
private RepeatedUseKEMSignedPreKeyStore keyStore;
@ -22,14 +20,14 @@ class RepeatedUseKEMSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTe
@RegisterExtension
static final DynamoDbExtension DYNAMO_DB_EXTENSION =
new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS);
new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS);
private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
@BeforeEach
void setUp() {
keyStore = new RepeatedUseKEMSignedPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName());
}
@Override

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

@ -315,6 +315,8 @@ class DeviceControllerTest {
verify(pendingDevicesManager).remove(AuthHelper.VALID_NUMBER);
verify(messagesManager).clear(eq(AuthHelper.VALID_UUID), eq(42L));
verify(clientPresenceManager).disconnectPresence(AuthHelper.VALID_UUID, Device.MASTER_ID);
verify(keysManager).storeEcSignedPreKeys(AuthHelper.VALID_UUID, Map.of(response.getDeviceId(), aciSignedPreKey.get()));
verify(keysManager).storeEcSignedPreKeys(AuthHelper.VALID_PNI, Map.of(response.getDeviceId(), pniSignedPreKey.get()));
verify(keysManager).storePqLastResort(AuthHelper.VALID_UUID, Map.of(response.getDeviceId(), aciPqLastResortPreKey.get()));
verify(keysManager).storePqLastResort(AuthHelper.VALID_PNI, Map.of(response.getDeviceId(), pniPqLastResortPreKey.get()));
}

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

@ -687,7 +687,7 @@ class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), listCaptor.capture(), isNull(), isNull());
verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
assertThat(listCaptor.getValue()).containsExactly(preKey);
@ -718,7 +718,7 @@ class KeysControllerTest {
ArgumentCaptor<List<ECPreKey>> ecCaptor = ArgumentCaptor.forClass(List.class);
ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(pqLastResortPreKey));
verify(KEYS).store(eq(AuthHelper.VALID_UUID), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(signedPreKey), eq(pqLastResortPreKey));
assertThat(ecCaptor.getValue()).containsExactly(preKey);
assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);
@ -820,7 +820,7 @@ class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), listCaptor.capture(), isNull(), isNull());
verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
assertThat(listCaptor.getValue()).containsExactly(preKey);
@ -852,7 +852,7 @@ class KeysControllerTest {
ArgumentCaptor<List<ECPreKey>> ecCaptor = ArgumentCaptor.forClass(List.class);
ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(pqLastResortPreKey));
verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), ecCaptor.capture(), pqCaptor.capture(), eq(signedPreKey), eq(pqLastResortPreKey));
assertThat(ecCaptor.getValue()).containsExactly(preKey);
assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);
@ -896,7 +896,7 @@ class KeysControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
ArgumentCaptor<List<ECPreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
verify(KEYS).store(eq(AuthHelper.DISABLED_UUID), eq(1L), listCaptor.capture(), isNull(), isNull());
verify(KEYS).store(eq(AuthHelper.DISABLED_UUID), eq(1L), listCaptor.capture(), isNull(), eq(signedPreKey), isNull());
List<ECPreKey> capturedList = listCaptor.getValue();
assertThat(capturedList.size()).isEqualTo(1);