Remove client public keys when deleting accounts/devices

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

@@ -587,7 +587,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         accountLockManager, keysManager, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client,
         clientPresenceManager,
-        registrationRecoveryPasswordsManager, accountLockExecutor, clientPresenceExecutor,
+        registrationRecoveryPasswordsManager, clientPublicKeysManager, accountLockExecutor, clientPresenceExecutor,
         clock);
     RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
     APNSender apnSender = new APNSender(apnSenderExecutor, config.getApnConfiguration());

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

@@ -23,7 +23,6 @@ import java.time.Duration;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -101,6 +100,7 @@ public class AccountsManager {
   private final SecureValueRecovery2Client secureValueRecovery2Client;
   private final ClientPresenceManager clientPresenceManager;
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager;
+  private final ClientPublicKeysManager clientPublicKeysManager;
   private final Executor accountLockExecutor;
   private final Executor clientPresenceExecutor;
   private final Clock clock;
@@ -141,6 +141,7 @@ public class AccountsManager {
       final SecureValueRecovery2Client secureValueRecovery2Client,
       final ClientPresenceManager clientPresenceManager,
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
+      final ClientPublicKeysManager clientPublicKeysManager,
       final Executor accountLockExecutor,
       final Executor clientPresenceExecutor,
       final Clock clock) {
@@ -155,6 +156,7 @@ public class AccountsManager {
     this.secureValueRecovery2Client = secureValueRecovery2Client;
     this.clientPresenceManager = clientPresenceManager;
     this.registrationRecoveryPasswordsManager = requireNonNull(registrationRecoveryPasswordsManager);
+    this.clientPublicKeysManager = clientPublicKeysManager;
     this.accountLockExecutor = accountLockExecutor;
     this.clientPresenceExecutor = clientPresenceExecutor;
     this.clock = requireNonNull(clock);
@@ -319,10 +321,15 @@ public class AccountsManager {
         .thenCompose(account -> {
           account.removeDevice(deviceId);
 
-          return accounts.updateTransactionallyAsync(account, keysManager.buildWriteItemsForRemovedDevice(
+          final List<TransactWriteItem> additionalWriteItems = new ArrayList<>(
-              account.getIdentifier(IdentityType.ACI),
+              keysManager.buildWriteItemsForRemovedDevice(
-              account.getIdentifier(IdentityType.PNI),
+                  account.getIdentifier(IdentityType.ACI),
-              deviceId))
+                  account.getIdentifier(IdentityType.PNI),
+                  deviceId));
+
+          additionalWriteItems.add(clientPublicKeysManager.buildTransactWriteItemForDeletion(account.getIdentifier(IdentityType.ACI), deviceId));
+
+          return accounts.updateTransactionallyAsync(account, additionalWriteItems)
               .thenApply(ignored -> account);
         })
         .thenCompose(updatedAccount -> redisDeleteAsync(updatedAccount).thenApply(ignored -> updatedAccount))
@@ -956,12 +963,18 @@ public class AccountsManager {
   }
 
   private CompletableFuture<Void> delete(final Account account) {
-    final List<TransactWriteItem> additionalWriteItems =
+    final List<TransactWriteItem> additionalWriteItems = Stream.concat(
-        account.getDevices().stream().flatMap(device -> keysManager.buildWriteItemsForRemovedDevice(
+            account.getDevices().stream()
-                account.getIdentifier(IdentityType.ACI),
+                .flatMap(device -> keysManager.buildWriteItemsForRemovedDevice(
-                account.getIdentifier(IdentityType.PNI),
+                        account.getIdentifier(IdentityType.ACI),
-                device.getId()).stream())
+                        account.getIdentifier(IdentityType.PNI),
-            .toList();
+                        device.getId())
+                    .stream()),
+            account.getDevices().stream()
+                .map(device -> clientPublicKeysManager.buildTransactWriteItemForDeletion(
+                    account.getIdentifier(IdentityType.ACI),
+                    device.getId())))
+        .toList();
 
     return CompletableFuture.allOf(
             secureStorageClient.deleteStoredData(account.getUuid()),

service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java

@@ -37,6 +37,8 @@ import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2
 import org.whispersystems.textsecuregcm.storage.AccountLockManager;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.ClientPublicKeys;
+import org.whispersystems.textsecuregcm.storage.ClientPublicKeysManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.KeysManager;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
@@ -145,6 +147,11 @@ record CommandDependencies(
     RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = new RegistrationRecoveryPasswordsManager(
         registrationRecoveryPasswords);
 
+    ClientPublicKeys clientPublicKeys =
+        new ClientPublicKeys(dynamoDbAsyncClient, configuration.getDynamoDbTables().getClientPublicKeys().getTableName());
+
+    ClientPublicKeysManager clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
+
     Accounts accounts = new Accounts(
         dynamoDbClient,
         dynamoDbAsyncClient,
@@ -197,7 +204,7 @@ record CommandDependencies(
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         accountLockManager, keys, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client, clientPresenceManager,
-        registrationRecoveryPasswordsManager, accountLockExecutor, clientPresenceExecutor,
+        registrationRecoveryPasswordsManager, clientPublicKeysManager, accountLockExecutor, clientPresenceExecutor,
         clock);
     RateLimiters rateLimiters = RateLimiters.createAndValidate(configuration.getLimitsConfiguration(),
         dynamicConfigurationManager, rateLimitersCluster);

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java

@@ -53,6 +53,7 @@ public class AccountCreationDeletionIntegrationTest {
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
       DynamoDbExtensionSchema.Tables.ACCOUNTS,
+      DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS,
       DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS,
       DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS_LOCK,
       DynamoDbExtensionSchema.Tables.NUMBERS,
@@ -74,6 +75,7 @@ public class AccountCreationDeletionIntegrationTest {
 
   private AccountsManager accountsManager;
   private KeysManager keysManager;
+  private ClientPublicKeysManager clientPublicKeysManager;
 
   record DeliveryChannels(boolean fetchesMessages, String apnsToken, String apnsVoipToken, String fcmToken) {}
 
@@ -93,6 +95,11 @@ public class AccountCreationDeletionIntegrationTest {
         DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
     );
 
+    final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
+        DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
+
+    clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
+
     final Accounts accounts = new Accounts(
         DYNAMO_DB_EXTENSION.getDynamoDbClient(),
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
@@ -142,6 +149,7 @@ public class AccountCreationDeletionIntegrationTest {
         svr2Client,
         mock(ClientPresenceManager.class),
         registrationRecoveryPasswordsManager,
+        clientPublicKeysManager,
         accountLockExecutor,
         clientPresenceExecutor,
         CLOCK);
@@ -451,6 +459,8 @@ public class AccountCreationDeletionIntegrationTest {
             aciPqLastResortPreKey,
             pniPqLastResortPreKey));
 
+    clientPublicKeysManager.setPublicKey(account.getIdentifier(IdentityType.ACI), Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
+
     final UUID aci = account.getIdentifier(IdentityType.ACI);
 
     assertTrue(accountsManager.getByAccountIdentifier(aci).isPresent());
@@ -462,6 +472,7 @@ public class AccountCreationDeletionIntegrationTest {
     assertFalse(keysManager.getEcSignedPreKey(account.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
     assertFalse(keysManager.getLastResort(account.getUuid(), Device.PRIMARY_ID).join().isPresent());
     assertFalse(keysManager.getLastResort(account.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
+    assertFalse(clientPublicKeysManager.findPublicKey(account.getUuid(), Device.PRIMARY_ID).join().isPresent());
   }
 
   @SuppressWarnings("OptionalUsedAsFieldOrParameterType")

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java

@@ -48,6 +48,7 @@ class AccountsManagerChangeNumberIntegrationTest {
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
       Tables.ACCOUNTS,
+      Tables.CLIENT_PUBLIC_KEYS,
       Tables.DELETED_ACCOUNTS,
       Tables.DELETED_ACCOUNTS_LOCK,
       Tables.NUMBERS,
@@ -87,6 +88,11 @@ class AccountsManagerChangeNumberIntegrationTest {
           Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
       );
 
+      final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
+          DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
+
+      final ClientPublicKeysManager clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
+
       final Accounts accounts = new Accounts(
           DYNAMO_DB_EXTENSION.getDynamoDbClient(),
           DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
@@ -137,6 +143,7 @@ class AccountsManagerChangeNumberIntegrationTest {
           svr2Client,
           clientPresenceManager,
           registrationRecoveryPasswordsManager,
+          clientPublicKeysManager,
           accountLockExecutor,
           clientPresenceExecutor,
           mock(Clock.class));

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java

@@ -131,6 +131,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
           mock(SecureValueRecovery2Client.class),
           mock(ClientPresenceManager.class),
           mock(RegistrationRecoveryPasswordsManager.class),
+          mock(ClientPublicKeysManager.class),
           mock(Executor.class),
           mock(Executor.class),
           mock(Clock.class)

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java

@@ -105,6 +105,7 @@ class AccountsManagerTest {
   private MessagesManager messagesManager;
   private ProfilesManager profilesManager;
   private ClientPresenceManager clientPresenceManager;
+  private ClientPublicKeysManager clientPublicKeysManager;
 
   private Map<String, UUID> phoneNumberIdentifiersByE164;
 
@@ -137,6 +138,7 @@ class AccountsManagerTest {
     messagesManager = mock(MessagesManager.class);
     profilesManager = mock(ProfilesManager.class);
     clientPresenceManager = mock(ClientPresenceManager.class);
+    clientPublicKeysManager = mock(ClientPublicKeysManager.class);
 
     final Executor clientPresenceExecutor = mock(Executor.class);
 
@@ -231,6 +233,7 @@ class AccountsManagerTest {
         svr2Client,
         clientPresenceManager,
         registrationRecoveryPasswordsManager,
+        clientPublicKeysManager,
         mock(Executor.class),
         clientPresenceExecutor,
         clock);
@@ -740,6 +743,7 @@ class AccountsManagerTest {
     verify(messagesManager, times(2)).clear(account.getUuid(), linkedDevice.getId());
     verify(keysManager, times(2)).deleteSingleUsePreKeys(account.getUuid(), linkedDevice.getId());
     verify(keysManager).buildWriteItemsForRemovedDevice(account.getUuid(), account.getPhoneNumberIdentifier(), linkedDevice.getId());
+    verify(clientPublicKeysManager).buildTransactWriteItemForDeletion(account.getUuid(), linkedDevice.getId());
     verify(clientPresenceManager).disconnectPresence(account.getUuid(), linkedDevice.getId());
   }
 

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java

@@ -143,6 +143,7 @@ class AccountsManagerUsernameIntegrationTest {
         mock(SecureValueRecovery2Client.class),
         mock(ClientPresenceManager.class),
         mock(RegistrationRecoveryPasswordsManager.class),
+        mock(ClientPublicKeysManager.class),
         Executors.newSingleThreadExecutor(),
         Executors.newSingleThreadExecutor(),
         mock(Clock.class));

service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java

@@ -42,6 +42,7 @@ public class AddRemoveDeviceIntegrationTest {
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
       DynamoDbExtensionSchema.Tables.ACCOUNTS,
+      DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS,
       DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS,
       DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS_LOCK,
       DynamoDbExtensionSchema.Tables.NUMBERS,
@@ -62,6 +63,7 @@ public class AddRemoveDeviceIntegrationTest {
   private ExecutorService clientPresenceExecutor;
 
   private KeysManager keysManager;
+  private ClientPublicKeysManager clientPublicKeysManager;
   private MessagesManager messagesManager;
   private AccountsManager accountsManager;
 
@@ -81,6 +83,11 @@ public class AddRemoveDeviceIntegrationTest {
         DynamoDbExtensionSchema.Tables.REPEATED_USE_KEM_SIGNED_PRE_KEYS.tableName()
     );
 
+    final ClientPublicKeys clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
+        DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName());
+
+    clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
+
     final Accounts accounts = new Accounts(
         DYNAMO_DB_EXTENSION.getDynamoDbClient(),
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
@@ -130,6 +137,7 @@ public class AddRemoveDeviceIntegrationTest {
         svr2Client,
         mock(ClientPresenceManager.class),
         registrationRecoveryPasswordsManager,
+        clientPublicKeysManager,
         accountLockExecutor,
         clientPresenceExecutor,
         CLOCK);
@@ -221,6 +229,9 @@ public class AddRemoveDeviceIntegrationTest {
 
     final byte addedDeviceId = updatedAccountAndDevice.second().getId();
 
+    clientPublicKeysManager.setPublicKey(account.getUuid(), Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
+    clientPublicKeysManager.setPublicKey(account.getUuid(), addedDeviceId, Curve.generateKeyPair().getPublicKey()).join();
+
     final Account updatedAccount = accountsManager.removeDevice(updatedAccountAndDevice.first(), addedDeviceId).join();
 
     assertEquals(1, updatedAccount.getDevices().size());
@@ -229,11 +240,13 @@ public class AddRemoveDeviceIntegrationTest {
     assertFalse(keysManager.getEcSignedPreKey(updatedAccount.getPhoneNumberIdentifier(), addedDeviceId).join().isPresent());
     assertFalse(keysManager.getLastResort(updatedAccount.getUuid(), addedDeviceId).join().isPresent());
     assertFalse(keysManager.getLastResort(updatedAccount.getPhoneNumberIdentifier(), addedDeviceId).join().isPresent());
+    assertFalse(clientPublicKeysManager.findPublicKey(updatedAccount.getUuid(), addedDeviceId).join().isPresent());
 
     assertTrue(keysManager.getEcSignedPreKey(updatedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getEcSignedPreKey(updatedAccount.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(updatedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(updatedAccount.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
+    assertTrue(clientPublicKeysManager.findPublicKey(updatedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
   }
 
   @Test
@@ -249,7 +262,6 @@ public class AddRemoveDeviceIntegrationTest {
     assertEquals(1, accountsManager.getByAccountIdentifier(account.getUuid()).orElseThrow().getDevices().size());
 
     final UUID aci = account.getIdentifier(IdentityType.ACI);
-    final UUID pni = account.getIdentifier(IdentityType.PNI);
 
     final Pair<Account, Device> updatedAccountAndDevice =
         accountsManager.addDevice(account, new DeviceSpec(
@@ -278,16 +290,21 @@ public class AddRemoveDeviceIntegrationTest {
 
     final Account retrievedAccount = accountsManager.getByAccountIdentifierAsync(aci).join().orElseThrow();
 
+    clientPublicKeysManager.setPublicKey(retrievedAccount.getUuid(), Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
+    clientPublicKeysManager.setPublicKey(retrievedAccount.getUuid(), addedDeviceId, Curve.generateKeyPair().getPublicKey()).join();
+
     assertEquals(2, retrievedAccount.getDevices().size());
 
     assertTrue(keysManager.getEcSignedPreKey(retrievedAccount.getUuid(), addedDeviceId).join().isPresent());
     assertTrue(keysManager.getEcSignedPreKey(retrievedAccount.getPhoneNumberIdentifier(), addedDeviceId).join().isPresent());
     assertTrue(keysManager.getLastResort(retrievedAccount.getUuid(), addedDeviceId).join().isPresent());
     assertTrue(keysManager.getLastResort(retrievedAccount.getPhoneNumberIdentifier(), addedDeviceId).join().isPresent());
+    assertTrue(clientPublicKeysManager.findPublicKey(retrievedAccount.getUuid(), addedDeviceId).join().isPresent());
 
     assertTrue(keysManager.getEcSignedPreKey(retrievedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getEcSignedPreKey(retrievedAccount.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(retrievedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
     assertTrue(keysManager.getLastResort(retrievedAccount.getPhoneNumberIdentifier(), Device.PRIMARY_ID).join().isPresent());
+    assertTrue(clientPublicKeysManager.findPublicKey(retrievedAccount.getUuid(), Device.PRIMARY_ID).join().isPresent());
   }
 }