Create separate key stores for different kinds of pre-keys

Author: jon-signal

LICENSE

@@ -296,7 +296,7 @@ commercial, industrial or non-consumer uses, unless such uses represent
 the only significant mode of use of the product.
 
   "Installation Information" for a User Product means any methods,
-procedures, authorization keys, or other information required to install
+procedures, authorization keysManager, or other information required to install
 and execute modified versions of a covered work in that User Product from
 a modified version of its Corresponding Source.  The information must
 suffice to ensure that the continued functioning of the modified object

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -176,7 +176,7 @@
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
-import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
 import org.whispersystems.textsecuregcm.storage.MessagePersister;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
 import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
@@ -345,10 +345,11 @@ public void run(WhisperServerConfiguration config, Environment environment) thro
         config.getDynamoDbTables().getPhoneNumberIdentifiers().getTableName());
     Profiles profiles = new Profiles(dynamoDbClient, dynamoDbAsyncClient,
         config.getDynamoDbTables().getProfiles().getTableName());
-    Keys keys = new Keys(dynamoDbClient,
+    KeysManager keys = new KeysManager(
+            dynamoDbAsyncClient,
         config.getDynamoDbTables().getEcKeys().getTableName(),
-        config.getDynamoDbTables().getPqKeys().getTableName(),
-        config.getDynamoDbTables().getPqLastResortKeys().getTableName());
+        config.getDynamoDbTables().getKemKeys().getTableName(),
+        config.getDynamoDbTables().getKemLastResortKeys().getTableName());
     MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         config.getDynamoDbTables().getMessages().getTableName(),
         config.getDynamoDbTables().getMessages().getExpiration(),

service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbTables.java

@@ -51,8 +51,8 @@ public Duration getExpiration() {
   private final Table deletedAccountsLock;
   private final IssuedReceiptsTableConfiguration issuedReceipts;
   private final Table ecKeys;
-  private final Table pqKeys;
-  private final Table pqLastResortKeys;
+  private final Table kemKeys;
+  private final Table kemLastResortKeys;
   private final TableWithExpiration messages;
   private final Table pendingAccounts;
   private final Table pendingDevices;
@@ -72,8 +72,8 @@ public DynamoDbTables(
       @JsonProperty("deletedAccountsLock") final Table deletedAccountsLock,
       @JsonProperty("issuedReceipts") final IssuedReceiptsTableConfiguration issuedReceipts,
       @JsonProperty("ecKeys") final Table ecKeys,
-      @JsonProperty("pqKeys") final Table pqKeys,
-      @JsonProperty("pqLastResortKeys") final Table pqLastResortKeys,
+      @JsonProperty("pqKeys") final Table kemKeys,
+      @JsonProperty("pqLastResortKeys") final Table kemLastResortKeys,
       @JsonProperty("messages") final TableWithExpiration messages,
       @JsonProperty("pendingAccounts") final Table pendingAccounts,
       @JsonProperty("pendingDevices") final Table pendingDevices,
@@ -92,8 +92,8 @@ public DynamoDbTables(
     this.deletedAccountsLock = deletedAccountsLock;
     this.issuedReceipts = issuedReceipts;
     this.ecKeys = ecKeys;
-    this.pqKeys = pqKeys;
-    this.pqLastResortKeys = pqLastResortKeys;
+    this.kemKeys = kemKeys;
+    this.kemLastResortKeys = kemLastResortKeys;
     this.messages = messages;
     this.pendingAccounts = pendingAccounts;
     this.pendingDevices = pendingDevices;
@@ -140,14 +140,14 @@ public Table getEcKeys() {
 
   @NotNull
   @Valid
-  public Table getPqKeys() {
-    return pqKeys;
+  public Table getKemKeys() {
+    return kemKeys;
   }
 
   @NotNull
   @Valid
-  public Table getPqLastResortKeys() {
-    return pqLastResortKeys;
+  public Table getKemLastResortKeys() {
+    return kemLastResortKeys;
   }
 
   @NotNull

service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java

@@ -51,7 +51,7 @@
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
-import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
 import org.whispersystems.textsecuregcm.util.Pair;
@@ -67,14 +67,14 @@ public class DeviceController {
   private final StoredVerificationCodeManager pendingDevices;
   private final AccountsManager       accounts;
   private final MessagesManager       messages;
-  private final Keys keys;
+  private final KeysManager keys;
   private final RateLimiters          rateLimiters;
   private final Map<String, Integer>  maxDeviceConfiguration;
 
   public DeviceController(StoredVerificationCodeManager pendingDevices,
       AccountsManager accounts,
       MessagesManager messages,
-      Keys keys,
+      KeysManager keys,
       RateLimiters rateLimiters,
       Map<String, Integer> maxDeviceConfiguration) {
     this.pendingDevices = pendingDevices;

service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

@@ -53,15 +53,15 @@
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
-import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 @Path("/v2/keys")
 @Tag(name = "Keys")
 public class KeysController {
 
   private final RateLimiters                rateLimiters;
-  private final Keys                        keys;
+  private final KeysManager keys;
   private final AccountsManager             accounts;
 
   private static final String IDENTITY_KEY_CHANGE_COUNTER_NAME = name(KeysController.class, "identityKeyChange");
@@ -70,7 +70,7 @@ public class KeysController {
   private static final String IDENTITY_TYPE_TAG_NAME = "identityType";
   private static final String HAS_IDENTITY_KEY_TAG_NAME = "hasIdentityKey";
 
-  public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts) {
+  public KeysController(RateLimiters rateLimiters, KeysManager keys, AccountsManager accounts) {
     this.rateLimiters = rateLimiters;
     this.keys = keys;
     this.accounts = accounts;

service/src/main/java/org/whispersystems/textsecuregcm/controllers/RegistrationController.java

@@ -48,7 +48,7 @@
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
-import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
 import org.whispersystems.textsecuregcm.util.HeaderUtils;
 import org.whispersystems.textsecuregcm.util.Util;
 
@@ -74,18 +74,18 @@ public class RegistrationController {
   private final AccountsManager accounts;
   private final PhoneVerificationTokenManager phoneVerificationTokenManager;
   private final RegistrationLockVerificationManager registrationLockVerificationManager;
-  private final Keys keys;
+  private final KeysManager keysManager;
   private final RateLimiters rateLimiters;
 
   public RegistrationController(final AccountsManager accounts,
                                 final PhoneVerificationTokenManager phoneVerificationTokenManager,
                                 final RegistrationLockVerificationManager registrationLockVerificationManager,
-                                final Keys keys,
+                                final KeysManager keysManager,
                                 final RateLimiters rateLimiters) {
     this.accounts = accounts;
     this.phoneVerificationTokenManager = phoneVerificationTokenManager;
     this.registrationLockVerificationManager = registrationLockVerificationManager;
-    this.keys = keys;
+    this.keysManager = keysManager;
     this.rateLimiters = rateLimiters;
   }
 
@@ -176,8 +176,8 @@ public AccountIdentityResponse register(
         registrationRequest.deviceActivationRequest().gcmToken().ifPresent(gcmRegistrationId ->
             device.setGcmId(gcmRegistrationId.gcmRegistrationId()));
 
-        keys.storePqLastResort(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciPqLastResortPreKey().get()));
-        keys.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniPqLastResortPreKey().get()));
+        keysManager.storePqLastResort(a.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().aciPqLastResortPreKey().get()));
+        keysManager.storePqLastResort(a.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.deviceActivationRequest().pniPqLastResortPreKey().get()));
       });
     }
 

service/src/main/java/org/whispersystems/textsecuregcm/entities/ChangeNumberRequest.java

@@ -43,7 +43,7 @@ public record ChangeNumberRequest(
     @NotEmpty byte[] pniIdentityKey,
 
     @Schema(description="""
-        A list of synchronization messages to send to companion devices to supply the private keys
+        A list of synchronization messages to send to companion devices to supply the private keysManager
         associated with the new identity key and their new prekeys.
         Exactly one message must be supplied for each enabled device other than the sending (primary) device.""")
     @NotNull @Valid List<@NotNull @Valid IncomingMessage> deviceMessages,

service/src/main/java/org/whispersystems/textsecuregcm/entities/ChangePhoneNumberRequest.java

@@ -36,7 +36,7 @@ public record ChangePhoneNumberRequest(
     @Nullable byte[] pniIdentityKey,
 
     @Schema(description="""
-        A list of synchronization messages to send to companion devices to supply the private keys
+        A list of synchronization messages to send to companion devices to supply the private keysManager
         associated with the new identity key and their new prekeys.
         Exactly one message must be supplied for each enabled device other than the sending (primary) device.""")
     @Nullable List<IncomingMessage> deviceMessages,

service/src/main/java/org/whispersystems/textsecuregcm/entities/PhoneNumberIdentityKeyDistributionRequest.java

@@ -30,7 +30,7 @@ public record PhoneNumberIdentityKeyDistributionRequest(
     @NotNull
     @Valid
     @Schema(description="""
-        A list of synchronization messages to send to companion devices to supply the private keys
+        A list of synchronization messages to send to companion devices to supply the private keysManager
         associated with the new identity key and their new prekeys.
         Exactly one message must be supplied for each enabled device other than the sending (primary) device.""")
     List<@NotNull @Valid IncomingMessage> deviceMessages,

service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java

@@ -90,7 +90,7 @@ public class AccountsManager {
   private final FaultTolerantRedisCluster cacheCluster;
   private final AccountLockManager accountLockManager;
   private final DeletedAccounts deletedAccounts;
-  private final Keys keys;
+  private final KeysManager keysManager;
   private final MessagesManager messagesManager;
   private final ProfilesManager profilesManager;
   private final StoredVerificationCodeManager pendingAccounts;
@@ -134,7 +134,7 @@ public AccountsManager(final Accounts accounts,
       final FaultTolerantRedisCluster cacheCluster,
       final AccountLockManager accountLockManager,
       final DeletedAccounts deletedAccounts,
-      final Keys keys,
+      final KeysManager keysManager,
       final MessagesManager messagesManager,
       final ProfilesManager profilesManager,
       final StoredVerificationCodeManager pendingAccounts,
@@ -150,7 +150,7 @@ public AccountsManager(final Accounts accounts,
     this.cacheCluster = cacheCluster;
     this.accountLockManager = accountLockManager;
     this.deletedAccounts = deletedAccounts;
-    this.keys = keys;
+    this.keysManager = keysManager;
     this.messagesManager = messagesManager;
     this.profilesManager = profilesManager;
     this.pendingAccounts = pendingAccounts;
@@ -223,8 +223,8 @@ public Account create(final String number,
         // 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);
-          keys.delete(actualUuid);
-          keys.delete(account.getPhoneNumberIdentifier());
+          keysManager.delete(actualUuid);
+          keysManager.delete(account.getPhoneNumberIdentifier());
           profilesManager.deleteAll(actualUuid);
           clientPresenceManager.disconnectAllPresencesForUuid(actualUuid);
         }
@@ -315,13 +315,13 @@ public Account changeNumber(final Account account,
 
       updatedAccount.set(numberChangedAccount);
 
-      keys.delete(phoneNumberIdentifier);
-      keys.delete(originalPhoneNumberIdentifier);
+      keysManager.delete(phoneNumberIdentifier);
+      keysManager.delete(originalPhoneNumberIdentifier);
 
       if (pniPqLastResortPreKeys != null) {
-        keys.storePqLastResort(
+        keysManager.storePqLastResort(
             phoneNumberIdentifier,
-            keys.getPqEnabledDevices(uuid).stream().collect(
+            keysManager.getPqEnabledDevices(uuid).stream().collect(
                 Collectors.toMap(
                     Function.identity(),
                     pniPqLastResortPreKeys::get)));
@@ -356,10 +356,10 @@ public Account updatePniKeys(final Account account,
     final UUID pni = account.getPhoneNumberIdentifier();
     final Account updatedAccount = update(account, a -> { return setPniKeys(a, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds); });
 
-    final List<Long> pqEnabledDeviceIDs = keys.getPqEnabledDevices(pni);
-    keys.delete(pni);
+    final List<Long> pqEnabledDeviceIDs = keysManager.getPqEnabledDevices(pni);
+    keysManager.delete(pni);
     if (pniPqLastResortPreKeys != null) {
-      keys.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
+      keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
     }
 
     return updatedAccount;
@@ -740,8 +740,8 @@ private void delete(final Account account) {
         account.getUuid());
 
     profilesManager.deleteAll(account.getUuid());
-    keys.delete(account.getUuid());
-    keys.delete(account.getPhoneNumberIdentifier());
+    keysManager.delete(account.getUuid());
+    keysManager.delete(account.getPhoneNumberIdentifier());
     messagesManager.clear(account.getUuid());
     messagesManager.clear(account.getPhoneNumberIdentifier());
     registrationRecoveryPasswordsManager.removeForNumber(account.getNumber());