Add an API endpoint for storing public keys

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -963,7 +963,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             zkAuthOperations, callingGenericZkSecretParams, clock),
         new ChallengeController(rateLimitChallengeManager, challengeConstraintChecker),
         new DeviceController(config.getLinkDeviceSecretConfiguration().secret().value(), accountsManager,
-            rateLimiters, rateLimitersCluster, config.getMaxDevices(), clock),
+            clientPublicKeysManager, rateLimiters, rateLimitersCluster, config.getMaxDevices(), clock),
         new DirectoryV2Controller(directoryV2CredentialsGenerator),
         new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),
             ReceiptCredentialPresentation::new),

service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java

@@ -26,6 +26,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import javax.annotation.Nullable;
 import javax.crypto.Mac;
 import javax.crypto.spec.SecretKeySpec;
@@ -56,11 +57,13 @@ import org.whispersystems.textsecuregcm.entities.DeviceInfoList;
 import org.whispersystems.textsecuregcm.entities.DeviceResponse;
 import org.whispersystems.textsecuregcm.entities.LinkDeviceRequest;
 import org.whispersystems.textsecuregcm.entities.PreKeySignatureValidator;
+import org.whispersystems.textsecuregcm.entities.SetPublicKeyRequest;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.ClientPublicKeysManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
 import org.whispersystems.textsecuregcm.storage.DeviceSpec;
@@ -76,6 +79,7 @@ public class DeviceController {
 
   private final Key verificationTokenKey;
   private final AccountsManager accounts;
+  private final ClientPublicKeysManager clientPublicKeysManager;
   private final RateLimiters rateLimiters;
   private final FaultTolerantRedisCluster usedTokenCluster;
   private final Map<String, Integer> maxDeviceConfiguration;
@@ -89,11 +93,13 @@ public class DeviceController {
 
   public DeviceController(byte[] linkDeviceSecret,
       AccountsManager accounts,
+      ClientPublicKeysManager clientPublicKeysManager,
       RateLimiters rateLimiters,
       FaultTolerantRedisCluster usedTokenCluster,
       Map<String, Integer> maxDeviceConfiguration, final Clock clock) {
     this.verificationTokenKey = new SecretKeySpec(linkDeviceSecret, VERIFICATION_TOKEN_ALGORITHM);
     this.accounts = accounts;
+    this.clientPublicKeysManager = clientPublicKeysManager;
     this.rateLimiters = rateLimiters;
     this.usedTokenCluster = usedTokenCluster;
     this.maxDeviceConfiguration = maxDeviceConfiguration;
@@ -278,6 +284,28 @@ public class DeviceController {
     accounts.updateDevice(auth.getAccount(), deviceId, d -> d.setCapabilities(capabilities));
   }
 
+  @PUT
+  @Produces(MediaType.APPLICATION_JSON)
+  @Path("/public_key")
+  @Operation(
+      summary = "Sets a public key for authentication",
+      description = """
+          Sets the authentication public key for the authenticated device. The public key will be used for
+          authentication in the nascent gRPC-over-Noise API. Existing devices must upload a public key before they can
+          use the gRPC-over-Noise API, and this endpoint exists to facilitate migration to the new API.
+          """
+  )
+  @ApiResponse(responseCode = "200", description = "Public key stored successfully")
+  @ApiResponse(responseCode = "401", description = "Account authentication check failed")
+  @ApiResponse(responseCode = "422", description = "Invalid request format")
+  public CompletableFuture<Void> setPublicKey(@Auth final AuthenticatedAccount auth,
+      final SetPublicKeyRequest setPublicKeyRequest) {
+
+    return clientPublicKeysManager.setPublicKey(auth.getAccount().getIdentifier(IdentityType.ACI),
+        auth.getAuthenticatedDevice().getId(),
+        setPublicKeyRequest.publicKey());
+  }
+
   private Mac getInitializedMac() {
     try {
       final Mac mac = Mac.getInstance(VERIFICATION_TOKEN_ALGORITHM);

service/src/main/java/org/whispersystems/textsecuregcm/entities/SetPublicKeyRequest.java

@@ -0,0 +1,17 @@
+package org.whispersystems.textsecuregcm.entities;
+
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import io.swagger.v3.oas.annotations.media.Schema;
+import org.signal.libsignal.protocol.ecc.ECPublicKey;
+import org.whispersystems.textsecuregcm.util.ECPublicKeyAdapter;
+
+public record SetPublicKeyRequest(
+    @JsonSerialize(using = ECPublicKeyAdapter.Serializer.class)
+    @JsonDeserialize(using = ECPublicKeyAdapter.Deserializer.class)
+    @Schema(type="string", description="""
+        The public key, serialized in libsignal's elliptic-curve public key format and then encoded as a standard (i.e.
+        not URL-safe), padded, base64-encoded string.
+        """)
+    ECPublicKey publicKey) {
+}

service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeys.java

@@ -9,11 +9,13 @@ import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
+import org.whispersystems.textsecuregcm.util.Util;
 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.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
 
 /**
@@ -35,6 +37,28 @@ public class ClientPublicKeys {
     this.tableName = tableName;
   }
 
+  /**
+   * Stores the given public key for the given account/device, overwriting any previously-stored public key. This method
+   * 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 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
+   */
+  CompletableFuture<Void> setPublicKey(final UUID accountIdentifier, final byte deviceId, final ECPublicKey publicKey) {
+    return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
+            .tableName(tableName)
+            .item(Map.of(
+                KEY_ACCOUNT_UUID, getPartitionKey(accountIdentifier),
+                KEY_DEVICE_ID, getSortKey(deviceId),
+                ATTR_PUBLIC_KEY, AttributeValues.fromByteArray(publicKey.serialize())))
+            .build())
+        .thenRun(Util.NOOP);
+  }
+
   /**
    * 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.

service/src/main/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysManager.java

@@ -18,6 +18,21 @@ public class ClientPublicKeysManager {
     this.clientPublicKeys = clientPublicKeys;
   }
 
+  /**
+   * Stores the given public key for the given account/device, overwriting any previously-stored public key. This method
+   * 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 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);
+  }
+
   /**
    * 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.