Add rotation_id to ZK Credential key

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

@@ -6,11 +6,14 @@
 package org.whispersystems.textsecuregcm.grpc;
 
 import com.google.protobuf.ByteString;
+import java.security.SecureRandom;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HexFormat;
 import java.util.List;
+import java.util.Objects;
 import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
 import org.signal.chat.account.ClearRegistrationLockRequest;
 import org.signal.chat.account.ClearRegistrationLockResponse;
 import org.signal.chat.account.ConfigureUnidentifiedAccessRequest;
@@ -62,6 +65,8 @@ import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
 
 public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase {
 
+  private static final SecureRandom SECURE_RANDOM = new SecureRandom();
+
   private final AccountsManager accountsManager;
   private final RateLimiters rateLimiters;
   private final UsernameHashZkProofVerifier usernameHashZkProofVerifier;
@@ -276,14 +281,24 @@ public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase {
     final byte[] zkCredentialKey = request.getPublicKey().toByteArray();
 
     if (Arrays.equals(authenticatedAccount.getZkCredentialKey(), zkCredentialKey)) {
-      return SetZkCredentialKeyResponse.getDefaultInstance();
+      return SetZkCredentialKeyResponse.newBuilder()
+          .setRotationId(Objects.requireNonNull(authenticatedAccount.getZkCredentialKeyRotationId()))
+          .build();
     }
 
     rateLimiters.getSetZkCredentialKeyLimiter().validate(authenticatedDevice.accountIdentifier());
 
-    accountsManager.update(authenticatedDevice.accountIdentifier(), account -> account.setZkCredentialKey(zkCredentialKey));
+    // It is technically fine from the credential's perspective if it is zero, but it's clearer to never have the default value
+    final long rotationId = SECURE_RANDOM.nextLong(1, Long.MAX_VALUE);
 
-    return SetZkCredentialKeyResponse.getDefaultInstance();
+    accountsManager.update(authenticatedDevice.accountIdentifier(), account -> {
+      account.setZkCredentialKey(zkCredentialKey);
+      account.setZkCredentialKeyRotationId(rotationId);
+    });
+
+    return SetZkCredentialKeyResponse.newBuilder()
+        .setRotationId(rotationId)
+        .build();
   }
 
   private Account getAuthenticatedAccount() {

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

@@ -17,7 +17,6 @@ import com.fasterxml.jackson.databind.JsonSerializer;
 import com.fasterxml.jackson.databind.SerializerProvider;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
-import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
 import java.io.IOException;
@@ -131,6 +130,10 @@ public class Account {
   @Nullable
   private byte[] zkCredentialKey;
 
+  @JsonProperty("zckr")
+  @Nullable
+  private Long zkCredentialKeyRotationId;
+
   @JsonProperty
   private int version;
 
@@ -566,6 +569,15 @@ public class Account {
     this.zkCredentialKey = zkCredentialKey;
   }
 
+  @Nullable
+  public Long getZkCredentialKeyRotationId() {
+    return zkCredentialKeyRotationId;
+  }
+
+  public void setZkCredentialKeyRotationId(@Nullable final Long zkCredentialKeyRotationId) {
+    this.zkCredentialKeyRotationId = zkCredentialKeyRotationId;
+  }
+
   public void markStale() {
     stale = true;
   }

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

@@ -275,4 +275,9 @@ message SetZkCredentialKeyRequest {
 }
 
 message SetZkCredentialKeyResponse {
+  // A random, non-zero, value that must be included in credential requests using the key.
+  //
+  // This value allows the server to ratchet the resulting binding identity,
+  // as reverting to the previous key will result in a new rotation ID.
+  uint64 rotation_id = 1;
 }

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

@@ -26,6 +26,7 @@ import java.util.HexFormat;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -55,6 +56,7 @@ 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.SetZkCredentialKeyResponse;
 import org.signal.chat.account.UsernameNotAvailable;
 import org.signal.chat.common.AccountIdentifiers;
 import org.signal.chat.errors.FailedPrecondition;
@@ -708,21 +710,29 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
   void setZkCredentialKey(final boolean matchesCurrentZkCredentialKey) {
 
     final byte[] publicKey = TestRandomUtil.nextBytes(33);
+    final long rotationId = ThreadLocalRandom.current().nextLong(1, Long.MAX_VALUE);
 
     final Account account = mock(Account.class);
 
     if (matchesCurrentZkCredentialKey) {
       when(account.getZkCredentialKey()).thenReturn(publicKey);
+      when(account.getZkCredentialKeyRotationId()).thenReturn(rotationId);
     }
 
     when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI))
         .thenReturn(Optional.of(account));
 
-    assertDoesNotThrow(() ->
+    final SetZkCredentialKeyResponse response = assertDoesNotThrow(() ->
         authenticatedServiceStub().setZkCredentialKey(SetZkCredentialKeyRequest.newBuilder()
                 .setPublicKey(ByteString.copyFrom(publicKey))
             .build()));
 
+    if (matchesCurrentZkCredentialKey) {
+      assertEquals(rotationId, response.getRotationId());
+    } else {
+      assertTrue(response.getRotationId() != 0);
+    }
+
     final int updateMethodCalls = matchesCurrentZkCredentialKey ? 0 : 1;
 
     verify(accountsManager, times(updateMethodCalls)).update(eq(AUTHENTICATED_ACI), any());