diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java index 79d111c9b..c1996874c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java @@ -410,6 +410,10 @@ public class Accounts { accountToCreate.setZkCredentialKey(existingAccount.getZkCredentialKey().orElse(null)); accountToCreate.setZkCredentialKeyRotationId(existingAccount.getZkCredentialKeyRotationId()); + // Carry over any existing TOTP keys to the new account; we don't need to copy the pending TOTP key since that's + // just a temporary holding place for essentially ephemeral data + accountToCreate.setTotpKeys(new HashMap<>(existingAccount.getTotpKeys())); + final List writeItems = new ArrayList<>(); // If we're reclaiming an account that already has a username, we'd like to give the re-registering client diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java index a6f8a36e2..1ec9b2396 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -23,6 +23,8 @@ import static org.mockito.Mockito.when; import static org.whispersystems.textsecuregcm.storage.ReceiptCredentialTestUtil.receiptPresentation; import static org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil.assertFailsWithCause; +import com.eatthepath.otp.HmacOneTimePasswordGenerator; +import com.eatthepath.otp.TimeBasedOneTimePasswordGenerator; import com.fasterxml.jackson.core.JsonProcessingException; import com.google.i18n.phonenumbers.PhoneNumberUtil; import jakarta.annotation.Nullable; @@ -112,7 +114,6 @@ class AccountsTest { private static final AtomicInteger ACCOUNT_COUNTER = new AtomicInteger(1); - @RegisterExtension static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension( Tables.ACCOUNTS, @@ -556,6 +557,46 @@ class AccountsTest { assertThat(reclaimed.getZkCredentialKey()).hasValue(existingAccount.getZkCredentialKey().orElseThrow()); } + @ParameterizedTest + @ValueSource(strings = {"+14151112222"}) + @NullSource + void testReclaimAccountPreservesTotpKeys(@Nullable final String number) throws Exception { + final UUID existingUuid = UUID.randomUUID(); + final byte[] accountRecoveryPassword = TestRandomUtil.nextBytes(16); + final Account existingAccount = + generateAccount(number, existingUuid, number == null ? null : UUID.randomUUID(), + List.of(generateDevice(DEVICE_ID_1)), accountRecoveryPassword); + + existingAccount.setTotpKeys(Map.of((byte) 1, new AnnotatedTotpKey(new TotpKey( + new TotpParameters( + TimeBasedOneTimePasswordGenerator.TOTP_ALGORITHM_HMAC_SHA1, + HmacOneTimePasswordGenerator.DEFAULT_PASSWORD_LENGTH, + TimeBasedOneTimePasswordGenerator.DEFAULT_TIME_STEP), + TestRandomUtil.nextBytes(16)), + TestRandomUtil.nextBytes(16)))); + + final ReceiptCredentialPresentation receiptCredentialPresentation = receiptPresentation(); + if (number != null) { + createAccount(existingAccount); + } else { + createNumberlessAccount(existingAccount, receiptCredentialPresentation, accountRecoveryPassword); + } + + final Account secondAccount = + generateAccount(number, UUID.randomUUID(), number == null ? null : UUID.randomUUID(), + List.of(generateDevice(DEVICE_ID_1)), accountRecoveryPassword); + + if (number != null) { + reclaimAccount(secondAccount); + } else { + reclaimNumberlessAccount(secondAccount, receiptCredentialPresentation, accountRecoveryPassword); + } + + final Account reclaimed = accounts.getByAccountIdentifier(existingUuid).orElseThrow(); + + assertThat(reclaimed.getTotpKeys()).isEqualTo(existingAccount.getTotpKeys()); + } + @ParameterizedTest @ValueSource(strings = {"+14151112222"}) @NullSource