accept encrypted username with confirm-username-hash requests

integration-tests/src/test/java/org/signal/integration/AccountTest.java

@@ -18,6 +18,7 @@
 import org.whispersystems.textsecuregcm.entities.AccountIdentifierResponse;
 import org.whispersystems.textsecuregcm.entities.AccountIdentityResponse;
 import org.whispersystems.textsecuregcm.entities.ConfirmUsernameHashRequest;
+import org.whispersystems.textsecuregcm.entities.EncryptedUsername;
 import org.whispersystems.textsecuregcm.entities.ReserveUsernameHashRequest;
 import org.whispersystems.textsecuregcm.entities.ReserveUsernameHashResponse;
 import org.whispersystems.textsecuregcm.entities.UsernameHashResponse;
@@ -87,9 +88,10 @@ private static void verifyFullUsernameLifecycle(final TestUser user) throws Base
         .orElseThrow();
 
     // confirm a username
-    final ConfirmUsernameHashRequest confirmUsernameHashRequest = new ConfirmUsernameHashRequest(
+   final ConfirmUsernameHashRequest confirmUsernameHashRequest = new ConfirmUsernameHashRequest(
         reservedUsername.getHash(),
-        reservedUsername.generateProof()
+        reservedUsername.generateProof(),
+        new EncryptedUsername("cluck cluck i'm a parrot".getBytes())
     );
     // try unauthorized
     Operations

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

@@ -350,17 +350,15 @@ public UsernameHashResponse confirmUsernameHash(
       throw new WebApplicationException(Response.status(422).build());
     }
 
-    // Whenever a valid request for a username change arrives,
-    // we're making sure to clear username link. This may happen before and username changes are written to the db
-    // but verifying zk proof means that request itself is valid from the client's perspective
-    clearUsernameLink(auth.getAccount());
-
     try {
-      final Account account = accounts.confirmReservedUsernameHash(auth.getAccount(), confirmRequest.usernameHash());
-      return account
-          .getUsernameHash()
-          .map(UsernameHashResponse::new)
-          .orElseThrow(() -> new IllegalStateException("Could not get username after setting"));
+      final Account account = accounts.confirmReservedUsernameHash(
+          auth.getAccount(),
+          confirmRequest.usernameHash(),
+          Optional.ofNullable(confirmRequest.encryptedUsername()).map(EncryptedUsername::usernameLinkEncryptedValue).orElse(null));
+      final UUID linkHandle = account.getUsernameLinkHandle();
+      return new UsernameHashResponse(
+          account.getUsernameHash().orElseThrow(() -> new IllegalStateException("Could not get username after setting")),
+          linkHandle == null ? null : new UsernameLinkHandle(linkHandle));
     } catch (final UsernameReservationNotFoundException e) {
       throw new WebApplicationException(Status.CONFLICT);
     } catch (final UsernameHashNotAvailableException e) {

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

@@ -5,12 +5,18 @@
 
 package org.whispersystems.textsecuregcm.entities;
 
+import javax.validation.Valid;
+
+import javax.annotation.Nullable;
+
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import org.whispersystems.textsecuregcm.controllers.AccountController;
 import org.whispersystems.textsecuregcm.util.ByteArrayBase64UrlAdapter;
 import org.whispersystems.textsecuregcm.util.ExactlySize;
 
+import io.swagger.v3.oas.annotations.media.Schema;
+
 public record ConfirmUsernameHashRequest(
     @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
     @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
@@ -19,5 +25,10 @@ public record ConfirmUsernameHashRequest(
 
     @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
     @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
-    byte[] zkProof
+    byte[] zkProof,
+
+    @Schema(description = "The encrypted username to be stored for username links")
+    @Nullable
+    @Valid
+    EncryptedUsername encryptedUsername
 ) {}

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

@@ -7,15 +7,24 @@
 
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import io.swagger.v3.oas.annotations.media.Schema;
 import org.whispersystems.textsecuregcm.controllers.AccountController;
 import org.whispersystems.textsecuregcm.util.ByteArrayBase64UrlAdapter;
 import org.whispersystems.textsecuregcm.util.ExactlySize;
+
+import javax.annotation.Nullable;
 import javax.validation.Valid;
 
 public record UsernameHashResponse(
     @Valid
     @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
     @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
     @ExactlySize(AccountController.USERNAME_HASH_LENGTH)
-    byte[] usernameHash
+    @Schema(description = "The hash of the confirmed username, as supplied in the request")
+    byte[] usernameHash,
+
+    @Nullable
+    @Valid
+    @Schema(description = "A handle that can be included in username links to retrieve the stored encrypted username")
+    UsernameLinkHandle usernameLinkHandle
 ) {}

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

@@ -30,6 +30,8 @@
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.util.AsyncTimerUtil;
@@ -386,15 +388,18 @@ public void reserveUsernameHash(
    * @param usernameHash believed to be available
    * @throws ContestedOptimisticLockException if the account has been updated or the username has taken by someone else
    */
-  public void confirmUsernameHash(final Account account, final byte[] usernameHash)
+  public void confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername)
       throws ContestedOptimisticLockException {
     final long startNanos = System.nanoTime();
 
     final Optional<byte[]> maybeOriginalUsernameHash = account.getUsernameHash();
     final Optional<byte[]> maybeOriginalReservationHash = account.getReservedUsernameHash();
+    final Optional<UUID> maybeOriginalUsernameLinkHandle = Optional.ofNullable(account.getUsernameLinkHandle());
+    final Optional<byte[]> maybeOriginalEncryptedUsername = account.getEncryptedUsername();
 
     account.setUsernameHash(usernameHash);
     account.setReservedUsernameHash(null);
+    account.setUsernameLinkDetails(encryptedUsername == null ? null : UUID.randomUUID(), encryptedUsername);
 
     boolean succeeded = false;
 
@@ -420,21 +425,32 @@ public void confirmUsernameHash(final Account account, final byte[] usernameHash
               .build())
           .build());
 
+      final StringBuilder updateExpr = new StringBuilder("SET #data = :data, #username_hash = :username_hash");
+      final Map<String, AttributeValue> expressionAttributeValues = new HashMap<>(Map.of(
+          ":data", AttributeValues.fromByteArray(SystemMapper.jsonMapper().writeValueAsBytes(account)),
+          ":username_hash", AttributeValues.fromByteArray(usernameHash),
+          ":version", AttributeValues.fromInt(account.getVersion()),
+          ":version_increment", AttributeValues.fromInt(1)));
+      if (account.getUsernameLinkHandle() != null) {
+        updateExpr.append(", #ul = :ul");
+        expressionAttributeValues.put(":ul", AttributeValues.fromUUID(account.getUsernameLinkHandle()));
+      } else {
+        updateExpr.append(" REMOVE #ul");
+      }
+      updateExpr.append(" ADD #version :version_increment");
+
       writeItems.add(
           TransactWriteItem.builder()
               .update(Update.builder()
                   .tableName(accountsTableName)
                   .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
-                  .updateExpression("SET #data = :data, #username_hash = :username_hash ADD #version :version_increment")
+                  .updateExpression(updateExpr.toString())
                   .conditionExpression("#version = :version")
                   .expressionAttributeNames(Map.of("#data", ATTR_ACCOUNT_DATA,
                       "#username_hash", ATTR_USERNAME_HASH,
+                      "#ul", ATTR_USERNAME_LINK_UUID,
                       "#version", ATTR_VERSION))
-                  .expressionAttributeValues(Map.of(
-                      ":data", AttributeValues.fromByteArray(SystemMapper.jsonMapper().writeValueAsBytes(account)),
-                      ":username_hash", AttributeValues.fromByteArray(usernameHash),
-                      ":version", AttributeValues.fromInt(account.getVersion()),
-                      ":version_increment", AttributeValues.fromInt(1)))
+                  .expressionAttributeValues(expressionAttributeValues)
                   .build())
               .build());
 
@@ -460,6 +476,7 @@ public void confirmUsernameHash(final Account account, final byte[] usernameHash
       if (!succeeded) {
         account.setUsernameHash(maybeOriginalUsernameHash.orElse(null));
         account.setReservedUsernameHash(maybeOriginalReservationHash.orElse(null));
+        account.setUsernameLinkDetails(maybeOriginalUsernameLinkHandle.orElse(null), maybeOriginalEncryptedUsername.orElse(null));
       }
       SET_USERNAME_TIMER.record(System.nanoTime() - startNanos, TimeUnit.NANOSECONDS);
     }

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

@@ -503,11 +503,12 @@ public void persistAccount(final Account account) throws UsernameHashNotAvailabl
    *
    * @param account the account to update
    * @param reservedUsernameHash the previously reserved username hash
+   * @param encryptedUsername the encrypted form of the previously reserved username for the username link
    * @return the updated account with the username hash field set
    * @throws UsernameHashNotAvailableException if the reserved username hash has been taken (because the reservation expired)
    * @throws UsernameReservationNotFoundException if `reservedUsernameHash` was not reserved in the account
    */
-  public Account confirmReservedUsernameHash(final Account account, final byte[] reservedUsernameHash) throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
+  public Account confirmReservedUsernameHash(final Account account, final byte[] reservedUsernameHash, @Nullable final byte[] encryptedUsername) throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
     if (!experimentEnrollmentManager.isEnrolled(account.getUuid(), USERNAME_EXPERIMENT_NAME)) {
       throw new UsernameHashNotAvailableException();
     }
@@ -532,7 +533,7 @@ public Account confirmReservedUsernameHash(final Account account, final byte[] r
           if (!accounts.usernameHashAvailable(Optional.of(account.getUuid()), reservedUsernameHash)) {
             throw new UsernameHashNotAvailableException();
           }
-          accounts.confirmUsernameHash(a, reservedUsernameHash);
+          accounts.confirmUsernameHash(a, reservedUsernameHash, encryptedUsername);
         },
         () -> accounts.getByAccountIdentifier(account.getUuid()).orElseThrow(),
         AccountChangeValidator.USERNAME_CHANGE_VALIDATOR);
@@ -731,6 +732,7 @@ private static Account cloneAccount(final Account account) {
     try {
       final Account clone = mapper.readValue(mapper.writeValueAsBytes(account), Account.class);
       clone.setUuid(account.getUuid());
+      clone.setUsernameLinkHandle(account.getUsernameLinkHandle());
 
       return clone;
     } catch (final IOException e) {

service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java

@@ -82,6 +82,12 @@ public void configure(Subparser subparser) {
         .required(true)
         .help("The username hash to assign");
 
+    subparser.addArgument("-e", "--encryptedUsername")
+        .dest("encryptedUsername")
+        .type(String.class)
+        .required(false)
+        .help("The encrypted username for the username link");
+
     subparser.addArgument("-a", "--aci")
         .dest("aci")
         .type(String.class)
@@ -210,14 +216,19 @@ protected void run(Environment environment, Namespace namespace,
         experimentEnrollmentManager, registrationRecoveryPasswordsManager, Clock.systemUTC());
 
     final String usernameHash = namespace.getString("usernameHash");
+    final String encryptedUsername = namespace.getString("encryptedUsername");
     final UUID accountIdentifier = UUID.fromString(namespace.getString("aci"));
 
     accountsManager.getByAccountIdentifier(accountIdentifier).ifPresentOrElse(account -> {
           try {
             final AccountsManager.UsernameReservation reservation = accountsManager.reserveUsernameHash(account,
                 List.of(Base64.getUrlDecoder().decode(usernameHash)));
-            final Account result = accountsManager.confirmReservedUsernameHash(account, Base64.getUrlDecoder().decode(usernameHash));
-            System.out.println("New username hash: " + usernameHash);
+            final Account result = accountsManager.confirmReservedUsernameHash(
+                account,
+                reservation.reservedUsernameHash(),
+                encryptedUsername == null ? null : Base64.getUrlDecoder().decode(encryptedUsername));
+            System.out.println("New username hash: " + Base64.getUrlEncoder().encodeToString(result.getUsernameHash().orElseThrow()));
+            System.out.println("New username link handle: " + result.getUsernameLinkHandle().toString());
           } catch (final UsernameHashNotAvailableException e) {
             throw new IllegalArgumentException("Username hash already taken");
           } catch (final UsernameReservationNotFoundException e) {