Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-19 21:58:08 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add PNI support to AccountLockManager

Browse Source
This commit is contained in:
Chris Eager
2024-11-22 10:11:27 -06:00
committed by Chris Eager
parent 73812b06be
commit eb55b80bdc
7 changed files with 92 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java
View File

@@ -8,11 +8,13 @@ import com.amazonaws.services.dynamodbv2.ReleaseLockOptions;
import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import java.util.stream.Stream;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
public class AccountLockManager {
@@ -43,28 +45,35 @@ public class AccountLockManager {
* account lifecycle changes (like deleting an account or changing a phone number). The given task runs once locks for
* all given phone numbers have been acquired, and the locks are released as soon as the task completes by any means.
*
* @param e164s the phone numbers for which to acquire a distributed, pessimistic lock
* @param task the task to execute once locks have been acquired
* @param e164s the phone numbers for which to acquire a distributed, pessimistic lock
* @param phoneNumberIdentifiers the phone number identifiers for which to acquire a distributed, pessimistic lock
* @param task the task to execute once locks have been acquired
* @param lockAcquisitionExecutor the executor on which to run blocking lock acquire/release tasks. this executor
* should not use virtual threads.
*
* @throws InterruptedException if interrupted while acquiring a lock
*/
public void withLock(final List<String> e164s, final Runnable task, final Executor lockAcquisitionExecutor) {
public void withLock(final List<String> e164s, final List<UUID> phoneNumberIdentifiers, final Runnable task,
final Executor lockAcquisitionExecutor) {
if (e164s.isEmpty()) {
throw new IllegalArgumentException("List of e164s to lock must not be empty");
}
if (phoneNumberIdentifiers.isEmpty()) {
throw new IllegalArgumentException("List of PNIs to lock must not be empty");
}
final List<LockItem> lockItems = new ArrayList<>(e164s.size());
final List<String> allIdentifiers = Stream.concat(e164s.stream(),
phoneNumberIdentifiers.stream().map(UUID::toString))
.toList();
final List<LockItem> lockItems = new ArrayList<>(allIdentifiers.size());
try {
// Offload the acquire/release tasks to the dedicated lock acquisition executor. The lock client performs blocking
// operations while holding locks which forces thread pinning when this method runs on a virtual thread.
// https://github.com/awslabs/amazon-dynamodb-lock-client/issues/97
CompletableFuture.runAsync(() -> {
for (final String e164 : e164s) {
for (final String identifier : allIdentifiers) {
try {
lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(e164)
lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(identifier)
.withAcquireReleasedLocksConsistently(true)
.build()));
} catch (final InterruptedException e) {
@@ -91,27 +100,32 @@ public class AccountLockManager {
* account lifecycle changes (like deleting an account or changing a phone number). The given task runs once locks for
* all given phone numbers have been acquired, and the locks are released as soon as the task completes by any means.
*
* @param e164s the phone numbers for which to acquire a distributed, pessimistic lock
* @param taskSupplier a supplier for the task to execute once locks have been acquired
* @param executor the executor on which to acquire and release locks
*
* @param e164s the phone numbers for which to acquire a distributed, pessimistic lock
* @param phoneNumberIdentifiers the phone number identifiers for which to acquire a distributed, pessimistic lock
* @param taskSupplier a supplier for the task to execute once locks have been acquired
* @param executor the executor on which to acquire and release locks
* @return a future that completes normally when the given task has executed successfully and all locks have been
* released; the returned future may fail with an {@link InterruptedException} if interrupted while acquiring a lock
*/
public <T> CompletableFuture<T> withLockAsync(final List<String> e164s,
final Supplier<CompletableFuture<T>> taskSupplier,
final Executor executor) {
public <T> CompletableFuture<T> withLockAsync(final List<String> e164s, final List<UUID> phoneNumberIdentifiers,
final Supplier<CompletableFuture<T>> taskSupplier, final Executor executor) {
if (e164s.isEmpty()) {
throw new IllegalArgumentException("List of e164s to lock must not be empty");
}
if (phoneNumberIdentifiers.isEmpty()) {
throw new IllegalArgumentException("List of PNIs to lock must not be empty");
}
final List<LockItem> lockItems = new ArrayList<>(e164s.size());
final List<String> allIdentifiers = Stream.concat(e164s.stream(),
phoneNumberIdentifiers.stream().map(UUID::toString))
.toList();
final List<LockItem> lockItems = new ArrayList<>(allIdentifiers.size());
return CompletableFuture.runAsync(() -> {
for (final String e164 : e164s) {
for (final String identifier : allIdentifiers) {
try {
lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(e164)
lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(identifier)
.withAcquireReleasedLocksConsistently(true)
.build()));
} catch (final InterruptedException e) {

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

@@ -271,15 +271,16 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
@Nullable final String userAgent) throws InterruptedException {
final Account account = new Account();
final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
return createTimer.record(() -> {
accountLockManager.withLock(List.of(number), () -> {
accountLockManager.withLock(List.of(number), List.of(phoneNumberIdentifier), () -> {
final Optional<UUID> maybeRecentlyDeletedAccountIdentifier =
accounts.findRecentlyDeletedAccountIdentifier(number);
// Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is
// re-registering.
account.setNumber(number, phoneNumberIdentifiers.getPhoneNumberIdentifier(number));
account.setNumber(number, phoneNumberIdentifier);
account.setUuid(maybeRecentlyDeletedAccountIdentifier.orElseGet(UUID::randomUUID));
account.setIdentityKey(aciIdentityKey);
account.setPhoneNumberIdentityKey(pniIdentityKey);
@@ -363,9 +364,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
// We exclude the primary device's repeated-use keys from deletion because new keys were provided as
// part of the account creation process, and we don't want to delete the keys that just got added.
return CompletableFuture.allOf(keysManager.deleteSingleUsePreKeys(aci),
keysManager.deleteSingleUsePreKeys(pni),
messagesManager.clear(aci),
profilesManager.deleteAll(aci));
keysManager.deleteSingleUsePreKeys(pni),
messagesManager.clear(aci),
profilesManager.deleteAll(aci));
})
.join();
}
@@ -375,7 +376,8 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
Tags tags = Tags.of(UserAgentTagUtil.getPlatformTag(userAgent),
Tag.of("type", accountCreationType),
Tag.of("hasPushToken", String.valueOf(
primaryDeviceSpec.apnRegistrationId().isPresent() || primaryDeviceSpec.gcmRegistrationId().isPresent())),
primaryDeviceSpec.apnRegistrationId().isPresent() || primaryDeviceSpec.gcmRegistrationId()
.isPresent())),
Tag.of("pushTokenType", pushTokenType));
if (StringUtils.isNotBlank(previousPushTokenType)) {
@@ -385,7 +387,8 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
Metrics.counter(CREATE_COUNTER_NAME, tags).increment();
accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword));
registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(),
registrationRecoveryPassword));
}, accountLockExecutor);
return account;
@@ -394,6 +397,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
public CompletableFuture<Pair<Account, Device>> addDevice(final Account account, final DeviceSpec deviceSpec, final String linkDeviceToken) {
return accountLockManager.withLockAsync(List.of(account.getNumber()),
List.of(account.getPhoneNumberIdentifier()),
() -> addDevice(account.getIdentifier(IdentityType.ACI), deviceSpec, linkDeviceToken, MAX_UPDATE_ATTEMPTS),
accountLockExecutor);
}
@@ -581,7 +585,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
throw new IllegalArgumentException("Cannot remove primary device");
}
return accountLockManager.withLockAsync(List.of(account.getNumber()),
return accountLockManager.withLockAsync(List.of(account.getNumber()), List.of(account.getPhoneNumberIdentifier()),
() -> removeDevice(account.getIdentifier(IdentityType.ACI), deviceId, MAX_UPDATE_ATTEMPTS),
accountLockExecutor);
}
@@ -647,8 +651,10 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
final AtomicReference<Account> updatedAccount = new AtomicReference<>();
final UUID targetPhoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber);
accountLockManager.withLock(List.of(account.getNumber(), targetNumber), () -> {
accountLockManager.withLock(List.of(account.getNumber(), targetNumber),
List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier), () -> {
redisDelete(account);
// There are three possible states for accounts associated with the target phone number:
@@ -674,15 +680,14 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
}
final UUID uuid = account.getUuid();
final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber);
CompletableFuture.allOf(
keysManager.deleteSingleUsePreKeys(phoneNumberIdentifier),
keysManager.deleteSingleUsePreKeys(targetPhoneNumberIdentifier),
keysManager.deleteSingleUsePreKeys(originalPhoneNumberIdentifier))
.join();
final Collection<TransactWriteItem> keyWriteItems =
buildPniKeyWriteItems(uuid, phoneNumberIdentifier, pniSignedPreKeys, pniPqLastResortPreKeys);
buildPniKeyWriteItems(uuid, targetPhoneNumberIdentifier, pniSignedPreKeys, pniPqLastResortPreKeys);
final Account numberChangedAccount = updateWithRetries(
account,
@@ -690,7 +695,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
setPniKeys(account, pniIdentityKey, pniRegistrationIds);
return true;
},
a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier, maybeDisplacedUuid, keyWriteItems),
a -> accounts.changeNumber(a, targetNumber, targetPhoneNumberIdentifier, maybeDisplacedUuid, keyWriteItems),
() -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
@@ -1220,7 +1225,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
public CompletableFuture<Void> delete(final Account account, final DeletionReason deletionReason) {
final Timer.Sample sample = Timer.start();
return accountLockManager.withLockAsync(List.of(account.getNumber()), () -> delete(account), accountLockExecutor)
return accountLockManager.withLockAsync(List.of(account.getNumber()), List.of(account.getPhoneNumberIdentifier()),
() -> delete(account),
accountLockExecutor)
.whenComplete((ignored, throwable) -> {
sample.stop(deleteTimer);

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

@@ -42,6 +42,7 @@ public class ClientPublicKeysManager {
*/
public CompletableFuture<Void> setPublicKey(final Account account, final byte deviceId, final ECPublicKey publicKey) {
return accountLockManager.withLockAsync(List.of(account.getNumber()),
List.of(account.getPhoneNumberIdentifier()),
() -> clientPublicKeys.setPublicKey(account.getIdentifier(IdentityType.ACI), deviceId, publicKey),
accountLockExecutor);
}