From 6c131936230b2b90937770b323c75ba8b23efdb7 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Wed, 17 Jan 2024 11:43:01 -0500 Subject: [PATCH] Add a data store for client public keys for transport-level authentication/encryption --- .../storage/ClientPublicKeys.java | 121 ++++++++++++++++++ .../storage/ClientPublicKeysManager.java | 63 +++++++++ .../storage/ClientPublicKeysTest.java | 48 +++++++ .../storage/DynamoDbExtensionSchema.java | 15 +++ 4 files changed, 247 insertions(+) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeys.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysManager.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeys.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeys.java new file mode 100644 index 000000000..c7546ad95 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeys.java @@ -0,0 +1,121 @@ +package org.whispersystems.textsecuregcm.storage; + +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import org.signal.libsignal.protocol.InvalidKeyException; +import org.signal.libsignal.protocol.ecc.ECPublicKey; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.util.AttributeValues; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.Delete; +import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.Put; +import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; + +/** + * Stores clients' public keys for transport-level authentication/encryption in a DynamoDB table. + */ +public class ClientPublicKeys { + + private final DynamoDbAsyncClient dynamoDbAsyncClient; + private final String tableName; + + static final String KEY_ACCOUNT_UUID = "U"; + static final String KEY_DEVICE_ID = "D"; + static final String ATTR_PUBLIC_KEY = "K"; + + private static final Logger log = LoggerFactory.getLogger(ClientPublicKeys.class); + + public ClientPublicKeys(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) { + this.dynamoDbAsyncClient = dynamoDbAsyncClient; + this.tableName = tableName; + } + + /** + * Builds a {@link TransactWriteItem} that will store a public key for the given account/device. Intended for use when + * adding devices to accounts or creating new accounts. + * + * @param accountIdentifier the identifier for 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 {@code TransactWriteItem} that will store the given public key for the given account/device + */ + TransactWriteItem buildTransactWriteItemForInsertion(final UUID accountIdentifier, + final byte deviceId, + final ECPublicKey publicKey) { + + return TransactWriteItem.builder() + .put(Put.builder() + .tableName(tableName) + .item(Map.of( + KEY_ACCOUNT_UUID, getPartitionKey(accountIdentifier), + KEY_DEVICE_ID, getSortKey(deviceId), + ATTR_PUBLIC_KEY, AttributeValues.fromByteArray(publicKey.serialize()))) + .build()) + .build(); + } + + /** + * Builds a {@link TransactWriteItem} that will remove the public key for the given account/device. Intended for + * use when removing devices from accounts or deleting/re-creating accounts. + * + * @param accountIdentifier the identifier for the target account + * @param deviceId the identifier for the target device + * + * @return a {@code TransactWriteItem} that will remove the public key for the given account/device + */ + TransactWriteItem buildTransactWriteItemForDeletion(final UUID accountIdentifier, final byte deviceId) { + return TransactWriteItem.builder() + .delete(Delete.builder() + .tableName(tableName) + .key(getPrimaryKey(accountIdentifier, deviceId)) + .build()) + .build(); + } + + /** + * Finds the public key for the given account/device. + * + * @param accountIdentifier the identifier for the target account + * @param deviceId the identifier for the target device + * + * @return a future that yields the Ed25519 public key for the given account/device, or empty if no public key was + * found + */ + CompletableFuture> findPublicKey(final UUID accountIdentifier, final byte deviceId) { + return dynamoDbAsyncClient.getItem(GetItemRequest.builder() + .tableName(tableName) + .consistentRead(true) + .key(getPrimaryKey(accountIdentifier, deviceId)) + .build()) + .thenApply(response -> Optional.of(response.item()) + .filter(item -> response.hasItem()) + .map(item -> { + try { + return new ECPublicKey(item.get(ATTR_PUBLIC_KEY).b().asByteArray()); + } catch (final InvalidKeyException e) { + log.warn("Invalid public key for {}:{}", accountIdentifier, deviceId, e); + return null; + } + })); + } + + private static Map getPrimaryKey(final UUID identifier, final byte deviceId) { + return Map.of( + KEY_ACCOUNT_UUID, getPartitionKey(identifier), + KEY_DEVICE_ID, getSortKey(deviceId)); + } + + private static AttributeValue getPartitionKey(final UUID accountUuid) { + return AttributeValues.fromUUID(accountUuid); + } + + private static AttributeValue getSortKey(final byte deviceId) { + return AttributeValues.fromInt(deviceId); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysManager.java new file mode 100644 index 000000000..48a15825b --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysManager.java @@ -0,0 +1,63 @@ +package org.whispersystems.textsecuregcm.storage; + +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import org.signal.libsignal.protocol.ecc.ECPublicKey; +import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; + +/** + * A client public key manager provides access to clients' public keys for use in transport-level authentication and + * encryption. + */ +public class ClientPublicKeysManager { + + private final ClientPublicKeys clientPublicKeys; + + public ClientPublicKeysManager(final ClientPublicKeys clientPublicKeys) { + this.clientPublicKeys = clientPublicKeys; + } + + /** + * Builds a {@link TransactWriteItem} that will store a public key for the given account/device. Intended for use when + * adding devices to accounts or creating new accounts. + * + * @param accountIdentifier the identifier for 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 {@code TransactWriteItem} that will store the given public key for the given account/device + */ + public TransactWriteItem buildTransactWriteItemForInsertion(final UUID accountIdentifier, + final byte deviceId, + final ECPublicKey publicKey) { + + return clientPublicKeys.buildTransactWriteItemForInsertion(accountIdentifier, deviceId, publicKey); + } + + /** + * Builds a {@link TransactWriteItem} that will remove the public key for the given account/device. Intended for use + * when removing devices from accounts or deleting/re-creating accounts. + * + * @param accountIdentifier the identifier for the target account + * @param deviceId the identifier for the target device + * + * @return a {@code TransactWriteItem} that will remove the public key for the given account/device + */ + public TransactWriteItem buildTransactWriteItemForDeletion(final UUID accountIdentifier, final byte deviceId) { + return clientPublicKeys.buildTransactWriteItemForDeletion(accountIdentifier, deviceId); + } + + /** + * Finds the public key for the given account/device. + * + * @param accountIdentifier the identifier for the target account + * @param deviceId the identifier for the target device + * + * @return a future that yields the Ed25519 public key for the given account/device, or empty if no public key was + * found + */ + public CompletableFuture> findPublicKey(final UUID accountIdentifier, final byte deviceId) { + return clientPublicKeys.findPublicKey(accountIdentifier, deviceId); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysTest.java new file mode 100644 index 000000000..856bdadb7 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysTest.java @@ -0,0 +1,48 @@ +package org.whispersystems.textsecuregcm.storage; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.Optional; +import java.util.UUID; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.signal.libsignal.protocol.ecc.Curve; +import org.signal.libsignal.protocol.ecc.ECPublicKey; +import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest; + +class ClientPublicKeysTest { + + private ClientPublicKeys clientPublicKeys; + + @RegisterExtension + static final DynamoDbExtension DYNAMO_DB_EXTENSION = + new DynamoDbExtension(DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS); + + @BeforeEach + void setUp() { + clientPublicKeys = new ClientPublicKeys(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), + DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName()); + } + + @Test + void buildTransactWriteItemForInsertionAndDeletion() { + final UUID accountIdentifier = UUID.randomUUID(); + final byte deviceId = Device.PRIMARY_ID; + final ECPublicKey publicKey = Curve.generateKeyPair().getPublicKey(); + + assertEquals(Optional.empty(), clientPublicKeys.findPublicKey(accountIdentifier, deviceId).join()); + + DYNAMO_DB_EXTENSION.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder() + .transactItems(clientPublicKeys.buildTransactWriteItemForInsertion(accountIdentifier, deviceId, publicKey)) + .build()); + + assertEquals(Optional.of(publicKey), clientPublicKeys.findPublicKey(accountIdentifier, deviceId).join()); + + DYNAMO_DB_EXTENSION.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder() + .transactItems(clientPublicKeys.buildTransactWriteItemForDeletion(accountIdentifier, deviceId)) + .build()); + + assertEquals(Optional.empty(), clientPublicKeys.findPublicKey(accountIdentifier, deviceId).join()); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java index 3d1c879ba..9127c1457 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtensionSchema.java @@ -326,6 +326,21 @@ public final class DynamoDbExtensionSchema { .build()), List.of()), + CLIENT_PUBLIC_KEYS("client_public_keys_test", + ClientPublicKeys.KEY_ACCOUNT_UUID, + ClientPublicKeys.KEY_DEVICE_ID, + List.of( + AttributeDefinition.builder() + .attributeName(ClientPublicKeys.KEY_ACCOUNT_UUID) + .attributeType(ScalarAttributeType.B) + .build(), + AttributeDefinition.builder() + .attributeName(ClientPublicKeys.KEY_DEVICE_ID) + .attributeType(ScalarAttributeType.N) + .build()), + List.of(), + List.of()), + USERNAMES("usernames_test", Accounts.ATTR_USERNAME_HASH, null,