From 38ca4e6b278cfe9978140b1a4c4c50fa5245a033 Mon Sep 17 00:00:00 2001 From: Jon Chambers <63609320+jon-signal@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:16:23 -0400 Subject: [PATCH] Add support for generating, storing, and managing TOTP keys --- pom.xml | 6 + service/pom.xml | 5 + .../textsecuregcm/storage/Account.java | 37 ++++ .../storage/AccountsManager.java | 142 ++++++++++++++ .../storage/AnnotatedTotpKey.java | 32 ++++ .../storage/TooManyTotpKeysException.java | 11 ++ .../textsecuregcm/storage/TotpKey.java | 34 ++++ .../textsecuregcm/storage/TotpParameters.java | 19 ++ .../textsecuregcm/storage/AccountTest.java | 21 +++ .../storage/AccountsManagerTest.java | 178 ++++++++++++++++++ 10 files changed, 485 insertions(+) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/TooManyTotpKeysException.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpParameters.java diff --git a/pom.xml b/pom.xml index 1e54c6873..bffbf2649 100644 --- a/pom.xml +++ b/pom.xml @@ -64,6 +64,7 @@ 4.4.16 4.5.14 2.22.2 + 1.0.0 2.3.0 3.0.2 2.4.10 @@ -191,6 +192,11 @@ pom import + + com.eatthepath + java-otp + ${java-otp.version} + com.eatthepath pushy diff --git a/service/pom.xml b/service/pom.xml index d90b48034..02ea62c0d 100644 --- a/service/pom.xml +++ b/service/pom.xml @@ -436,6 +436,11 @@ lettuce-core + + com.eatthepath + java-otp + + com.eatthepath pushy diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java index 75bbc1847..5f0311abe 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java +++ b/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 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 getPendingTotpKey() { + requireNotStale(); + return Optional.ofNullable(pendingTotpKey); + } + + public int getNextTotpKeyId() { + requireNotStale(); + + return totpKeys.keySet().stream() + .mapToInt(i -> i) + .max() + .orElse(-1) + 1; + } + + public Map getTotpKeys() { + requireNotStale(); + return totpKeys; + } + + public void setTotpKeys(final Map totpKeys) { + requireNotStale(); + this.totpKeys = totpKeys; + } + public void markStale() { stale = true; } 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 b4225a89a..7b7737296 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/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 implemen private final ScheduledExecutorService retryExecutor; private final Clock clock; + private final KeyGenerator totpKeyGenerator; + private final TimeBasedOneTimePasswordGenerator totpGenerator; + private final Key verificationTokenKey; private final FaultTolerantPubSubConnection pubSubConnection; @@ -194,6 +204,17 @@ public class AccountsManager extends RedisPubSubAdapter 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 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 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 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 confirmPendingTotpKey(final UUID accountIdentifier, + final int oneTimePassword, + final Instant timestamp, + final byte[] metadataCiphertext) { + + final Optional maybeAccount = accounts.getByAccountIdentifier(accountIdentifier); + + if (maybeAccount.isEmpty()) { + return Optional.empty(); + } + + final Optional 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 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)); + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AnnotatedTotpKey.java new file mode 100644 index 000000000..d960ad769 --- /dev/null +++ b/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(); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/TooManyTotpKeysException.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/TooManyTotpKeysException.java new file mode 100644 index 000000000..cf06c8df3 --- /dev/null +++ b/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 { +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpKey.java new file mode 100644 index 000000000..6a721493d --- /dev/null +++ b/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); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpParameters.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/TotpParameters.java new file mode 100644 index 000000000..3919f5a4e --- /dev/null +++ b/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) { +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountTest.java index e08a66c8d..74bceb657 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountTest.java +++ b/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()); + } } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java index e6869675e..e16b0eef4 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java +++ b/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); + } + } }