Preserve TOTP keys when reclaiming accounts

This commit is contained in:
Jon Chambers
2026-08-26 16:01:00 -04:00
committed by Jon Chambers
parent 31b34ed06f
commit eb0ffb0c4d
2 changed files with 46 additions and 1 deletions
@@ -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<TransactWriteItem> writeItems = new ArrayList<>();
// If we're reclaiming an account that already has a username, we'd like to give the re-registering client
@@ -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