Add a data store for client public keys for transport-level authentication/encryption

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<Optional<ECPublicKey>> 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<String, AttributeValue> 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);
+  }
+}

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<Optional<ECPublicKey>> findPublicKey(final UUID accountIdentifier, final byte deviceId) {
+    return clientPublicKeys.findPublicKey(accountIdentifier, deviceId);
+  }
+}

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());
+  }
+}

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,