Skip to content

Commit 04aa528

Browse files
eager-signaleager-signal
committed
Add Accounts.SetZkCredentialKey
1 parent 1b5c602 commit 04aa528

7 files changed

Lines changed: 114 additions & 3 deletions

File tree

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import com.google.protobuf.ByteString;
99
import java.util.ArrayList;
10+
import java.util.Arrays;
1011
import java.util.HexFormat;
1112
import java.util.List;
1213
import java.util.UUID;
@@ -34,6 +35,8 @@
3435
import org.signal.chat.account.SetRegistrationRecoveryPasswordResponse;
3536
import org.signal.chat.account.SetUsernameLinkRequest;
3637
import org.signal.chat.account.SetUsernameLinkResponse;
38+
import org.signal.chat.account.SetZkCredentialKeyRequest;
39+
import org.signal.chat.account.SetZkCredentialKeyResponse;
3740
import org.signal.chat.account.SimpleAccountsGrpc;
3841
import org.signal.chat.account.UsernameNotAvailable;
3942
import org.signal.chat.common.AccountIdentifiers;
@@ -267,6 +270,24 @@ public SetRegistrationRecoveryPasswordResponse setRegistrationRecoveryPassword(f
267270
return SetRegistrationRecoveryPasswordResponse.getDefaultInstance();
268271
}
269272

273+
@Override
274+
public SetZkCredentialKeyResponse setZkCredentialKey(final SetZkCredentialKeyRequest request) throws RateLimitExceededException {
275+
final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
276+
277+
final Account authenticatedAccount = getAuthenticatedAccount();
278+
final byte[] zkCredentialKey = request.getPublicKey().toByteArray();
279+
280+
if (Arrays.equals(authenticatedAccount.getZkCredentialKey(), zkCredentialKey)) {
281+
return SetZkCredentialKeyResponse.getDefaultInstance();
282+
}
283+
284+
rateLimiters.getSetZkCredentialKeyLimiter().validate(authenticatedDevice.accountIdentifier());
285+
286+
accountsManager.update(authenticatedAccount, account -> account.setZkCredentialKey(zkCredentialKey));
287+
288+
return SetZkCredentialKeyResponse.getDefaultInstance();
289+
}
290+
270291
private Account getAuthenticatedAccount() {
271292
return getAuthenticatedAccount(AuthenticationUtil.requireAuthenticatedDevice());
272293
}

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public enum For implements RateLimiterDescriptor {
3636
USERNAME_LINK_LOOKUP_PER_IP("usernameLinkLookupPerIp", new RateLimiterConfig(100, Duration.ofSeconds(15), true)),
3737
CHECK_ACCOUNT_EXISTENCE("checkAccountExistence", new RateLimiterConfig(1000, Duration.ofSeconds(4), true)),
3838
REGISTRATION("registration", new RateLimiterConfig(6, Duration.ofSeconds(30), false)),
39+
SET_ZK_CREDENTIAL_KEY("setZkCredentialKey", new RateLimiterConfig(5, Duration.ofDays(7), false)),
3940
VERIFICATION_PUSH_CHALLENGE("verificationPushChallenge", new RateLimiterConfig(5, Duration.ofSeconds(30), false)),
4041
VERIFICATION_CAPTCHA("verificationCaptcha", new RateLimiterConfig(10, Duration.ofSeconds(30), false)),
4142
RATE_LIMIT_RESET("rateLimitReset", new RateLimiterConfig(2, Duration.ofHours(12), false)),
@@ -106,8 +107,8 @@ public RateLimiter getVerifyDeviceLimiter() {
106107
return forDescriptor(For.VERIFY_DEVICE);
107108
}
108109

109-
public RateLimiter getMessagesLimiter() {
110-
return forDescriptor(For.MESSAGES);
110+
public RateLimiter getSetZkCredentialKeyLimiter() {
111+
return forDescriptor(For.SET_ZK_CREDENTIAL_KEY);
111112
}
112113

113114
public RateLimiter getPreKeysLimiter() {
@@ -162,6 +163,10 @@ public RateLimiter getRegistrationLimiter() {
162163
return forDescriptor(For.REGISTRATION);
163164
}
164165

166+
public RateLimiter getMessagesLimiter() {
167+
return forDescriptor(For.MESSAGES);
168+
}
169+
165170
public RateLimiter getRateLimitResetLimiter() {
166171
return forDescriptor(For.RATE_LIMIT_RESET);
167172
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,10 @@ public class Account {
112112
@Nullable
113113
private BackupVoucher backupVoucher;
114114

115+
@JsonProperty("zck")
116+
@Nullable
117+
private byte[] zkCredentialKey;
118+
115119
@JsonProperty
116120
private int version;
117121

@@ -536,6 +540,15 @@ public void setUsernameHolds(final List<UsernameHold> usernameHolds) {
536540
this.usernameHolds = usernameHolds;
537541
}
538542

543+
@Nullable
544+
public byte[] getZkCredentialKey() {
545+
return zkCredentialKey;
546+
}
547+
548+
public void setZkCredentialKey(@Nullable final byte[] zkCredentialKey) {
549+
this.zkCredentialKey = zkCredentialKey;
550+
}
551+
539552
public void markStale() {
540553
stale = true;
541554
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,9 @@ CompletionStage<Void> reclaimAccount(final Account existingAccount,
326326
// Carry over the existing backup voucher to the new account
327327
accountToCreate.setBackupVoucher(existingAccount.getBackupVoucher());
328328

329+
// Carry over the existing ZK credential key to the new account
330+
accountToCreate.setZkCredentialKey(existingAccount.getZkCredentialKey());
331+
329332
final List<TransactWriteItem> writeItems = new ArrayList<>();
330333

331334
// 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‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ service Accounts {
5757

5858
// Sets the registration recovery password for the authenticated account.
5959
rpc SetRegistrationRecoveryPassword(SetRegistrationRecoveryPasswordRequest) returns (SetRegistrationRecoveryPasswordResponse) {}
60+
61+
// Store a public key used to issue and verify zero-knowledge (anonymous) credentials for the account.
62+
rpc SetZkCredentialKey(SetZkCredentialKeyRequest) returns (SetZkCredentialKeyResponse) {}
6063
}
6164

6265
// Provides methods for looking up Signal accounts. Callers must not provide
@@ -265,3 +268,11 @@ message LookupUsernameLinkResponse {
265268
errors.NotFound not_found = 2 [(tag.reason) = "not_found"];
266269
}
267270
}
271+
272+
message SetZkCredentialKeyRequest {
273+
// A serialized Ristretto key with a one-byte type prefix
274+
bytes public_key = 1 [(require.exactlySize) = 33];
275+
}
276+
277+
message SetZkCredentialKeyResponse {
278+
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,13 @@
88
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
99
import static org.junit.jupiter.api.Assertions.assertEquals;
1010
import static org.junit.jupiter.api.Assertions.assertTrue;
11+
import static org.mockito.AdditionalMatchers.aryEq;
1112
import static org.mockito.ArgumentMatchers.any;
1213
import static org.mockito.ArgumentMatchers.eq;
1314
import static org.mockito.Mockito.doThrow;
1415
import static org.mockito.Mockito.mock;
1516
import static org.mockito.Mockito.never;
17+
import static org.mockito.Mockito.times;
1618
import static org.mockito.Mockito.verify;
1719
import static org.mockito.Mockito.when;
1820

@@ -54,6 +56,7 @@
5456
import org.signal.chat.account.SetRegistrationRecoveryPasswordRequest;
5557
import org.signal.chat.account.SetUsernameLinkRequest;
5658
import org.signal.chat.account.SetUsernameLinkResponse;
59+
import org.signal.chat.account.SetZkCredentialKeyRequest;
5760
import org.signal.chat.account.UsernameNotAvailable;
5861
import org.signal.chat.common.AccountIdentifiers;
5962
import org.signal.chat.errors.FailedPrecondition;
@@ -109,6 +112,8 @@ protected AccountsGrpcService createServiceBeforeEachTest() {
109112
when(rateLimiters.getUsernameSetLimiter()).thenReturn(rateLimiter);
110113
when(rateLimiters.getUsernameLinkOperationLimiter()).thenReturn(rateLimiter);
111114

115+
when(rateLimiters.getSetZkCredentialKeyLimiter()).thenReturn(rateLimiter);
116+
112117
when(registrationRecoveryPasswordsManager.store(any(), any()))
113118
.thenReturn(CompletableFuture.completedFuture(null));
114119

@@ -712,4 +717,53 @@ void setRegistrationRecoveryPasswordMissingPassword() {
712717
() -> authenticatedServiceStub().setRegistrationRecoveryPassword(
713718
SetRegistrationRecoveryPasswordRequest.newBuilder().build()));
714719
}
720+
721+
@ParameterizedTest
722+
@ValueSource(booleans = {true, false})
723+
void setZkCredentialKey(final boolean matchesCurrentZkCredentialKey) {
724+
725+
final byte[] publicKey = TestRandomUtil.nextBytes(33);
726+
727+
final Account account = mock(Account.class);
728+
729+
if (matchesCurrentZkCredentialKey) {
730+
when(account.getZkCredentialKey()).thenReturn(publicKey);
731+
}
732+
733+
when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI))
734+
.thenReturn(Optional.of(account));
735+
736+
assertDoesNotThrow(() ->
737+
authenticatedServiceStub().setZkCredentialKey(SetZkCredentialKeyRequest.newBuilder()
738+
.setPublicKey(ByteString.copyFrom(publicKey))
739+
.build()));
740+
741+
final int updateMethodCalls = matchesCurrentZkCredentialKey ? 0 : 1;
742+
743+
verify(accountsManager, times(updateMethodCalls)).update(eq(account), any(Consumer.class));
744+
verify(account, times(updateMethodCalls)).setZkCredentialKey(aryEq(publicKey));
745+
}
746+
747+
@Test
748+
void setZkCredentialKeyRateLimited() throws Exception {
749+
750+
final byte[] publicKey = TestRandomUtil.nextBytes(33);
751+
final Duration retryDuration = Duration.ofDays(1);
752+
753+
final Account account = mock(Account.class);
754+
when(account.getUuid()).thenReturn(AUTHENTICATED_ACI);
755+
756+
when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI))
757+
.thenReturn(Optional.of(account));
758+
doThrow(new RateLimitExceededException(retryDuration))
759+
.when(rateLimiter).validate(AUTHENTICATED_ACI);
760+
761+
GrpcTestUtils.assertRateLimitExceeded(retryDuration, () ->
762+
authenticatedServiceStub().setZkCredentialKey(SetZkCredentialKeyRequest.newBuilder()
763+
.setPublicKey(ByteString.copyFrom(publicKey))
764+
.build()));
765+
766+
verify(accountsManager, never()).update(any(Account.class), any(Consumer.class));
767+
verify(account, never()).setZkCredentialKey(any());
768+
}
715769
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,7 @@ void testReclaimAccountPreservesFields() {
479479

480480
// the backup credential request and share-set are always preserved across account reclaims
481481
existingAccount.setBackupCredentialRequests(TestRandomUtil.nextBytes(32), TestRandomUtil.nextBytes(32));
482+
existingAccount.setZkCredentialKey(TestRandomUtil.nextBytes(32));
482483
createAccount(existingAccount);
483484
final Account secondAccount =
484485
generateAccount(e164, UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
@@ -490,6 +491,7 @@ void testReclaimAccountPreservesFields() {
490491
.isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MESSAGES).orElseThrow());
491492
assertThat(reclaimed.getBackupCredentialRequest(BackupCredentialType.MEDIA).orElseThrow())
492493
.isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MEDIA).orElseThrow());
494+
assertThat(reclaimed.getZkCredentialKey()).isEqualTo(existingAccount.getZkCredentialKey());
493495
}
494496

495497
@Test
@@ -500,9 +502,11 @@ void testReclaimAccount() throws UsernameHashNotAvailableException {
500502
final UUID existingPni = UUID.randomUUID();
501503
final Account existingAccount = generateAccount(e164, existingUuid, existingPni, List.of(device));
502504

503-
// Backup vouchers should be carried over accross re-registration
505+
// Backup vouchers should be carried over across re-registration
504506
final Account.BackupVoucher bv = new Account.BackupVoucher(1, Instant.now().plus(Duration.ofDays(1)));
505507
existingAccount.setBackupVoucher(bv);
508+
// ZK credential keys should be carried over across re-registration
509+
existingAccount.setZkCredentialKey(TestRandomUtil.nextBytes(32));
506510

507511
createAccount(existingAccount);
508512

0 commit comments

Comments
 (0)