Add support for generating, storing, and managing TOTP keys

pom.xml

@@ -64,6 +64,7 @@
     <httpcore.version>4.4.16</httpcore.version>
     <httpclient.version>4.5.14</httpclient.version>
     <jackson.version>2.22.2</jackson.version>
+    <java-otp.version>1.0.0</java-otp.version>
     <junit-pioneer.version>2.3.0</junit-pioneer.version>
     <jsr305.version>3.0.2</jsr305.version>
     <kotlin.version>2.4.10</kotlin.version>
@@ -191,6 +192,11 @@
         <type>pom</type>
         <scope>import</scope>
       </dependency>
+      <dependency>
+        <groupId>com.eatthepath</groupId>
+        <artifactId>java-otp</artifactId>
+        <version>${java-otp.version}</version>
+      </dependency>
       <dependency>
         <groupId>com.eatthepath</groupId>
         <artifactId>pushy</artifactId>

service/pom.xml

@@ -436,6 +436,11 @@
       <artifactId>lettuce-core</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>com.eatthepath</groupId>
+      <artifactId>java-otp</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>com.eatthepath</groupId>
       <artifactId>pushy</artifactId>

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

@@ -24,6 +24,7 @@ import java.util.Base64;
 import java.util.Collections;
 import java.util.HexFormat;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
@@ -155,6 +156,13 @@ public class Account {
   @Nullable
   private byte[] authCredentialSalt;
 
+  @JsonProperty("pendingTotp")
+  @Nullable
+  private TotpKey pendingTotpKey;
+
+  @JsonProperty("totp")
+  private Map<Integer, AnnotatedTotpKey> totpKeys = Collections.emptyMap();
+
   @JsonIgnore
   private boolean stale;
 
@@ -650,6 +658,35 @@ public class Account {
     this.authCredentialSalt = authCredentialSalt;
   }
 
+  public void setPendingTotpKey(@Nullable final TotpKey pendingTotpKey) {
+    requireNotStale();
+    this.pendingTotpKey = pendingTotpKey;
+  }
+
+  public Optional<TotpKey> getPendingTotpKey() {
+    requireNotStale();
+    return Optional.ofNullable(pendingTotpKey);
+  }
+
+  public int getNextTotpKeyId() {
+    requireNotStale();
+
+    return totpKeys.keySet().stream()
+        .mapToInt(i -> i)
+        .max()
+        .orElse(-1) + 1;
+  }
+
+  public Map<Integer, AnnotatedTotpKey> getTotpKeys() {
+    requireNotStale();
+    return totpKeys;
+  }
+
+  public void setTotpKeys(final Map<Integer, AnnotatedTotpKey> totpKeys) {
+    requireNotStale();
+    this.totpKeys = totpKeys;
+  }
+
   public void markStale() {
     stale = true;
   }

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

@@ -8,6 +8,8 @@ package org.whispersystems.textsecuregcm.storage;
 import static java.util.Objects.requireNonNull;
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
+import com.eatthepath.otp.HmacOneTimePasswordGenerator;
+import com.eatthepath.otp.TimeBasedOneTimePasswordGenerator;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectWriter;
 import com.google.common.annotations.VisibleForTesting;
@@ -38,6 +40,7 @@ import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -51,6 +54,7 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -58,7 +62,9 @@ import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import javax.crypto.KeyGenerator;
 import javax.crypto.Mac;
+import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
 import org.apache.commons.lang3.StringUtils;
 import org.signal.libsignal.protocol.IdentityKey;
@@ -82,6 +88,7 @@ import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecoveryClient;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
+import org.whispersystems.textsecuregcm.util.NoStackTraceRuntimeException;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.RegistrationIdValidator;
 import org.whispersystems.textsecuregcm.util.ResilienceUtil;
@@ -141,6 +148,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   private final ScheduledExecutorService retryExecutor;
   private final Clock clock;
 
+  private final KeyGenerator totpKeyGenerator;
+  private final TimeBasedOneTimePasswordGenerator totpGenerator;
+
   private final Key verificationTokenKey;
 
   private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
@@ -194,6 +204,17 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   @VisibleForTesting
   static final String LINK_DEVICE_VERIFICATION_TOKEN_ALGORITHM = "HmacSHA256";
 
+  @VisibleForTesting
+  static final int TOTP_KEY_LENGTH_BITS = 256;
+
+  @VisibleForTesting
+  static final TotpParameters TOTP_PARAMETERS = new TotpParameters(
+      TimeBasedOneTimePasswordGenerator.TOTP_ALGORITHM_HMAC_SHA256,
+      HmacOneTimePasswordGenerator.DEFAULT_PASSWORD_LENGTH,
+      TimeBasedOneTimePasswordGenerator.DEFAULT_TIME_STEP);
+
+  public static final int MAX_TOTP_KEYS = 2;
+
   public enum DeletionReason {
     ADMIN_DELETED("admin"),
     EXPIRED      ("expired"),
@@ -262,6 +283,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
                                   int registrationId) {
   }
 
+  private static class UncheckedTooManyTotpKeysException extends NoStackTraceRuntimeException {
+  }
+
   public AccountsManager(final Accounts accounts,
       final PhoneNumberIdentifiers phoneNumberIdentifiers,
       final FaultTolerantRedisClusterClient cacheCluster,
@@ -305,6 +329,21 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       throw new IllegalArgumentException(e);
     }
 
+    try {
+      this.totpKeyGenerator = KeyGenerator.getInstance(TOTP_PARAMETERS.algorithm());
+      totpKeyGenerator.init(TOTP_KEY_LENGTH_BITS);
+    } catch (final NoSuchAlgorithmException e) {
+      throw new AssertionError("Every implementation of the Java platform is required to support the HmacSHA256 KeyGenerator algorithm", e);
+    }
+
+    try {
+      this.totpGenerator = new TimeBasedOneTimePasswordGenerator(TOTP_PARAMETERS.timeStep(),
+          TOTP_PARAMETERS.passwordLength(),
+          TOTP_PARAMETERS.algorithm());
+    } catch (final NoSuchAlgorithmException e) {
+      throw new AssertionError("Every implementation of the Java platform is required to support the HmacSHA256 MAC algorithm", e);
+    }
+
     this.pubSubConnection = pubSubRedisClient.createPubSubConnection();
   }
 
@@ -1968,4 +2007,107 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       throw e;
     }
   }
+
+  /// Generates and stores a pending TOTP key for the identified account. Accounts may have at most one pending TOTP
+  /// key.
+  ///
+  /// @param accountIdentifier the identifier of the account for which to generate and store a pending TOTP key
+  ///
+  /// @return the generated pending TOTP key
+  ///
+  /// @see [#confirmPendingTotpKey(UUID, int, Instant, byte[])
+  ///
+  /// @throws TooManyTotpKeysException if the target account already has at least [#MAX_TOTP_KEYS] TOTP keys
+  public TotpKey generatePendingTotpKey(final UUID accountIdentifier) throws TooManyTotpKeysException {
+    final SecretKey secretKey = totpKeyGenerator.generateKey();
+    final TotpKey pendingTotpKey = new TotpKey(TOTP_PARAMETERS, secretKey.getEncoded());
+
+    try {
+      update(accountIdentifier, account -> {
+        if (account.getTotpKeys().size() >= MAX_TOTP_KEYS) {
+          throw new UncheckedTooManyTotpKeysException();
+        }
+
+        account.setPendingTotpKey(pendingTotpKey);
+      });
+    } catch (final UncheckedTooManyTotpKeysException _) {
+      throw new TooManyTotpKeysException();
+    }
+
+    return pendingTotpKey;
+  }
+
+  /// Verifies that a caller has stored a copy of their pending TOTP key and can use it to generate one-time passwords,
+  /// then stores the key to the caller's account record.
+  ///
+  /// @param accountIdentifier the identifier of the account for which to confirm a pending TOTP key
+  /// @param oneTimePassword the one-time password the caller derived from the pending TOTP key
+  /// @param timestamp the time at which the user submitted the one-time password
+  ///
+  /// @return the account-specific ID for the confirmed key if the given one-time password is valid for either a pending
+  /// TOTP password for the given account or for a one-time password previously verified for the given account or empty
+  /// otherwise
+  ///
+  /// @see [#generatePendingTotpKey(UUID)
+  public Optional<Integer> confirmPendingTotpKey(final UUID accountIdentifier,
+      final int oneTimePassword,
+      final Instant timestamp,
+      final byte[] metadataCiphertext) {
+
+    final Optional<Account> maybeAccount = accounts.getByAccountIdentifier(accountIdentifier);
+
+    if (maybeAccount.isEmpty()) {
+      return Optional.empty();
+    }
+
+    final Optional<TotpKey> maybePendingTotpKey = maybeAccount.flatMap(Account::getPendingTotpKey);
+
+    if (maybePendingTotpKey.isPresent()) {
+      final TotpKey pendingTotpKey = maybePendingTotpKey.get();
+
+      try {
+        if (totpGenerator.validateOneTimePassword(pendingTotpKey, timestamp, oneTimePassword)) {
+          final AtomicInteger keyId = new AtomicInteger();
+
+          update(accountIdentifier, account -> {
+            final Map<Integer, AnnotatedTotpKey> updatedTotpKeys = new HashMap<>(account.getTotpKeys());
+
+            keyId.set(account.getNextTotpKeyId());
+
+            updatedTotpKeys.put(keyId.get(),
+                new AnnotatedTotpKey(new TotpKey(pendingTotpKey.totpParameters(), pendingTotpKey.encodedKey()), metadataCiphertext));
+
+            account.setPendingTotpKey(null);
+            account.setTotpKeys(updatedTotpKeys);
+          });
+
+          return Optional.of(keyId.get());
+        }
+      } catch (final InvalidKeyException e) {
+        ImpossibleEvents.logImpossible(logger, "Invalid pending TOTP key for account {}", accountIdentifier, e);
+      }
+    }
+
+    // Either there was no pending TOTP password for the given account identifier or the given one-time password
+    // wasn't valid for the pending key. Either way, see if it's a valid one-time password for a key stored on the
+    // account record in case the caller is retrying a dropped request (in which case we've stored a previously
+    // pending key on the account record).
+    //
+    // It's possible (though unlikely) that more than one key will produce the same one-time password at a given
+    // instant. To compensate, we just check the key with the highest ID (i.e. the most recent). It's also theoretically
+    // possible that a user will have iterated through so many keys that they've wrapped around into negative integers,
+    // but that's not really a practical concern.
+    return getByAccountIdentifier(accountIdentifier)
+        .flatMap(account -> account.getTotpKeys().entrySet().stream()
+            .max(Map.Entry.comparingByKey())
+            .filter(entry -> {
+              try {
+                return totpGenerator.validateOneTimePassword(entry.getValue(), timestamp, oneTimePassword);
+              } catch (final InvalidKeyException e) {
+                ImpossibleEvents.logImpossible(logger, "Invalid TOTP key for account {}", accountIdentifier, e);
+                return false;
+              }
+            })
+            .map(Map.Entry::getKey));
+  }
 }

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonUnwrapped;
+import javax.crypto.SecretKey;
+
+public record AnnotatedTotpKey(@JsonUnwrapped
+                               TotpKey totpKey,
+
+                               @JsonProperty("metadata")
+                               byte[] metadataCiphertext) implements SecretKey {
+
+  @Override
+  public String getAlgorithm() {
+    return totpKey().getAlgorithm();
+  }
+
+  @Override
+  public String getFormat() {
+    return totpKey.getFormat();
+  }
+
+  @Override
+  public byte[] getEncoded() {
+    return totpKey().getEncoded();
+  }
+}

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

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import org.whispersystems.textsecuregcm.util.NoStackTraceException;
+
+public class TooManyTotpKeysException extends NoStackTraceException {
+}

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+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;
+
+public record TotpKey(@JsonUnwrapped
+                      TotpParameters totpParameters,
+
+                      @JsonProperty("key")
+                      byte[] encodedKey) implements SecretKey {
+
+  @Override
+  public String getAlgorithm() {
+    return totpParameters().algorithm();
+  }
+
+  @Override
+  public String getFormat() {
+    return "RAW";
+  }
+
+  @Override
+  public byte[] getEncoded() {
+    return Arrays.copyOf(encodedKey(), encodedKey().length);
+  }
+}

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.time.Duration;
+
+public record TotpParameters(@JsonProperty("alg")
+                             String algorithm,
+
+                             @JsonProperty("len")
+                             int passwordLength,
+
+                             @JsonProperty("step")
+                             Duration timeStep) {
+}

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

@@ -16,6 +16,8 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.whispersystems.textsecuregcm.tests.util.DevicesHelper.createDevice;
 
+import com.eatthepath.otp.HmacOneTimePasswordGenerator;
+import com.eatthepath.otp.TimeBasedOneTimePasswordGenerator;
 import com.fasterxml.jackson.annotation.JsonFilter;
 import java.lang.annotation.Annotation;
 import java.nio.charset.StandardCharsets;
@@ -25,6 +27,7 @@ import java.util.Base64;
 import java.util.Collections;
 import java.util.HexFormat;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
@@ -278,4 +281,22 @@ class AccountTest {
       final Account deserializedBase64Cpv = SystemMapper.jsonMapper().readValue(String.format(jsonTemplate, Base64.getEncoder().encodeToString(version)), Account.class);
       assertThat(deserializedBase64Cpv.getCurrentProfileVersion()).isPresent().hasValue(version);
   }
+
+  @Test
+  void getNextTotpKeyId() {
+    assertEquals(0, new Account().getNextTotpKeyId());
+
+    final AnnotatedTotpKey totpKey = new AnnotatedTotpKey(new TotpKey(
+        new TotpParameters(
+            TimeBasedOneTimePasswordGenerator.TOTP_ALGORITHM_HMAC_SHA256,
+            HmacOneTimePasswordGenerator.DEFAULT_PASSWORD_LENGTH,
+            TimeBasedOneTimePasswordGenerator.DEFAULT_TIME_STEP),
+        TestRandomUtil.nextBytes(16)),
+        TestRandomUtil.nextBytes(16));
+
+    final Account accountWithTotpKey = new Account();
+    accountWithTotpKey.setTotpKeys(Map.of(0, totpKey));
+
+    assertEquals(1, accountWithTotpKey.getNextTotpKeyId());
+  }
 }

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

@@ -33,6 +33,7 @@ import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import com.eatthepath.otp.TimeBasedOneTimePasswordGenerator;
 import com.google.i18n.phonenumbers.PhoneNumberUtil;
 import io.lettuce.core.RedisException;
 import io.lettuce.core.api.async.RedisAsyncCommands;
@@ -41,6 +42,7 @@ import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -59,13 +61,18 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
+import javax.crypto.KeyGenerator;
 import javax.crypto.spec.SecretKeySpec;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.function.Executable;
@@ -1918,4 +1925,175 @@ class AccountsManagerTest {
     when(accounts.getByAccountIdentifierAsync(account.getAccountIdentifier()))
         .thenReturn(CompletableFuture.completedFuture(Optional.of(account)));
   }
+
+  @Nested
+  class Totp {
+
+    @Test
+    void generatePendingTotpKey() throws TooManyTotpKeysException {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      final TotpKey pendingTotpKey = accountsManager.generatePendingTotpKey(accountIdentifier);
+
+      verify(account).setPendingTotpKey(pendingTotpKey);
+    }
+
+    @Test
+    void generatePendingTotpKeyTooManyConfirmedKeys() {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(account.getTotpKeys()).thenReturn(IntStream.range(0, AccountsManager.MAX_TOTP_KEYS)
+          .boxed()
+          .collect(Collectors.toMap(keyId -> keyId, _ -> new AnnotatedTotpKey(
+              new TotpKey(AccountsManager.TOTP_PARAMETERS, TestRandomUtil.nextBytes(16)),
+              TestRandomUtil.nextBytes(16)))));
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      assertThrows(TooManyTotpKeysException.class, () -> accountsManager.generatePendingTotpKey(accountIdentifier));
+      verify(account, never()).setPendingTotpKey(any());
+    }
+
+    @Test
+    void confirmPendingTotpKey() throws InvalidKeyException, TooManyTotpKeysException, NoSuchAlgorithmException {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      final TotpKey pendingTotpKey = accountsManager.generatePendingTotpKey(accountIdentifier);
+      final int nextTotpKeyId = ThreadLocalRandom.current().nextInt();
+
+      when(account.getPendingTotpKey()).thenReturn(Optional.of(pendingTotpKey));
+      when(account.getNextTotpKeyId()).thenReturn(nextTotpKeyId);
+
+      final TimeBasedOneTimePasswordGenerator totpGenerator =
+          new TimeBasedOneTimePasswordGenerator(AccountsManager.TOTP_PARAMETERS.timeStep(),
+              AccountsManager.TOTP_PARAMETERS.passwordLength(),
+              AccountsManager.TOTP_PARAMETERS.algorithm());
+
+      final Instant timestamp = Instant.now();
+
+      assertEquals(Optional.of(nextTotpKeyId), accountsManager.confirmPendingTotpKey(accountIdentifier,
+          totpGenerator.generateOneTimePassword(pendingTotpKey, timestamp),
+          timestamp,
+          TestRandomUtil.nextBytes(16)));
+
+      verify(account).setPendingTotpKey(null);
+    }
+
+    @Test
+    void confirmPendingTotpKeyPreviouslyConfirmed() throws InvalidKeyException, NoSuchAlgorithmException {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      final AnnotatedTotpKey confirmedTotpKey;
+      {
+        final KeyGenerator totpKeyGenerator = KeyGenerator.getInstance(AccountsManager.TOTP_PARAMETERS.algorithm());
+        totpKeyGenerator.init(AccountsManager.TOTP_KEY_LENGTH_BITS);
+
+        confirmedTotpKey = new AnnotatedTotpKey(
+            new TotpKey(AccountsManager.TOTP_PARAMETERS, totpKeyGenerator.generateKey().getEncoded()),
+            TestRandomUtil.nextBytes(16));
+      }
+
+      final int keyId = 17;
+
+      when(account.getPendingTotpKey()).thenReturn(Optional.empty());
+      when(account.getTotpKeys()).thenReturn(Map.of(
+          keyId - 1, confirmedTotpKey,
+          keyId, confirmedTotpKey));
+
+      final TimeBasedOneTimePasswordGenerator totpGenerator =
+          new TimeBasedOneTimePasswordGenerator(AccountsManager.TOTP_PARAMETERS.timeStep(),
+              AccountsManager.TOTP_PARAMETERS.passwordLength(),
+              AccountsManager.TOTP_PARAMETERS.algorithm());
+
+      final Instant timestamp = Instant.now();
+
+      assertEquals(Optional.of(keyId), accountsManager.confirmPendingTotpKey(accountIdentifier,
+          totpGenerator.generateOneTimePassword(confirmedTotpKey, timestamp),
+          timestamp,
+          TestRandomUtil.nextBytes(16)));
+    }
+
+    @Test
+    void confirmPendingTotpKeyNoKeys() throws InvalidKeyException, TooManyTotpKeysException, NoSuchAlgorithmException {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      final TotpKey pendingTotpKey = accountsManager.generatePendingTotpKey(accountIdentifier);
+
+      when(account.getPendingTotpKey()).thenReturn(Optional.empty());
+      when(account.getTotpKeys()).thenReturn(Collections.emptyMap());
+
+      final TimeBasedOneTimePasswordGenerator totpGenerator =
+          new TimeBasedOneTimePasswordGenerator(AccountsManager.TOTP_PARAMETERS.timeStep(),
+              AccountsManager.TOTP_PARAMETERS.passwordLength(),
+              AccountsManager.TOTP_PARAMETERS.algorithm());
+
+      final Instant timestamp = Instant.now();
+
+      assertEquals(Optional.empty(), accountsManager.confirmPendingTotpKey(accountIdentifier,
+          totpGenerator.generateOneTimePassword(pendingTotpKey, timestamp),
+          timestamp,
+          TestRandomUtil.nextBytes(16)));
+    }
+
+    @Test
+    void confirmPendingTotpKeyIncorrectPassword()
+        throws InvalidKeyException, TooManyTotpKeysException, NoSuchAlgorithmException {
+      final UUID accountIdentifier = UUID.randomUUID();
+
+      final Account account = mock(Account.class);
+      when(account.getAccountIdentifier()).thenReturn(accountIdentifier);
+
+      when(accounts.getByAccountIdentifier(accountIdentifier))
+          .thenReturn(Optional.of(account));
+
+      final TotpKey pendingTotpKey = accountsManager.generatePendingTotpKey(accountIdentifier);
+      final int nextTotpKeyId = ThreadLocalRandom.current().nextInt();
+
+      when(account.getPendingTotpKey()).thenReturn(Optional.of(pendingTotpKey));
+      when(account.getNextTotpKeyId()).thenReturn(nextTotpKeyId);
+
+      final TimeBasedOneTimePasswordGenerator totpGenerator =
+          new TimeBasedOneTimePasswordGenerator(AccountsManager.TOTP_PARAMETERS.timeStep(),
+              AccountsManager.TOTP_PARAMETERS.passwordLength(),
+              AccountsManager.TOTP_PARAMETERS.algorithm());
+
+      final Instant timestamp = Instant.now();
+      final int incorrectPassword = totpGenerator.generateOneTimePassword(pendingTotpKey, timestamp) + 1;
+
+      assertEquals(Optional.empty(), accountsManager.confirmPendingTotpKey(accountIdentifier,
+          incorrectPassword,
+          timestamp,
+          TestRandomUtil.nextBytes(16)));
+
+      verify(account, never()).setPendingTotpKey(null);
+    }
+  }
 }