From 739ed56b4c4e872059be8db34b318cbfb70b905f Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Fri, 22 Nov 2024 15:57:46 -0600
Subject: [PATCH] =?UTF-8?q?Don=E2=80=99t=20use=20an=20existing=20record's?=
 =?UTF-8?q?=20number=20in=20`AccountsManager`=20re-registration=20handling?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../textsecuregcm/storage/AccountsManager.java   |  9 +++------
 .../storage/AccountsManagerTest.java             | 16 ++++++++++++----
 2 files changed, 15 insertions(+), 10 deletions(-)

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 8c34d49f3..30f71ed6a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
@@ -271,17 +271,17 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       @Nullable final String userAgent) throws InterruptedException {
 
     final Account account = new Account();
-    final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
+    final UUID pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
 
     return createTimer.record(() -> {
-      accountLockManager.withLock(List.of(phoneNumberIdentifier), () -> {
+      accountLockManager.withLock(List.of(pni), () -> {
         final Optional<UUID> maybeRecentlyDeletedAccountIdentifier =
             accounts.findRecentlyDeletedAccountIdentifier(number);
 
         // Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is
         // re-registering.
-        account.setNumber(number, phoneNumberIdentifier);
         account.setUuid(maybeRecentlyDeletedAccountIdentifier.orElseGet(UUID::randomUUID));
+        account.setNumber(number, pni);
         account.setIdentityKey(aciIdentityKey);
         account.setPhoneNumberIdentityKey(pniIdentityKey);
         account.addDevice(primaryDeviceSpec.toDevice(Device.PRIMARY_ID, clock));
@@ -325,10 +325,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
           }
 
           final UUID aci = e.getExistingAccount().getIdentifier(IdentityType.ACI);
-          final UUID pni = e.getExistingAccount().getIdentifier(IdentityType.PNI);
-
           account.setUuid(aci);
-          account.setNumber(e.getExistingAccount().getNumber(), pni);
 
           final List<TransactWriteItem> additionalWriteItems = Stream.concat(
                   keysManager.buildWriteItemsForNewDevice(account.getIdentifier(IdentityType.ACI),
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 4e4aa9f34..84797715c 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java
@@ -66,6 +66,7 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.stubbing.Answer;
@@ -837,11 +838,17 @@ class AccountsManagerTest {
     verifyNoInteractions(profilesManager);
   }
 
-  @Test
-  void testReregisterAccount() throws InterruptedException, AccountAlreadyExistsException {
+  @ParameterizedTest
+  @CsvSource({
+      "+18005550123, +18005550123",
+      // the canonical form of numbers may change over time, so an existing account might have not-identical e164 that
+      // maps to the same PNI, and the number used by the caller must be present on the re-registered account
+      "+2290123456789, +22923456789"
+  })
+  void testReregisterAccount(final String e164, final String existingAccountE164)
+      throws InterruptedException, AccountAlreadyExistsException {
     final UUID existingUuid = UUID.randomUUID();
 
-    final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null);
 
     when(accounts.create(any(), any()))
@@ -851,7 +858,7 @@ class AccountsManagerTest {
           final Account existingAccount = mock(Account.class);
           when(existingAccount.getUuid()).thenReturn(existingUuid);
           when(existingAccount.getIdentifier(IdentityType.ACI)).thenReturn(existingUuid);
-          when(existingAccount.getNumber()).thenReturn(e164);
+          when(existingAccount.getNumber()).thenReturn(existingAccountE164);
           when(existingAccount.getPhoneNumberIdentifier()).thenReturn(requestedAccount.getIdentifier(IdentityType.PNI));
           when(existingAccount.getIdentifier(IdentityType.PNI)).thenReturn(requestedAccount.getIdentifier(IdentityType.PNI));
           when(existingAccount.getPrimaryDevice()).thenReturn(mock(Device.class));
@@ -864,6 +871,7 @@ class AccountsManagerTest {
     final Account reregisteredAccount = createAccount(e164, attributes);
 
     assertTrue(phoneNumberIdentifiersByE164.containsKey(e164));
+    assertEquals(e164, reregisteredAccount.getNumber());
 
     verify(accounts)
         .create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())), any());