From 2b3a6e1aa502e2a79b68b2bc9e6e2834c1a6aa06 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 17 Aug 2026 20:20:51 -0400 Subject: [PATCH] Add gRPC endpoints for working with TOTP keys --- .../grpc/AccountsGrpcService.java | 123 ++++++++++- .../storage/AccountsManager.java | 2 +- .../storage/AnnotatedTotpKey.java | 16 ++ .../textsecuregcm/storage/TotpKey.java | 18 +- .../main/proto/org/signal/chat/account.proto | 125 +++++++++++ .../grpc/AccountsGrpcServiceTest.java | 195 ++++++++++++++++++ 6 files changed, 475 insertions(+), 4 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java index a15e3b33f..cddede3e7 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java @@ -14,10 +14,12 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.HexFormat; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.UUID; import java.util.function.Function; import java.util.stream.Collectors; @@ -28,6 +30,8 @@ import org.signal.chat.account.ClearRegistrationLockRequest; import org.signal.chat.account.ClearRegistrationLockResponse; import org.signal.chat.account.ConfigureUnidentifiedAccessRequest; import org.signal.chat.account.ConfigureUnidentifiedAccessResponse; +import org.signal.chat.account.ConfirmTotpKeyRequest; +import org.signal.chat.account.ConfirmTotpKeyResponse; import org.signal.chat.account.ConfirmUsernameHashRequest; import org.signal.chat.account.ConfirmUsernameHashResponse; import org.signal.chat.account.DeleteAccountRequest; @@ -37,6 +41,8 @@ import org.signal.chat.account.DeleteUsernameHashResponse; import org.signal.chat.account.DeleteUsernameLinkRequest; import org.signal.chat.account.DeleteUsernameLinkResponse; import org.signal.chat.account.ExternalServiceCredentials; +import org.signal.chat.account.GenerateTotpKeyRequest; +import org.signal.chat.account.GenerateTotpKeyResponse; import org.signal.chat.account.GetAccountDataReportRequest; import org.signal.chat.account.GetAccountDataReportResponse; import org.signal.chat.account.GetAccountIdentityRequest; @@ -45,7 +51,11 @@ import org.signal.chat.account.GetCapabilitiesRequest; import org.signal.chat.account.GetCapabilitiesResponse; import org.signal.chat.account.GetEntitlementsRequest; import org.signal.chat.account.GetEntitlementsResponse; +import org.signal.chat.account.ListTotpKeysRequest; +import org.signal.chat.account.ListTotpKeysResponse; import org.signal.chat.account.RegistrationLockFailure; +import org.signal.chat.account.RemoveTotpKeyRequest; +import org.signal.chat.account.RemoveTotpKeyResponse; import org.signal.chat.account.ReserveUsernameHashRequest; import org.signal.chat.account.ReserveUsernameHashResponse; import org.signal.chat.account.SetDiscoverableByPhoneNumberRequest; @@ -54,15 +64,19 @@ import org.signal.chat.account.SetRegistrationLockRequest; import org.signal.chat.account.SetRegistrationLockResponse; import org.signal.chat.account.SetRegistrationRecoveryPasswordRequest; import org.signal.chat.account.SetRegistrationRecoveryPasswordResponse; +import org.signal.chat.account.SetTotpKeyMetadataRequest; +import org.signal.chat.account.SetTotpKeyMetadataResponse; 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.StaleDevices; +import org.signal.chat.account.TotpParameters; import org.signal.chat.account.UsernameNotAvailable; import org.signal.chat.common.AccountIdentifiers; import org.signal.chat.errors.FailedPrecondition; +import org.signal.chat.errors.NotFound; import org.signal.chat.messages.SendMessageType; import org.signal.libsignal.protocol.IdentityKey; import org.signal.libsignal.protocol.InvalidKeyException; @@ -90,16 +104,21 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters; import org.whispersystems.textsecuregcm.push.MessageTooLargeException; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.AnnotatedTotpKey; import org.whispersystems.textsecuregcm.storage.ChangeNumberManager; import org.whispersystems.textsecuregcm.storage.DeviceCapability; import org.whispersystems.textsecuregcm.storage.PhoneNumberRecoveryPasswordsManager; +import org.whispersystems.textsecuregcm.storage.TooManyTotpKeysException; +import org.whispersystems.textsecuregcm.storage.TotpKey; import org.whispersystems.textsecuregcm.storage.UsernameHashNotAvailableException; import org.whispersystems.textsecuregcm.storage.UsernameReservationNotFoundException; +import org.whispersystems.textsecuregcm.util.NoStackTraceRuntimeException; import org.whispersystems.textsecuregcm.util.RegistrationIdValidator; import org.whispersystems.textsecuregcm.util.SystemMapper; import org.whispersystems.textsecuregcm.util.UUIDUtil; import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; +import javax.annotation.Nullable; public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase { @@ -112,6 +131,9 @@ public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase { private final Clock clock; private final ChangeNumberManager changeNumberManager; + private static class TotpKeyNotFoundException extends NoStackTraceRuntimeException { + } + public AccountsGrpcService(final AccountsManager accountsManager, final RateLimiters rateLimiters, final UsernameHashZkProofVerifier usernameHashZkProofVerifier, @@ -534,6 +556,105 @@ public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase { return GetCapabilitiesResponse.newBuilder().setCapabilities(builder.build()).build(); } + @Override + public GenerateTotpKeyResponse generateTotpKey(final GenerateTotpKeyRequest request) { + try { + final TotpKey pendingTotpKey = + accountsManager.generatePendingTotpKey(AuthenticationUtil.requireAuthenticatedDevice().accountIdentifier()); + + return GenerateTotpKeyResponse.newBuilder() + .setKeyGenerated(GenerateTotpKeyResponse.KeyGenerated.newBuilder() + .setKey(ByteString.copyFrom(pendingTotpKey.getEncoded())) + .setTotpParameters(toGrpcTotpParameters(pendingTotpKey.totpParameters())) + .build()) + .build(); + } catch (final TooManyTotpKeysException e) { + return GenerateTotpKeyResponse.newBuilder() + .setTooManyTotpKeys(FailedPrecondition.getDefaultInstance()) + .build(); + } + } + + @Override + public ConfirmTotpKeyResponse confirmTotpKey(final ConfirmTotpKeyRequest request) { + final Optional maybeConfirmedTotpKeyId = + accountsManager.confirmPendingTotpKey(AuthenticationUtil.requireAuthenticatedDevice().accountIdentifier(), + request.getOneTimePassword(), + clock.instant(), + request.getMetadataCiphertext().toByteArray()); + + return maybeConfirmedTotpKeyId.map(keyId -> ConfirmTotpKeyResponse.newBuilder() + .setKeyConfirmed(ConfirmTotpKeyResponse.KeyConfirmed.newBuilder() + .setKeyId(keyId) + .build()) + .build()).orElseGet(() -> ConfirmTotpKeyResponse.newBuilder() + .setOneTimePasswordNotVerified(FailedPrecondition.getDefaultInstance()) + .build()); + } + + @Override + public ListTotpKeysResponse listTotpKeys(final ListTotpKeysRequest request) { + final ListTotpKeysResponse.Builder responseBuilder = ListTotpKeysResponse.newBuilder(); + + getAuthenticatedAccount().getTotpKeys().forEach((keyId, totpKey) -> { + assert totpKey.metadataCiphertext() != null; + + responseBuilder.putKeys(keyId, ListTotpKeysResponse.TotpKeyMetadata.newBuilder() + .setMetadataCiphertext(ByteString.copyFrom(totpKey.metadataCiphertext())) + .setTotpParameters(toGrpcTotpParameters(totpKey.totpKey().totpParameters())) + .build()); + }); + + return responseBuilder.build(); + } + + @Override + public SetTotpKeyMetadataResponse setTotpKeyMetadata(final SetTotpKeyMetadataRequest request) { + try { + accountsManager.update(AuthenticationUtil.requireAuthenticatedDevice().accountIdentifier(), account -> { + final Map totpKeys = new HashMap<>(account.getTotpKeys()); + @Nullable final AnnotatedTotpKey existingKey = totpKeys.get(request.getKeyId()); + + if (existingKey == null) { + throw new TotpKeyNotFoundException(); + } + + totpKeys.put(request.getKeyId(), + new AnnotatedTotpKey(existingKey.totpKey(), request.getMetadataCiphertext().toByteArray())); + + account.setTotpKeys(totpKeys); + }); + + return SetTotpKeyMetadataResponse.newBuilder() + .setMetadataUpdated(SetTotpKeyMetadataResponse.MetadataUpdated.getDefaultInstance()) + .build(); + } catch (final TotpKeyNotFoundException _) { + return SetTotpKeyMetadataResponse.newBuilder() + .setKeyNotFound(NotFound.getDefaultInstance()) + .build(); + } + } + + @Override + public RemoveTotpKeyResponse removeTotpKey(final RemoveTotpKeyRequest request) { + accountsManager.update(AuthenticationUtil.requireAuthenticatedDevice().accountIdentifier(), account -> { + final Map totpKeys = new HashMap<>(account.getTotpKeys()); + totpKeys.remove(request.getKeyId()); + + account.setTotpKeys(totpKeys); + }); + + return RemoveTotpKeyResponse.getDefaultInstance(); + } + + private static TotpParameters toGrpcTotpParameters(final org.whispersystems.textsecuregcm.storage.TotpParameters totpParameters) { + return TotpParameters.newBuilder() + .setAlgorithm(totpParameters.algorithm()) + .setPasswordLength(totpParameters.passwordLength()) + .setTimeStepSeconds(Math.toIntExact(totpParameters.timeStep().toSeconds())) + .build(); + } + private static AccountIdentifiers buildAccountIdentifiers(final Account account) { final AccountIdentifiers.Builder accountIdentifiersBuilder = AccountIdentifiers.newBuilder() .addServiceIdentifiers(GrpcServiceIdentifierUtil.toGrpcServiceIdentifier(new AciServiceIdentifier(account.getAccountIdentifier()))); @@ -584,7 +705,7 @@ public class AccountsGrpcService extends SimpleAccountsGrpc.AccountsImplBase { } private static StatusRuntimeException invalidSignatureException(final String fieldName) { - return GrpcExceptions.fieldViolation(fieldName, "pre-key signature did not match PNI identity key"); + return GrpcExceptions.fieldViolation(fieldName, "pre-key signature did not match PNI identity key"); } private static Map transformDeviceMap(final Map byDeviceId, final Function f) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index 7b7737296..3d22b6237 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -208,7 +208,7 @@ public class AccountsManager extends RedisPubSubAdapter implemen static final int TOTP_KEY_LENGTH_BITS = 256; @VisibleForTesting - static final TotpParameters TOTP_PARAMETERS = new TotpParameters( + public static final TotpParameters TOTP_PARAMETERS = new TotpParameters( TimeBasedOneTimePasswordGenerator.TOTP_ALGORITHM_HMAC_SHA256, HmacOneTimePasswordGenerator.DEFAULT_PASSWORD_LENGTH, TimeBasedOneTimePasswordGenerator.DEFAULT_TIME_STEP); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java index d960ad769..39e211fe4 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java @@ -7,6 +7,8 @@ package org.whispersystems.textsecuregcm.storage; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonUnwrapped; +import java.util.Arrays; +import java.util.Objects; import javax.crypto.SecretKey; public record AnnotatedTotpKey(@JsonUnwrapped @@ -29,4 +31,18 @@ public record AnnotatedTotpKey(@JsonUnwrapped public byte[] getEncoded() { return totpKey().getEncoded(); } + + @Override + public boolean equals(final Object o) { + if (!(o instanceof AnnotatedTotpKey(TotpKey key, byte[] ciphertext))) { + return false; + } + + return Objects.equals(totpKey, key) && Objects.deepEquals(metadataCiphertext, ciphertext); + } + + @Override + public int hashCode() { + return Objects.hash(totpKey, Arrays.hashCode(metadataCiphertext)); + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java index 6a721493d..ad61aa391 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java @@ -7,9 +7,9 @@ package org.whispersystems.textsecuregcm.storage; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonUnwrapped; -import javax.annotation.Nullable; -import javax.crypto.SecretKey; import java.util.Arrays; +import java.util.Objects; +import javax.crypto.SecretKey; public record TotpKey(@JsonUnwrapped TotpParameters totpParameters, @@ -31,4 +31,18 @@ public record TotpKey(@JsonUnwrapped public byte[] getEncoded() { return Arrays.copyOf(encodedKey(), encodedKey().length); } + + @Override + public boolean equals(final Object o) { + if (!(o instanceof TotpKey(TotpParameters parameters, byte[] key))) { + return false; + } + + return Objects.deepEquals(encodedKey, key) && Objects.equals(totpParameters, parameters); + } + + @Override + public int hashCode() { + return Objects.hash(totpParameters, Arrays.hashCode(encodedKey)); + } } diff --git a/service/src/main/proto/org/signal/chat/account.proto b/service/src/main/proto/org/signal/chat/account.proto index 4da06bd91..c12755b16 100644 --- a/service/src/main/proto/org/signal/chat/account.proto +++ b/service/src/main/proto/org/signal/chat/account.proto @@ -77,6 +77,25 @@ service Accounts { // Gets the capabilities for the authenticated account. rpc GetCapabilities(GetCapabilitiesRequest) returns (GetCapabilitiesResponse) {} + + // Generates and stores a pending TOTP key for the authenticated account. + // To "activate" the key, callers must call the `ConfirmTotpKey` endpoint. + rpc GenerateTotpKey(GenerateTotpKeyRequest) returns (GenerateTotpKeyResponse) {} + + // Confirms that the caller has stored and can derive one-time passwords from + // a pending TOTP key generated via `GenerateTotpKey` and stores/activates the + // key for the caller's account + rpc ConfirmTotpKey(ConfirmTotpKeyRequest) returns (ConfirmTotpKeyResponse) {} + + // Returns a list of confirmed TOTP keys for the authenticated account + rpc ListTotpKeys(ListTotpKeysRequest) returns (ListTotpKeysResponse) {} + + // Updates encrypted, user-supplied metadata (e.g. a human-readable name and + // creation timestamp) for an existing, confirmed TOTP key + rpc SetTotpKeyMetadata(SetTotpKeyMetadataRequest) returns (SetTotpKeyMetadataResponse) {} + + // Removes a TOTP from the authenticated account + rpc RemoveTotpKey(RemoveTotpKeyRequest) returns (RemoveTotpKeyResponse) {} } // Provides methods for looking up Signal accounts. Callers must not provide @@ -481,3 +500,109 @@ message GetCapabilitiesAnonymousResponse { errors.FailedUnidentifiedAuthorization failed_unidentified_authorization = 3 [(tag.reason) = "failed_unidentified_authorization"]; } } + +message TotpParameters { + // The HMAC algorithm (e.g. "HmacSHA256") used by the TOTP generator + string algorithm = 1; + + // The length of one-time passwords (in decimal digits) produced and expected + // by the TOTP generator + uint32 password_length = 2; + + // The time step (in seconds) used by the TOTP generator + uint32 time_step_seconds = 3; +} + +message GenerateTotpKeyRequest { +} + +message GenerateTotpKeyResponse { + message KeyGenerated { + // The raw TOTP key + bytes key = 1; + + // The TOTP parameters associated with the generated key + TotpParameters totp_parameters = 2; + } + + oneof response { + // A new, pending TOTP key has been generated and added to the authenticated + // account + KeyGenerated key_generated = 1; + + // The authenticated account already has too many TOTP keys, and the caller + // must remove one before adding more + errors.FailedPrecondition too_many_totp_keys = 2 [(tag.reason) = "too_many_totp_keys"]; + } +} + +message ConfirmTotpKeyRequest { + // A one-time password derived from the current pending TOTP key + uint32 one_time_password = 1; + + // The ciphertext of user-provided metadata (presumably including a + // human-readable name and creation timestamp) to be attached to the + // newly-confirmed key + bytes metadata_ciphertext = 2 [(require.exactlySize) = 160]; +} + +message ConfirmTotpKeyResponse { + message KeyConfirmed { + // The account-specific identifier for the newly-confirmed TOTP key + uint32 key_id = 1; + } + + oneof response { + // The provided one-time password was accepted and the pending TOTP key was + // stored with the provided name ciphertext + KeyConfirmed key_confirmed = 1; + + // The provided one-time password was not valid for any reason (including + // incorrect passwords, misaligned clocks, or missing account records) + errors.FailedPrecondition one_time_password_not_verified = 2 [(tag.reason) = "one_time_password_not_verified"]; + } +} + +message ListTotpKeysRequest {} + +message ListTotpKeysResponse { + message TotpKeyMetadata { + // The user-provided ciphertext for metadata associated with this TOTP key + bytes metadata_ciphertext = 1; + + // The TOTP parameters associated with this key + TotpParameters totp_parameters = 2; + } + + map keys = 1; +} + +message SetTotpKeyMetadataRequest { + // The account-specific identifier of the TOTP key to modify + uint32 key_id = 1; + + // The ciphertext of the new user-provided metadata to be attached to the + // identified key + bytes metadata_ciphertext = 2 [(require.exactlySize) = 160]; +} + +message SetTotpKeyMetadataResponse { + message MetadataUpdated { + } + + oneof response { + // New metadata was stored for the identified TOTP key + MetadataUpdated metadata_updated = 1; + + // No TOTP was found with the given ID + errors.NotFound key_not_found = 2 [(tag.reason) = "key_not_found"]; + } +} + +message RemoveTotpKeyRequest { + // The account-specific identifier of the TOTP key to remove + uint32 key_id = 1; +} + +message RemoveTotpKeyResponse { +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java index 1373f4e2d..95f850a76 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java @@ -6,11 +6,13 @@ package org.whispersystems.textsecuregcm.grpc; import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doThrow; @@ -27,12 +29,15 @@ import io.grpc.Status; import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.Collections; import java.util.HexFormat; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.ThreadLocalRandom; +import java.util.stream.Collectors; import java.util.stream.Stream; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -49,6 +54,8 @@ import org.signal.chat.account.ChangeNumberResponse; import org.signal.chat.account.ClearRegistrationLockRequest; import org.signal.chat.account.ClearRegistrationLockResponse; import org.signal.chat.account.ConfigureUnidentifiedAccessRequest; +import org.signal.chat.account.ConfirmTotpKeyRequest; +import org.signal.chat.account.ConfirmTotpKeyResponse; import org.signal.chat.account.ConfirmUsernameHashRequest; import org.signal.chat.account.ConfirmUsernameHashResponse; import org.signal.chat.account.DeleteAccountRequest; @@ -56,6 +63,8 @@ import org.signal.chat.account.DeleteAccountResponse; import org.signal.chat.account.DeleteUsernameHashRequest; import org.signal.chat.account.DeleteUsernameLinkRequest; import org.signal.chat.account.ExternalServiceCredentials; +import org.signal.chat.account.GenerateTotpKeyRequest; +import org.signal.chat.account.GenerateTotpKeyResponse; import org.signal.chat.account.GetAccountDataReportRequest; import org.signal.chat.account.GetAccountDataReportResponse; import org.signal.chat.account.GetAccountIdentityRequest; @@ -64,18 +73,24 @@ import org.signal.chat.account.GetCapabilitiesRequest; import org.signal.chat.account.GetCapabilitiesResponse; import org.signal.chat.account.GetEntitlementsRequest; import org.signal.chat.account.GetEntitlementsResponse; +import org.signal.chat.account.ListTotpKeysRequest; +import org.signal.chat.account.ListTotpKeysResponse; import org.signal.chat.account.RegistrationLockFailure; +import org.signal.chat.account.RemoveTotpKeyRequest; import org.signal.chat.account.ReserveUsernameHashRequest; import org.signal.chat.account.ReserveUsernameHashResponse; import org.signal.chat.account.SetDiscoverableByPhoneNumberRequest; import org.signal.chat.account.SetRegistrationLockRequest; import org.signal.chat.account.SetRegistrationLockResponse; import org.signal.chat.account.SetRegistrationRecoveryPasswordRequest; +import org.signal.chat.account.SetTotpKeyMetadataRequest; +import org.signal.chat.account.SetTotpKeyMetadataResponse; 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.StaleDevices; +import org.signal.chat.account.TotpParameters; import org.signal.chat.account.UsernameNotAvailable; import org.signal.chat.common.AccountIdentifiers; import org.signal.chat.common.EcSignedPreKey; @@ -108,11 +123,14 @@ import org.whispersystems.textsecuregcm.push.MessageTooLargeException; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountBadge; import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.AnnotatedTotpKey; import org.whispersystems.textsecuregcm.storage.ChangeNumberManager; import org.whispersystems.textsecuregcm.storage.Device; import org.whispersystems.textsecuregcm.storage.DeviceCapability; import org.whispersystems.textsecuregcm.storage.KeyIdUtil; import org.whispersystems.textsecuregcm.storage.PhoneNumberRecoveryPasswordsManager; +import org.whispersystems.textsecuregcm.storage.TooManyTotpKeysException; +import org.whispersystems.textsecuregcm.storage.TotpKey; import org.whispersystems.textsecuregcm.storage.UsernameHashNotAvailableException; import org.whispersystems.textsecuregcm.storage.UsernameReservationNotFoundException; import org.whispersystems.textsecuregcm.tests.util.AccountsHelper; @@ -143,6 +161,8 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest authenticatedServiceStub().confirmTotpKey(ConfirmTotpKeyRequest.newBuilder() + .setOneTimePassword(123456) + .setMetadataCiphertext(ByteString.copyFrom(TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE + 1))) + .build())); + } + + @Test + void confirmTotpKeyPasswordNotVerified() { + when(accountsManager.confirmPendingTotpKey(any(), anyInt(), any(), any())) + .thenReturn(Optional.empty()); + + final ConfirmTotpKeyResponse response = authenticatedServiceStub().confirmTotpKey(ConfirmTotpKeyRequest.newBuilder() + .setOneTimePassword(123456) + .setMetadataCiphertext(ByteString.copyFrom(TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE))) + .build()); + + assertEquals(ConfirmTotpKeyResponse.ResponseCase.ONE_TIME_PASSWORD_NOT_VERIFIED, response.getResponseCase()); + } + + @Test + void listTotpKeys() { + final Map totpKeys = Map.of( + 1, generateRandomAnnotatedTotpKey(), + 2, generateRandomAnnotatedTotpKey()); + + final Account account = mock(Account.class); + when(account.getTotpKeys()).thenReturn(totpKeys); + + when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)) + .thenReturn(Optional.of(account)); + + final ListTotpKeysResponse response = + authenticatedServiceStub().listTotpKeys(ListTotpKeysRequest.getDefaultInstance()); + + final TotpParameters expectedTotpParameters = TotpParameters.newBuilder() + .setAlgorithm(AccountsManager.TOTP_PARAMETERS.algorithm()) + .setPasswordLength(AccountsManager.TOTP_PARAMETERS.passwordLength()) + .setTimeStepSeconds(Math.toIntExact(AccountsManager.TOTP_PARAMETERS.timeStep().toSeconds())) + .build(); + + final Map expectedTotpKeys = totpKeys.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> ListTotpKeysResponse.TotpKeyMetadata.newBuilder() + .setTotpParameters(expectedTotpParameters) + .setMetadataCiphertext(ByteString.copyFrom(entry.getValue().metadataCiphertext())) + .build())); + + assertEquals(expectedTotpKeys, response.getKeysMap()); + } + + @Test + void setTotpKeyMetadata() { + final int keyId = ThreadLocalRandom.current().nextInt(); + final byte[] updatedMetadata = TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE); + final AnnotatedTotpKey existingTotpKey = generateRandomAnnotatedTotpKey(); + + final Account account = mock(Account.class); + when(account.getTotpKeys()).thenReturn(Map.of(keyId, existingTotpKey)); + + when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)) + .thenReturn(Optional.of(account)); + + final SetTotpKeyMetadataResponse response = + authenticatedServiceStub().setTotpKeyMetadata(SetTotpKeyMetadataRequest.newBuilder() + .setKeyId(keyId) + .setMetadataCiphertext(ByteString.copyFrom(updatedMetadata)) + .build()); + + assertEquals(SetTotpKeyMetadataResponse.ResponseCase.METADATA_UPDATED, response.getResponseCase()); + verify(account).setTotpKeys(Map.of(keyId, new AnnotatedTotpKey(existingTotpKey.totpKey(), updatedMetadata))); + } + + @Test + void setTotpKeyMetadataKeyNotFound() { + final Account account = mock(Account.class); + when(account.getTotpKeys()).thenReturn(Collections.emptyMap()); + + when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)) + .thenReturn(Optional.of(account)); + + final SetTotpKeyMetadataResponse response = + authenticatedServiceStub().setTotpKeyMetadata(SetTotpKeyMetadataRequest.newBuilder() + .setKeyId(ThreadLocalRandom.current().nextInt()) + .setMetadataCiphertext(ByteString.copyFrom(TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE))) + .build()); + + assertEquals(SetTotpKeyMetadataResponse.ResponseCase.KEY_NOT_FOUND, response.getResponseCase()); + } + + @Test + void setTotpKeyMetadataIncorrectMetadataSize() { + GrpcTestUtils.assertStatusInvalidArgument( + () -> authenticatedServiceStub().setTotpKeyMetadata(SetTotpKeyMetadataRequest.newBuilder() + .setKeyId(1) + .setMetadataCiphertext(ByteString.copyFrom(TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE + 1))) + .build())); + } + + @Test + void removeTotpKey() { + final int retainedKeyId = 1; + final int removedKeyId = 2; + + final Map initialKeys = Map.of( + retainedKeyId, generateRandomAnnotatedTotpKey(), + removedKeyId, generateRandomAnnotatedTotpKey()); + + final Account account = mock(Account.class); + when(account.getTotpKeys()).thenReturn(initialKeys); + + when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)) + .thenReturn(Optional.of(account)); + + //noinspection ResultOfMethodCallIgnored + authenticatedServiceStub().removeTotpKey(RemoveTotpKeyRequest.newBuilder() + .setKeyId(removedKeyId) + .build()); + + verify(account).setTotpKeys(Map.of(retainedKeyId, initialKeys.get(retainedKeyId))); + } + + private static AnnotatedTotpKey generateRandomAnnotatedTotpKey() { + return new AnnotatedTotpKey(new TotpKey(AccountsManager.TOTP_PARAMETERS, TestRandomUtil.nextBytes(32)), + TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE)); + } }