Make key deletion operations asynchronous

This commit is contained in:
Jon Chambers 2023-06-26 11:48:21 -04:00 committed by Jon Chambers
parent f709b00be3
commit a0d6146ff5
8 changed files with 49 additions and 23 deletions

View File

@ -111,11 +111,14 @@ public class DeviceController {
throw new WebApplicationException(Response.Status.UNAUTHORIZED);
}
final CompletableFuture<Void> deleteKeysFuture = keys.delete(account.getUuid(), deviceId);
messages.clear(account.getUuid(), deviceId);
account = accounts.update(account, a -> a.removeDevice(deviceId));
keys.delete(account.getUuid(), deviceId);
// ensure any messages that came in after the first clear() are also removed
messages.clear(account.getUuid(), deviceId);
deleteKeysFuture.join();
}
@Timed
@ -336,10 +339,13 @@ public class DeviceController {
final Account updatedAccount = accounts.update(account, a -> {
device.setId(a.getNextDeviceId());
final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
keys.delete(a.getUuid(), device.getId()),
keys.delete(a.getPhoneNumberIdentifier(), device.getId()));
messages.clear(a.getUuid(), device.getId());
keys.delete(a.getUuid(), device.getId());
keys.delete(a.getPhoneNumberIdentifier(), device.getId());
deleteKeysFuture.join();
maybeDeviceActivationRequest.ifPresent(deviceActivationRequest -> CompletableFuture.allOf(
keys.storeEcSignedPreKeys(a.getUuid(),

View File

@ -228,10 +228,15 @@ public class AccountsManager {
// confident that everything has already been deleted. In the second case, though, we're taking over an existing
// account and need to clear out messages and keys that may have been stored for the old account.
if (!originalUuid.equals(actualUuid)) {
final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
keysManager.delete(actualUuid),
keysManager.delete(account.getPhoneNumberIdentifier()));
messagesManager.clear(actualUuid);
keysManager.delete(actualUuid);
keysManager.delete(account.getPhoneNumberIdentifier());
profilesManager.deleteAll(actualUuid);
deleteKeysFuture.join();
clientPresenceManager.disconnectAllPresencesForUuid(actualUuid);
}
@ -324,8 +329,10 @@ public class AccountsManager {
updatedAccount.set(numberChangedAccount);
keysManager.delete(phoneNumberIdentifier);
keysManager.delete(originalPhoneNumberIdentifier);
CompletableFuture.allOf(
keysManager.delete(phoneNumberIdentifier),
keysManager.delete(originalPhoneNumberIdentifier))
.join();
keysManager.storeEcSignedPreKeys(phoneNumberIdentifier, pniSignedPreKeys);
@ -369,9 +376,9 @@ public class AccountsManager {
final List<Long> pqEnabledDeviceIDs = keysManager.getPqEnabledDevices(pni).join();
keysManager.delete(pni);
keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys);
keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys).join();
if (pniPqLastResortPreKeys != null) {
keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get))).join();
}
return updatedAccount;
@ -755,13 +762,16 @@ public class AccountsManager {
final CompletableFuture<Void> deleteSecureValueRecoveryServiceDataFuture = secureValueRecovery2Client.deleteBackups(
account.getUuid());
final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
keysManager.delete(account.getUuid()),
keysManager.delete(account.getPhoneNumberIdentifier()));
profilesManager.deleteAll(account.getUuid());
keysManager.delete(account.getUuid());
keysManager.delete(account.getPhoneNumberIdentifier());
messagesManager.clear(account.getUuid());
messagesManager.clear(account.getPhoneNumberIdentifier());
registrationRecoveryPasswordsManager.removeForNumber(account.getNumber());
deleteKeysFuture.join();
deleteStorageServiceDataFuture.join();
deleteBackupServiceDataFuture.join();
deleteSecureValueRecoveryServiceDataFuture.join();

View File

@ -122,25 +122,23 @@ public class KeysManager {
return pqPreKeys.getCount(identifier, deviceId);
}
public void delete(final UUID accountUuid) {
CompletableFuture.allOf(
public CompletableFuture<Void> delete(final UUID accountUuid) {
return CompletableFuture.allOf(
ecPreKeys.delete(accountUuid),
pqPreKeys.delete(accountUuid),
dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys()
? ecSignedPreKeys.delete(accountUuid)
: CompletableFuture.completedFuture(null),
pqLastResortKeys.delete(accountUuid))
.join();
pqLastResortKeys.delete(accountUuid));
}
public void delete(final UUID accountUuid, final long deviceId) {
CompletableFuture.allOf(
public CompletableFuture<Void> delete(final UUID accountUuid, final long deviceId) {
return CompletableFuture.allOf(
ecPreKeys.delete(accountUuid, deviceId),
pqPreKeys.delete(accountUuid, deviceId),
dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys()
? ecSignedPreKeys.delete(accountUuid, deviceId)
: CompletableFuture.completedFuture(null),
pqLastResortKeys.delete(accountUuid, deviceId))
.join();
pqLastResortKeys.delete(accountUuid, deviceId));
}
}

View File

@ -65,7 +65,7 @@ public class UnlinkDeviceCommand extends EnvironmentCommand<WhisperServerConfigu
account = deps.accountsManager().update(account, a -> a.removeDevice(deviceId));
System.out.format("Removing keys for device %s::%d\n", aci, deviceId);
deps.keysManager().delete(account.getUuid(), deviceId);
deps.keysManager().delete(account.getUuid(), deviceId).join();
System.out.format("Clearing additional messages for %s::%d\n", aci, deviceId);
deps.messagesManager().clear(account.getUuid(), deviceId);

View File

@ -100,13 +100,16 @@ class AccountsManagerChangeNumberIntegrationTest {
final PhoneNumberIdentifiers phoneNumberIdentifiers =
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), Tables.PNI.tableName());
final KeysManager keysManager = mock(KeysManager.class);
when(keysManager.delete(any())).thenReturn(CompletableFuture.completedFuture(null));
accountsManager = new AccountsManager(
accounts,
phoneNumberIdentifiers,
CACHE_CLUSTER_EXTENSION.getRedisCluster(),
accountLockManager,
deletedAccounts,
mock(KeysManager.class),
keysManager,
mock(MessagesManager.class),
mock(ProfilesManager.class),
mock(StoredVerificationCodeManager.class),

View File

@ -157,6 +157,8 @@ class AccountsManagerTest {
return null;
}).when(accountLockManager).withLock(any(), any());
when(keysManager.delete(any())).thenReturn(CompletableFuture.completedFuture(null));
accountsManager = new AccountsManager(
accounts,
phoneNumberIdentifiers,
@ -785,6 +787,7 @@ class AccountsManagerTest {
final IdentityKey pniIdentityKey = new IdentityKey(Curve.generateKeyPair().getPublicKey());
when(keysManager.getPqEnabledDevices(any())).thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
final Account updatedAccount = accountsManager.updatePniKeys(account, pniIdentityKey, newSignedKeys, null, newRegistrationIds);
@ -829,6 +832,8 @@ class AccountsManagerTest {
UUID oldPni = account.getPhoneNumberIdentifier();
when(keysManager.getPqEnabledDevices(oldPni)).thenReturn(CompletableFuture.completedFuture(List.of(1L)));
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
Map<Long, ECSignedPreKey> oldSignedPreKeys = account.getDevices().stream().collect(Collectors.toMap(Device::getId, Device::getSignedPreKey));

View File

@ -185,7 +185,7 @@ class KeysManagerTest {
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
keysManager.delete(ACCOUNT_UUID);
keysManager.delete(ACCOUNT_UUID).join();
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID).join());
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());
@ -222,7 +222,7 @@ class KeysManagerTest {
assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
keysManager.delete(ACCOUNT_UUID, DEVICE_ID);
keysManager.delete(ACCOUNT_UUID, DEVICE_ID).join();
assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID).join());
assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());

View File

@ -8,6 +8,7 @@ import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
@ -155,6 +156,9 @@ class DeviceControllerTest {
when(accountsManager.getByE164(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(maxedAccount));
AccountsHelper.setupMockUpdate(accountsManager);
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
when(keysManager.delete(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
}
@AfterEach