0
0
mirror of https://github.com/signalapp/Signal-Server.git synced 2024-09-20 03:52:16 +02:00

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,

View File

@ -827,6 +827,6 @@ class DeviceControllerTest {
assertEquals(204, response.getStatus());
}
verify(clientPublicKeysManager).setPublicKey(AuthHelper.VALID_UUID, AuthHelper.VALID_DEVICE.getId(), request.publicKey());
verify(clientPublicKeysManager).setPublicKey(AuthHelper.VALID_ACCOUNT, AuthHelper.VALID_DEVICE.getId(), request.publicKey());
}
}

View File

@ -98,8 +98,6 @@ public class AccountCreationDeletionIntegrationTest {
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(),
@ -115,6 +113,8 @@ public class AccountCreationDeletionIntegrationTest {
final AccountLockManager accountLockManager = new AccountLockManager(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS_LOCK.tableName());
clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
final SecureStorageClient secureStorageClient = mock(SecureStorageClient.class);
when(secureStorageClient.deleteStoredData(any())).thenReturn(CompletableFuture.completedFuture(null));
@ -459,7 +459,7 @@ public class AccountCreationDeletionIntegrationTest {
aciPqLastResortPreKey,
pniPqLastResortPreKey));
clientPublicKeysManager.setPublicKey(account.getIdentifier(IdentityType.ACI), Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
clientPublicKeysManager.setPublicKey(account, Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
final UUID aci = account.getIdentifier(IdentityType.ACI);

View File

@ -91,8 +91,6 @@ class AccountsManagerChangeNumberIntegrationTest {
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(),
@ -108,6 +106,9 @@ class AccountsManagerChangeNumberIntegrationTest {
final AccountLockManager accountLockManager = new AccountLockManager(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
Tables.DELETED_ACCOUNTS_LOCK.tableName());
final ClientPublicKeysManager clientPublicKeysManager =
new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
final SecureStorageClient secureStorageClient = mock(SecureStorageClient.class);
when(secureStorageClient.deleteStoredData(any())).thenReturn(CompletableFuture.completedFuture(null));

View File

@ -86,8 +86,6 @@ public class AddRemoveDeviceIntegrationTest {
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(),
@ -103,6 +101,8 @@ public class AddRemoveDeviceIntegrationTest {
final AccountLockManager accountLockManager = new AccountLockManager(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
DynamoDbExtensionSchema.Tables.DELETED_ACCOUNTS_LOCK.tableName());
clientPublicKeysManager = new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
final SecureStorageClient secureStorageClient = mock(SecureStorageClient.class);
when(secureStorageClient.deleteStoredData(any())).thenReturn(CompletableFuture.completedFuture(null));
@ -229,8 +229,8 @@ 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();
clientPublicKeysManager.setPublicKey(account, Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
clientPublicKeysManager.setPublicKey(account, addedDeviceId, Curve.generateKeyPair().getPublicKey()).join();
final Account updatedAccount = accountsManager.removeDevice(updatedAccountAndDevice.first(), addedDeviceId).join();
@ -290,8 +290,8 @@ 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();
clientPublicKeysManager.setPublicKey(retrievedAccount, Device.PRIMARY_ID, Curve.generateKeyPair().getPublicKey()).join();
clientPublicKeysManager.setPublicKey(retrievedAccount, addedDeviceId, Curve.generateKeyPair().getPublicKey()).join();
assertEquals(2, retrievedAccount.getDevices().size());