Add gRPC endpoints for working with TOTP keys

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<Integer> 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<Integer, AnnotatedTotpKey> 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<Integer, AnnotatedTotpKey> 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 <T, U> Map<Byte, U> transformDeviceMap(final Map<Integer, T> byDeviceId, final Function<T, U> f) {

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

@@ -208,7 +208,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> 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);

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));
+  }
 }

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));
+  }
 }

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<int32, TotpKeyMetadata> 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 {
+}

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<AccountsGrpcService, Ac
   @Mock
   private ChangeNumberManager changeNumberManager;
 
+  private static final int TOTP_KEY_METADATA_SIZE = 160;
+
   @Override
   protected AccountsGrpcService createServiceBeforeEachTest() {
     AccountsHelper.setupMockUpdate(accountsManager);
@@ -1096,4 +1116,179 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
         .addCapabilities(org.signal.chat.common.DeviceCapability.DEVICE_CAPABILITY_USERNAME_CHANGE_SYNC_MESSAGE)
         .build(), response.getCapabilities());
   }
+
+  @Test
+  void generateTotpKey() throws TooManyTotpKeysException {
+    final byte[] encodedKey = TestRandomUtil.nextBytes(16);
+
+    when(accountsManager.generatePendingTotpKey(AUTHENTICATED_ACI))
+        .thenReturn(new TotpKey(AccountsManager.TOTP_PARAMETERS, encodedKey));
+
+    final GenerateTotpKeyResponse response =
+        authenticatedServiceStub().generateTotpKey(GenerateTotpKeyRequest.getDefaultInstance());
+
+    assertEquals(GenerateTotpKeyResponse.ResponseCase.KEY_GENERATED, response.getResponseCase());
+
+    assertArrayEquals(encodedKey, response.getKeyGenerated().getKey().toByteArray());
+    assertEquals(AccountsManager.TOTP_PARAMETERS.algorithm(), response.getKeyGenerated().getTotpParameters().getAlgorithm());
+    assertEquals(AccountsManager.TOTP_PARAMETERS.passwordLength(), response.getKeyGenerated().getTotpParameters().getPasswordLength());
+    assertEquals(AccountsManager.TOTP_PARAMETERS.timeStep().toSeconds(), response.getKeyGenerated().getTotpParameters().getTimeStepSeconds());
+  }
+
+  @Test
+  void generateTotpKeyTooManyKeys() throws TooManyTotpKeysException {
+    when(accountsManager.generatePendingTotpKey(AUTHENTICATED_ACI))
+        .thenThrow(TooManyTotpKeysException.class);
+
+    final GenerateTotpKeyResponse response =
+        authenticatedServiceStub().generateTotpKey(GenerateTotpKeyRequest.getDefaultInstance());
+
+    assertEquals(GenerateTotpKeyResponse.ResponseCase.TOO_MANY_TOTP_KEYS, response.getResponseCase());
+  }
+
+  @Test
+  void confirmTotpKey() {
+    final int keyId = 17;
+    final int oneTimePassword = 123456;
+    final byte[] metadataCiphertext = TestRandomUtil.nextBytes(TOTP_KEY_METADATA_SIZE);
+
+    when(accountsManager.confirmPendingTotpKey(AUTHENTICATED_ACI, oneTimePassword, testClock.instant(), metadataCiphertext))
+        .thenReturn(Optional.of(keyId));
+
+    final ConfirmTotpKeyResponse response = authenticatedServiceStub().confirmTotpKey(ConfirmTotpKeyRequest.newBuilder()
+            .setOneTimePassword(oneTimePassword)
+            .setMetadataCiphertext(ByteString.copyFrom(metadataCiphertext))
+        .build());
+
+    assertEquals(ConfirmTotpKeyResponse.ResponseCase.KEY_CONFIRMED, response.getResponseCase());
+    assertEquals(keyId, response.getKeyConfirmed().getKeyId());
+  }
+
+  @Test
+  void confirmTotpKeyIncorrectMetadataSize() {
+    GrpcTestUtils.assertStatusInvalidArgument(
+        () -> 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<Integer, AnnotatedTotpKey> 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<Integer, ListTotpKeysResponse.TotpKeyMetadata> 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<Integer, AnnotatedTotpKey> 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));
+  }
 }