Perform device list validations in the scope of a pessimistic account lock

This commit is contained in:
Jon Chambers 2025-05-12 14:06:01 -04:00 committed by Jon Chambers
parent a4b98f38a6
commit 93ba6616d1
1 changed files with 31 additions and 18 deletions

View File

@ -673,8 +673,6 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
return account; return account;
} }
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
try { try {
return accountLockManager.withLock(List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier), return accountLockManager.withLock(List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier),
() -> changeNumber(account, targetNumber, targetPhoneNumberIdentifier, pniIdentityKey, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds), accountLockExecutor); () -> changeNumber(account, targetNumber, targetPhoneNumberIdentifier, pniIdentityKey, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds), accountLockExecutor);
@ -696,7 +694,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
final IdentityKey pniIdentityKey, final IdentityKey pniIdentityKey,
final Map<Byte, ECSignedPreKey> pniSignedPreKeys, final Map<Byte, ECSignedPreKey> pniSignedPreKeys,
final Map<Byte, KEMSignedPreKey> pniPqLastResortPreKeys, final Map<Byte, KEMSignedPreKey> pniPqLastResortPreKeys,
final Map<Byte, Integer> pniRegistrationIds) { final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException {
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier(); final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier();
@ -751,24 +751,37 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
final Map<Byte, KEMSignedPreKey> pniPqLastResortPreKeys, final Map<Byte, KEMSignedPreKey> pniPqLastResortPreKeys,
final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException { final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException {
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds); try {
return accountLockManager.withLock(List.of(account.getIdentifier(IdentityType.PNI)), () -> {
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
final UUID aci = account.getIdentifier(IdentityType.ACI); final UUID aci = account.getIdentifier(IdentityType.ACI);
final UUID pni = account.getIdentifier(IdentityType.PNI); final UUID pni = account.getIdentifier(IdentityType.PNI);
final Collection<TransactWriteItem> keyWriteItems = final Collection<TransactWriteItem> keyWriteItems =
buildPniKeyWriteItems(pni, pniSignedPreKeys, pniPqLastResortPreKeys); buildPniKeyWriteItems(pni, pniSignedPreKeys, pniPqLastResortPreKeys);
return redisDeleteAsync(account) return redisDeleteAsync(account)
.thenCompose(ignored -> keysManager.deleteSingleUsePreKeys(pni)) .thenCompose(ignored -> keysManager.deleteSingleUsePreKeys(pni))
.thenCompose(ignored -> updateTransactionallyWithRetriesAsync(account, .thenCompose(ignored -> updateTransactionallyWithRetriesAsync(account,
a -> setPniKeys(a, pniIdentityKey, pniRegistrationIds), a -> setPniKeys(a, pniIdentityKey, pniRegistrationIds),
accounts::updateTransactionallyAsync, accounts::updateTransactionallyAsync,
() -> accounts.getByAccountIdentifierAsync(aci).thenApply(Optional::orElseThrow), () -> accounts.getByAccountIdentifierAsync(aci).thenApply(Optional::orElseThrow),
a -> keyWriteItems, a -> keyWriteItems,
AccountChangeValidator.GENERAL_CHANGE_VALIDATOR, AccountChangeValidator.GENERAL_CHANGE_VALIDATOR,
MAX_UPDATE_ATTEMPTS)) MAX_UPDATE_ATTEMPTS))
.join(); .join();
}, accountLockExecutor);
} catch (final Exception e) {
if (e instanceof MismatchedDevicesException mismatchedDevicesException) {
throw mismatchedDevicesException;
} else if (e instanceof RuntimeException runtimeException) {
throw runtimeException;
}
logger.error("Unexpected exception when updating PNI key material", e);
throw new RuntimeException(e);
}
} }
private Collection<TransactWriteItem> buildPniKeyWriteItems( private Collection<TransactWriteItem> buildPniKeyWriteItems(