From 4fc394936741bdf2ba9a6a641f57ae9350056d2d Mon Sep 17 00:00:00 2001
From: Katherine Yen <katherine@signal.org>
Date: Thu, 9 Feb 2023 09:02:53 -0800
Subject: [PATCH] Add zkproof validation in username flow

---
 pom.xml                                       |  2 +-
 .../textsecuregcm/WhisperServerService.java   |  4 +-
 .../controllers/AccountController.java        | 11 +++
 .../entities/ConfirmUsernameHashRequest.java  |  6 +-
 .../util/UsernameHashZkProofVerifier.java     | 10 +++
 .../controllers/AccountControllerTest.java    | 74 ++++++++++++++++---
 6 files changed, 95 insertions(+), 12 deletions(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/util/UsernameHashZkProofVerifier.java

diff --git a/pom.xml b/pom.xml
index b2a0bd118..afc784a65 100644
--- a/pom.xml
+++ b/pom.xml
@@ -296,7 +296,7 @@
       <dependency>
         <groupId>org.signal</groupId>
         <artifactId>libsignal-server</artifactId>
-        <version>0.21.1</version>
+        <version>0.22.0</version>
       </dependency>
       <dependency>
         <groupId>org.apache.logging.log4j</groupId>
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 843c468a7..f7f262472 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -213,6 +213,7 @@ import org.whispersystems.textsecuregcm.subscriptions.StripeManager;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
 import org.whispersystems.textsecuregcm.util.HostnameUtil;
+import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
 import org.whispersystems.textsecuregcm.util.logging.LoggingUnhandledExceptionMapper;
 import org.whispersystems.textsecuregcm.util.logging.UncaughtExceptionHandler;
 import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener;
@@ -475,6 +476,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
     RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords);
+    UsernameHashZkProofVerifier usernameHashZkProofVerifier = new UsernameHashZkProofVerifier();
 
     RegistrationServiceClient  registrationServiceClient  = new RegistrationServiceClient(config.getRegistrationServiceConfiguration().getHost(), config.getRegistrationServiceConfiguration().getPort(), config.getRegistrationServiceConfiguration().getApiKey(), config.getRegistrationServiceConfiguration().getRegistrationCaCertificate(), registrationCallbackExecutor);
     SecureBackupClient         secureBackupClient         = new SecureBackupClient(backupCredentialsGenerator, backupServiceExecutor, config.getSecureBackupServiceConfiguration());
@@ -679,7 +681,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new AccountController(pendingAccountsManager, accountsManager, rateLimiters,
             registrationServiceClient, dynamicConfigurationManager, turnTokenGenerator, config.getTestDevices(),
             captchaChecker, pushNotificationManager, changeNumberManager, registrationLockVerificationManager,
-            registrationRecoveryPasswordsManager, clock));
+            registrationRecoveryPasswordsManager, usernameHashZkProofVerifier, clock));
 
     environment.jersey().register(new KeysController(rateLimiters, keys, accountsManager));
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
index c91f2839c..b17a4d5d0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
@@ -54,6 +54,7 @@ import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
 import org.apache.commons.lang3.StringUtils;
+import org.signal.libsignal.usernames.BaseUsernameException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
@@ -107,6 +108,7 @@ import org.whispersystems.textsecuregcm.util.HeaderUtils;
 import org.whispersystems.textsecuregcm.util.ImpossiblePhoneNumberException;
 import org.whispersystems.textsecuregcm.util.NonNormalizedPhoneNumberException;
 import org.whispersystems.textsecuregcm.util.Optionals;
+import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
 import org.whispersystems.textsecuregcm.util.Util;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
@@ -160,6 +162,7 @@ public class AccountController {
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager;
   private final ChangeNumberManager changeNumberManager;
   private final Clock clock;
+  private final UsernameHashZkProofVerifier usernameHashZkProofVerifier;
 
 
   @VisibleForTesting
@@ -178,6 +181,7 @@ public class AccountController {
       ChangeNumberManager changeNumberManager,
       RegistrationLockVerificationManager registrationLockVerificationManager,
       RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
+      UsernameHashZkProofVerifier usernameHashZkProofVerifier,
       Clock clock
   ) {
     this.pendingAccounts = pendingAccounts;
@@ -192,6 +196,7 @@ public class AccountController {
     this.registrationLockVerificationManager = registrationLockVerificationManager;
     this.changeNumberManager = changeNumberManager;
     this.registrationRecoveryPasswordsManager = registrationRecoveryPasswordsManager;
+    this.usernameHashZkProofVerifier = usernameHashZkProofVerifier;
     this.clock = clock;
   }
 
@@ -712,6 +717,12 @@ public class AccountController {
       @NotNull @Valid ConfirmUsernameHashRequest confirmRequest) throws RateLimitExceededException {
     rateLimiters.getUsernameSetLimiter().validate(auth.getAccount().getUuid());
 
+    try {
+      usernameHashZkProofVerifier.verifyProof(confirmRequest.zkProof(), confirmRequest.usernameHash());
+    } catch (final BaseUsernameException e) {
+      throw new WebApplicationException(Response.status(422).build());
+    }
+
     try {
       final Account account = accounts.confirmReservedUsernameHash(auth.getAccount(), confirmRequest.usernameHash());
       return account
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/entities/ConfirmUsernameHashRequest.java b/service/src/main/java/org/whispersystems/textsecuregcm/entities/ConfirmUsernameHashRequest.java
index 0ead70093..eb481aa3e 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/ConfirmUsernameHashRequest.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/ConfirmUsernameHashRequest.java
@@ -15,5 +15,9 @@ public record ConfirmUsernameHashRequest(
     @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
     @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
     @ExactlySize(AccountController.USERNAME_HASH_LENGTH)
-    byte[] usernameHash
+    byte[] usernameHash,
+
+    @JsonSerialize(using = ByteArrayBase64UrlAdapter.Serializing.class)
+    @JsonDeserialize(using = ByteArrayBase64UrlAdapter.Deserializing.class)
+    byte[] zkProof
 ) {}
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/UsernameHashZkProofVerifier.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/UsernameHashZkProofVerifier.java
new file mode 100644
index 000000000..4570060be
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/UsernameHashZkProofVerifier.java
@@ -0,0 +1,10 @@
+package org.whispersystems.textsecuregcm.util;
+
+import org.signal.libsignal.usernames.BaseUsernameException;
+import org.signal.libsignal.usernames.Username;
+
+public class UsernameHashZkProofVerifier {
+  public void verifyProof(byte[] proof, byte[] hash) throws BaseUsernameException {
+    Username.verifyProof(proof, hash);
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
index 760c1096e..fb549b092 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
@@ -64,6 +64,7 @@ import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.stubbing.Answer;
+import org.signal.libsignal.usernames.BaseUsernameException;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccount;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialsGenerator;
@@ -119,6 +120,7 @@ import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 import org.whispersystems.textsecuregcm.util.MockUtils;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 import org.whispersystems.textsecuregcm.util.TestClock;
+import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
 
 @ExtendWith(DropwizardExtensionsSupport.class)
 class AccountControllerTest {
@@ -142,7 +144,8 @@ class AccountControllerTest {
   private static final byte[] USERNAME_HASH_2 = Base64.getUrlDecoder().decode(BASE_64_URL_USERNAME_HASH_2);
   private static final byte[] INVALID_USERNAME_HASH = Base64.getDecoder().decode(INVALID_BASE_64_URL_USERNAME_HASH);
   private static final byte[] TOO_SHORT_USERNAME_HASH = Base64.getUrlDecoder().decode(TOO_SHORT_BASE_64_URL_USERNAME_HASH);
-
+  private static final String BASE_64_URL_ZK_PROOF = "2kambOgmdeeIO0faCMgR6HR4G2BQ5bnhXdIe9ZuZY0NmQXSra5BzDBQ7jzy1cvoEqUHYLpBYMrXudkYPJaWoQg";
+  private static final byte[] ZK_PROOF = Base64.getUrlDecoder().decode(BASE_64_URL_ZK_PROOF);
   private static final UUID   SENDER_REG_LOCK_UUID = UUID.randomUUID();
   private static final UUID   SENDER_TRANSFER_UUID = UUID.randomUUID();
 
@@ -179,10 +182,10 @@ class AccountControllerTest {
   private static RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
       RegistrationRecoveryPasswordsManager.class);
   private static ClientPresenceManager clientPresenceManager = mock(ClientPresenceManager.class);
+  private static final UsernameHashZkProofVerifier usernameZkProofVerifier = mock(UsernameHashZkProofVerifier.class);
   private static TestClock testClock = TestClock.now();
 
   private static DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
-
   private byte[] registration_lock_key = new byte[32];
 
   private static final SecureBackupServiceConfiguration BACKUP_CFG = MockUtils.buildMock(
@@ -220,6 +223,7 @@ class AccountControllerTest {
           changeNumberManager,
           registrationLockVerificationManager,
           registrationRecoveryPasswordsManager,
+          usernameZkProofVerifier,
           testClock))
       .build();
 
@@ -356,7 +360,8 @@ class AccountControllerTest {
         captchaChecker,
         pushNotificationManager,
         changeNumberManager,
-        clientPresenceManager);
+        clientPresenceManager,
+        usernameZkProofVerifier);
 
     clearInvocations(AuthHelper.DISABLED_DEVICE);
   }
@@ -1751,7 +1756,8 @@ class AccountControllerTest {
   }
 
   @Test
-  void testCommitUsername() throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
+  void testConfirmUsernameHash()
+      throws UsernameHashNotAvailableException, UsernameReservationNotFoundException, BaseUsernameException {
     Account account = mock(Account.class);
     when(account.getUsernameHash()).thenReturn(Optional.of(USERNAME_HASH_1));
     when(accountsManager.confirmReservedUsernameHash(any(), eq(USERNAME_HASH_1))).thenReturn(account);
@@ -1760,13 +1766,15 @@ class AccountControllerTest {
             .target("/v1/accounts/username_hash/confirm")
             .request()
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1)));
+            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1, ZK_PROOF)));
     assertThat(response.getStatus()).isEqualTo(200);
     assertArrayEquals(response.readEntity(UsernameHashResponse.class).usernameHash(), USERNAME_HASH_1);
+    verify(usernameZkProofVerifier).verifyProof(ZK_PROOF, USERNAME_HASH_1);
   }
 
   @Test
-  void testCommitUnreservedUsername() throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
+  void testConfirmUnreservedUsernameHash()
+      throws UsernameHashNotAvailableException, UsernameReservationNotFoundException, BaseUsernameException {
     when(accountsManager.confirmReservedUsernameHash(any(), eq(USERNAME_HASH_1)))
         .thenThrow(new UsernameReservationNotFoundException());
     Response response =
@@ -1774,12 +1782,14 @@ class AccountControllerTest {
             .target("/v1/accounts/username_hash/confirm")
             .request()
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1)));
+            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1, ZK_PROOF)));
     assertThat(response.getStatus()).isEqualTo(409);
+    verify(usernameZkProofVerifier).verifyProof(ZK_PROOF, USERNAME_HASH_1);
   }
 
   @Test
-  void testCommitLapsedUsername() throws UsernameHashNotAvailableException, UsernameReservationNotFoundException {
+  void testConfirmLapsedUsernameHash()
+      throws UsernameHashNotAvailableException, UsernameReservationNotFoundException, BaseUsernameException {
     when(accountsManager.confirmReservedUsernameHash(any(), eq(USERNAME_HASH_1)))
         .thenThrow(new UsernameHashNotAvailableException());
     Response response =
@@ -1787,8 +1797,54 @@ class AccountControllerTest {
             .target("/v1/accounts/username_hash/confirm")
             .request()
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1)));
+            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1, ZK_PROOF)));
     assertThat(response.getStatus()).isEqualTo(410);
+    verify(usernameZkProofVerifier).verifyProof(ZK_PROOF, USERNAME_HASH_1);
+  }
+
+  @Test
+  void testConfirmUsernameHashInvalidBase64UrlEncoding() {
+    Response response =
+        resources.getJerseyTest()
+            .target("/v1/accounts/username_hash/confirm")
+            .request()
+            .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+            .put(Entity.json(
+                // Has '+' and '='characters which are invalid in base64url
+                """
+                  {
+                    "usernameHash": "jh1jJ50oGn9wUXAFNtDus6AJgWOQ6XbZzF+wCv7OOQs=",
+                    "zkProof": "iYXE0QPK60PS3lGa-xdNv0GlXA3B03xQLzltSf-2xmscyS_8fjy5H9ymfaEr62PcVY7tsWhWjOOvcCnhmP_HS="
+                  }
+                """));
+    assertThat(response.getStatus()).isEqualTo(422);
+    verifyNoInteractions(usernameZkProofVerifier);
+  }
+
+  @Test
+  void testConfirmUsernameHashInvalidHashSize() {
+    byte[] usernameHash = new byte[31];
+    Response response =
+        resources.getJerseyTest()
+            .target("/v1/accounts/username_hash/confirm")
+            .request()
+            .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+            .put(Entity.json(new ConfirmUsernameHashRequest(usernameHash, ZK_PROOF)));
+    assertThat(response.getStatus()).isEqualTo(422);
+    verifyNoInteractions(usernameZkProofVerifier);
+  }
+
+  @Test
+  void testCommitUsernameHashWithInvalidProof() throws BaseUsernameException {
+    doThrow(new BaseUsernameException("invalid username")).when(usernameZkProofVerifier).verifyProof(eq(ZK_PROOF), eq(USERNAME_HASH_1));
+    Response response =
+        resources.getJerseyTest()
+            .target("/v1/accounts/username_hash/confirm")
+            .request()
+            .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+            .put(Entity.json(new ConfirmUsernameHashRequest(USERNAME_HASH_1, ZK_PROOF)));
+    assertThat(response.getStatus()).isEqualTo(422);
+    verify(usernameZkProofVerifier).verifyProof(ZK_PROOF, USERNAME_HASH_1);
   }
 
   @Test