Add Accounts.SetZkCredentialKey

service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.grpc;
 
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HexFormat;
 import java.util.List;
 import java.util.UUID;
@@ -34,6 +35,8 @@ import org.signal.chat.account.SetRegistrationRecoveryPasswordRequest;
 import org.signal.chat.account.SetRegistrationRecoveryPasswordResponse;
 import org.signal.chat.account.SetUsernameLinkRequest;
 import org.signal.chat.account.SetUsernameLinkResponse;
+import org.signal.chat.account.SetZkCredentialKeyRequest;
+import org.signal.chat.account.SetZkCredentialKeyResponse;
 import org.signal.chat.account.SimpleAccountsGrpc;
 import org.signal.chat.account.UsernameNotAvailable;
 import org.signal.chat.common.AccountIdentifiers;
@@ -267,6 +270,24 @@ public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase {
     return SetRegistrationRecoveryPasswordResponse.getDefaultInstance();
   }
 
+  @Override
+  public SetZkCredentialKeyResponse setZkCredentialKey(final SetZkCredentialKeyRequest request) throws RateLimitExceededException {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    final Account authenticatedAccount = getAuthenticatedAccount();
+    final byte[] zkCredentialKey = request.getPublicKey().toByteArray();
+
+    if (Arrays.equals(authenticatedAccount.getZkCredentialKey(), zkCredentialKey)) {
+      return SetZkCredentialKeyResponse.getDefaultInstance();
+    }
+
+    rateLimiters.getSetZkCredentialKeyLimiter().validate(authenticatedDevice.accountIdentifier());
+
+    accountsManager.update(authenticatedAccount, account -> account.setZkCredentialKey(zkCredentialKey));
+
+    return SetZkCredentialKeyResponse.getDefaultInstance();
+  }
+
   private Account getAuthenticatedAccount() {
     return getAuthenticatedAccount(AuthenticationUtil.requireAuthenticatedDevice());
   }

service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java

@@ -36,6 +36,7 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
     USERNAME_LINK_LOOKUP_PER_IP("usernameLinkLookupPerIp", new RateLimiterConfig(100, Duration.ofSeconds(15), true)),
     CHECK_ACCOUNT_EXISTENCE("checkAccountExistence", new RateLimiterConfig(1000, Duration.ofSeconds(4), true)),
     REGISTRATION("registration", new RateLimiterConfig(6, Duration.ofSeconds(30), false)),
+    SET_ZK_CREDENTIAL_KEY("setZkCredentialKey", new RateLimiterConfig(5, Duration.ofDays(7), false)),
     VERIFICATION_PUSH_CHALLENGE("verificationPushChallenge", new RateLimiterConfig(5, Duration.ofSeconds(30), false)),
     VERIFICATION_CAPTCHA("verificationCaptcha", new RateLimiterConfig(10, Duration.ofSeconds(30), false)),
     RATE_LIMIT_RESET("rateLimitReset", new RateLimiterConfig(2, Duration.ofHours(12), false)),
@@ -106,8 +107,8 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
     return forDescriptor(For.VERIFY_DEVICE);
   }
 
-  public RateLimiter getMessagesLimiter() {
-    return forDescriptor(For.MESSAGES);
+  public RateLimiter getSetZkCredentialKeyLimiter() {
+    return forDescriptor(For.SET_ZK_CREDENTIAL_KEY);
   }
 
   public RateLimiter getPreKeysLimiter() {
@@ -162,6 +163,10 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
     return forDescriptor(For.REGISTRATION);
   }
 
+  public RateLimiter getMessagesLimiter() {
+    return forDescriptor(For.MESSAGES);
+  }
+
   public RateLimiter getRateLimitResetLimiter() {
     return forDescriptor(For.RATE_LIMIT_RESET);
   }

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

@@ -112,6 +112,10 @@ public class Account {
   @Nullable
   private BackupVoucher backupVoucher;
 
+  @JsonProperty("zck")
+  @Nullable
+  private byte[] zkCredentialKey;
+
   @JsonProperty
   private int version;
 
@@ -536,6 +540,15 @@ public class Account {
     this.usernameHolds = usernameHolds;
   }
 
+  @Nullable
+  public byte[] getZkCredentialKey() {
+    return zkCredentialKey;
+  }
+
+  public void setZkCredentialKey(@Nullable final byte[] zkCredentialKey) {
+    this.zkCredentialKey = zkCredentialKey;
+  }
+
   public void markStale() {
     stale = true;
   }

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

@@ -326,6 +326,9 @@ public class Accounts {
       // Carry over the existing backup voucher to the new account
       accountToCreate.setBackupVoucher(existingAccount.getBackupVoucher());
 
+      // Carry over the existing ZK credential key to the new account
+      accountToCreate.setZkCredentialKey(existingAccount.getZkCredentialKey());
+
       final List<TransactWriteItem> writeItems = new ArrayList<>();
 
       // If we're reclaiming an account that already has a username, we'd like to give the re-registering client

service/src/main/proto/org/signal/chat/account.proto

@@ -57,6 +57,9 @@ service Accounts {
 
   // Sets the registration recovery password for the authenticated account.
   rpc SetRegistrationRecoveryPassword(SetRegistrationRecoveryPasswordRequest) returns (SetRegistrationRecoveryPasswordResponse) {}
+
+  // Store a public key used to issue and verify zero-knowledge (anonymous) credentials for the account.
+  rpc SetZkCredentialKey(SetZkCredentialKeyRequest) returns (SetZkCredentialKeyResponse) {}
 }
 
 // Provides methods for looking up Signal accounts. Callers must not provide
@@ -265,3 +268,11 @@ message LookupUsernameLinkResponse {
     errors.NotFound not_found = 2 [(tag.reason) = "not_found"];
   }
 }
+
+message SetZkCredentialKeyRequest {
+  // A serialized Ristretto key with a one-byte type prefix
+  bytes public_key = 1 [(require.exactlySize) = 33];
+}
+
+message SetZkCredentialKeyResponse {
+}

service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java

@@ -8,11 +8,13 @@ package org.whispersystems.textsecuregcm.grpc;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.AdditionalMatchers.aryEq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -54,6 +56,7 @@ import org.signal.chat.account.SetRegistrationLockResponse;
 import org.signal.chat.account.SetRegistrationRecoveryPasswordRequest;
 import org.signal.chat.account.SetUsernameLinkRequest;
 import org.signal.chat.account.SetUsernameLinkResponse;
+import org.signal.chat.account.SetZkCredentialKeyRequest;
 import org.signal.chat.account.UsernameNotAvailable;
 import org.signal.chat.common.AccountIdentifiers;
 import org.signal.chat.errors.FailedPrecondition;
@@ -109,6 +112,8 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
     when(rateLimiters.getUsernameSetLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getUsernameLinkOperationLimiter()).thenReturn(rateLimiter);
 
+    when(rateLimiters.getSetZkCredentialKeyLimiter()).thenReturn(rateLimiter);
+
     when(registrationRecoveryPasswordsManager.store(any(), any()))
         .thenReturn(CompletableFuture.completedFuture(null));
 
@@ -712,4 +717,53 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
         () -> authenticatedServiceStub().setRegistrationRecoveryPassword(
             SetRegistrationRecoveryPasswordRequest.newBuilder().build()));
   }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void setZkCredentialKey(final boolean matchesCurrentZkCredentialKey) {
+
+    final byte[] publicKey = TestRandomUtil.nextBytes(33);
+
+    final Account account = mock(Account.class);
+
+    if (matchesCurrentZkCredentialKey) {
+      when(account.getZkCredentialKey()).thenReturn(publicKey);
+    }
+
+    when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI))
+        .thenReturn(Optional.of(account));
+
+    assertDoesNotThrow(() ->
+        authenticatedServiceStub().setZkCredentialKey(SetZkCredentialKeyRequest.newBuilder()
+                .setPublicKey(ByteString.copyFrom(publicKey))
+            .build()));
+
+    final int updateMethodCalls = matchesCurrentZkCredentialKey ? 0 : 1;
+
+    verify(accountsManager, times(updateMethodCalls)).update(eq(account), any(Consumer.class));
+    verify(account, times(updateMethodCalls)).setZkCredentialKey(aryEq(publicKey));
+  }
+
+  @Test
+  void setZkCredentialKeyRateLimited() throws Exception {
+
+    final byte[] publicKey = TestRandomUtil.nextBytes(33);
+    final Duration retryDuration = Duration.ofDays(1);
+
+    final Account account = mock(Account.class);
+    when(account.getUuid()).thenReturn(AUTHENTICATED_ACI);
+
+    when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI))
+        .thenReturn(Optional.of(account));
+    doThrow(new RateLimitExceededException(retryDuration))
+        .when(rateLimiter).validate(AUTHENTICATED_ACI);
+
+    GrpcTestUtils.assertRateLimitExceeded(retryDuration, () ->
+        authenticatedServiceStub().setZkCredentialKey(SetZkCredentialKeyRequest.newBuilder()
+            .setPublicKey(ByteString.copyFrom(publicKey))
+            .build()));
+
+    verify(accountsManager, never()).update(any(Account.class), any(Consumer.class));
+    verify(account, never()).setZkCredentialKey(any());
+  }
 }

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

@@ -479,6 +479,7 @@ class AccountsTest {
 
     // the backup credential request and share-set are always preserved across account reclaims
     existingAccount.setBackupCredentialRequests(TestRandomUtil.nextBytes(32), TestRandomUtil.nextBytes(32));
+    existingAccount.setZkCredentialKey(TestRandomUtil.nextBytes(32));
     createAccount(existingAccount);
     final Account secondAccount =
         generateAccount(e164, UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
@@ -490,6 +491,7 @@ class AccountsTest {
         .isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MESSAGES).orElseThrow());
     assertThat(reclaimed.getBackupCredentialRequest(BackupCredentialType.MEDIA).orElseThrow())
         .isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MEDIA).orElseThrow());
+    assertThat(reclaimed.getZkCredentialKey()).isEqualTo(existingAccount.getZkCredentialKey());
   }
 
   @Test
@@ -500,9 +502,11 @@ class AccountsTest {
     final UUID existingPni = UUID.randomUUID();
     final Account existingAccount = generateAccount(e164, existingUuid, existingPni, List.of(device));
 
-    // Backup vouchers should be carried over accross re-registration
+    // Backup vouchers should be carried over across re-registration
     final Account.BackupVoucher bv = new Account.BackupVoucher(1, Instant.now().plus(Duration.ofDays(1)));
     existingAccount.setBackupVoucher(bv);
+    // ZK credential keys should be carried over across re-registration
+    existingAccount.setZkCredentialKey(TestRandomUtil.nextBytes(32));
 
     createAccount(existingAccount);