Set client public keys in the scope of a pessimistic account lock

This commit is contained in:
Jon Chambers
2024-05-20 10:48:16 -04:00
committed by Jon Chambers
parent 0e43524dac
commit 7980da9ce5
8 changed files with 35 additions and 20 deletions

View File

@@ -582,7 +582,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
messageDeletionAsyncExecutor);
AccountLockManager accountLockManager = new AccountLockManager(dynamoDbClient,
config.getDynamoDbTables().getDeletedAccountsLock().getTableName());
ClientPublicKeysManager clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
ClientPublicKeysManager clientPublicKeysManager =
new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
accountLockManager, keysManager, messagesManager, profilesManager,
secureStorageClient, secureValueRecovery2Client,

View File

@@ -301,7 +301,7 @@ public class DeviceController {
public CompletableFuture<Void> setPublicKey(@Auth final AuthenticatedAccount auth,
final SetPublicKeyRequest setPublicKeyRequest) {
return clientPublicKeysManager.setPublicKey(auth.getAccount().getIdentifier(IdentityType.ACI),
return clientPublicKeysManager.setPublicKey(auth.getAccount(),
auth.getAuthenticatedDevice().getId(),
setPublicKeyRequest.publicKey());
}

View File

@@ -1,9 +1,12 @@
package org.whispersystems.textsecuregcm.storage;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import org.signal.libsignal.protocol.ecc.ECPublicKey;
import org.whispersystems.textsecuregcm.identity.IdentityType;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
/**
@@ -14,8 +17,16 @@ public class ClientPublicKeysManager {
private final ClientPublicKeys clientPublicKeys;
public ClientPublicKeysManager(final ClientPublicKeys clientPublicKeys) {
private final AccountLockManager accountLockManager;
private final Executor accountLockExecutor;
public ClientPublicKeysManager(final ClientPublicKeys clientPublicKeys,
final AccountLockManager accountLockManager,
final Executor accountLockExecutor) {
this.clientPublicKeys = clientPublicKeys;
this.accountLockManager = accountLockManager;
this.accountLockExecutor = accountLockExecutor;
}
/**
@@ -23,14 +34,16 @@ public class ClientPublicKeysManager {
* is intended for use for adding public keys to existing accounts/devices as a migration step. Callers should use
* {@link #buildTransactWriteItemForInsertion(UUID, byte, ECPublicKey)} instead when creating new accounts/devices.
*
* @param accountIdentifier the identifier for the target account
* @param account the target account
* @param deviceId the identifier for the target device
* @param publicKey the public key to store for the target account/device
* @return a future that completes when the given key has been stored
*/
public CompletableFuture<Void> setPublicKey(final UUID accountIdentifier, final byte deviceId, final ECPublicKey publicKey) {
return clientPublicKeys.setPublicKey(accountIdentifier, deviceId, publicKey);
public CompletableFuture<Void> setPublicKey(final Account account, final byte deviceId, final ECPublicKey publicKey) {
return accountLockManager.withLockAsync(List.of(account.getNumber()),
() -> clientPublicKeys.setPublicKey(account.getIdentifier(IdentityType.ACI), deviceId, publicKey),
accountLockExecutor);
}
/**

View File

@@ -150,8 +150,6 @@ record CommandDependencies(
ClientPublicKeys clientPublicKeys =
new ClientPublicKeys(dynamoDbAsyncClient, configuration.getDynamoDbTables().getClientPublicKeys().getTableName());
ClientPublicKeysManager clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys);
Accounts accounts = new Accounts(
dynamoDbClient,
dynamoDbAsyncClient,
@@ -201,6 +199,8 @@ record CommandDependencies(
reportMessageManager, messageDeletionExecutor);
AccountLockManager accountLockManager = new AccountLockManager(dynamoDbClient,
configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
ClientPublicKeysManager clientPublicKeysManager =
new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
accountLockManager, keys, messagesManager, profilesManager,
secureStorageClient, secureValueRecovery2Client, clientPresenceManager,