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.

service/src/test/java/org/whispersystems/textsecuregcm/controllers/DeviceControllerTest.java

@@ -61,6 +61,7 @@ import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.LinkDeviceRequest;
+import org.whispersystems.textsecuregcm.entities.SetPublicKeyRequest;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -68,6 +69,7 @@ import org.whispersystems.textsecuregcm.mappers.DeviceLimitExceededExceptionMapp
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 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;
@@ -85,6 +87,7 @@ import org.whispersystems.textsecuregcm.util.VerificationCode;
 class DeviceControllerTest {
 
   private static AccountsManager accountsManager = mock(AccountsManager.class);
+  private static ClientPublicKeysManager clientPublicKeysManager = mock(ClientPublicKeysManager.class);
   private static RateLimiters rateLimiters = mock(RateLimiters.class);
   private static RateLimiter rateLimiter = mock(RateLimiter.class);
   private static RedisAdvancedClusterCommands<String, String> commands = mock(RedisAdvancedClusterCommands.class);
@@ -101,6 +104,7 @@ class DeviceControllerTest {
   private static DeviceController deviceController = new DeviceController(
       generateLinkDeviceSecret(),
       accountsManager,
+      clientPublicKeysManager,
       rateLimiters,
       RedisClusterHelper.builder()
           .stringCommands(commands)
@@ -147,6 +151,9 @@ class DeviceControllerTest {
     when(accountsManager.getByE164(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(account));
     when(accountsManager.getByE164(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(maxedAccount));
 
+    when(clientPublicKeysManager.setPublicKey(any(), anyByte(), any()))
+        .thenReturn(CompletableFuture.completedFuture(null));
+
     AccountsHelper.setupMockUpdate(accountsManager);
   }
 
@@ -807,4 +814,20 @@ class DeviceControllerTest {
         Arguments.of("e552603a-1492-4de6-872d-bac19a2825b4.1691096565171:This is not valid base64", tokenTimestamp)
     );
   }
+
+  @Test
+  void setPublicKey() {
+    final SetPublicKeyRequest request = new SetPublicKeyRequest(Curve.generateKeyPair().getPublicKey());
+
+    try (final Response response = resources.getJerseyTest()
+        .target("/v1/devices/public_key")
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .put(Entity.entity(request, MediaType.APPLICATION_JSON_TYPE))) {
+
+      assertEquals(204, response.getStatus());
+    }
+
+    verify(clientPublicKeysManager).setPublicKey(AuthHelper.VALID_UUID, AuthHelper.VALID_DEVICE.getId(), request.publicKey());
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/ClientPublicKeysTest.java

@@ -45,4 +45,17 @@ class ClientPublicKeysTest {
 
     assertEquals(Optional.empty(), clientPublicKeys.findPublicKey(accountIdentifier, deviceId).join());
   }
+
+  @Test
+  void setPublicKey() {
+    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());
+
+    clientPublicKeys.setPublicKey(accountIdentifier, deviceId, publicKey).join();
+
+    assertEquals(Optional.of(publicKey), clientPublicKeys.findPublicKey(accountIdentifier, deviceId).join());
+  }
 }