Fix username restore issue.

app/src/main/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageController.kt

@@ -79,6 +79,7 @@ import org.thoughtcrime.securesms.jobs.CheckKeyTransparencyJob
 import org.thoughtcrime.securesms.jobs.DirectoryRefreshJob
 import org.thoughtcrime.securesms.jobs.LocalBackupRestoreMediaJob
 import org.thoughtcrime.securesms.jobs.PreKeysSyncJob
+import org.thoughtcrime.securesms.jobs.ReclaimUsernameAndLinkJob
 import org.thoughtcrime.securesms.jobs.RefreshOwnProfileJob
 import org.thoughtcrime.securesms.jobs.RotateCertificateJob
 import org.thoughtcrime.securesms.keyvalue.Completed
@@ -284,6 +285,11 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
   }
 
   override suspend fun onRegistrationFlowFinished() = withContext(Dispatchers.Default) {
+    if (SignalStore.misc.needsUsernameRestore) {
+      Log.i(TAG, "[onRegistrationFlowFinished] Username reclaim is still pending. Enqueuing a job to handle it.")
+      AppDependencies.jobManager.add(ReclaimUsernameAndLinkJob())
+    }
+
     if (Environment.MOCK_PHONE_NUMBERLESS_REGISTRATION) {
       Log.w(TAG, "[onRegistrationFlowFinished] Mocking a phone-number-less account. Wiping all local knowledge of the E164 and PNI.")
 
@@ -808,6 +814,10 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
       // the remote metadata during the first BackupMessagesJob) we must not start a new chain, or the existing remote backup becomes unrestorable.
       Log.i(TAG, "[applyAccountData] Re-registration. Marking that we need to restore SVRB secrets before backing up.")
       SignalStore.backup.backupSecretRestoreRequired = true
+
+      // Registering releases any username we previously held, so it has to be re-reserved once storage service tells us what it was.
+      Log.i(TAG, "[applyAccountData] Re-registration. Marking that we need to reclaim our username and link.")
+      SignalStore.misc.needsUsernameRestore = true
     }
 
     SignalStore.account.setServicePassword(accountData.servicePassword)

app/src/test/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageControllerTest.kt

@@ -49,6 +49,7 @@ import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.jobmanager.runJobBlocking
 import org.thoughtcrime.securesms.jobs.DirectoryRefreshJob
 import org.thoughtcrime.securesms.jobs.PreKeysSyncJob
+import org.thoughtcrime.securesms.jobs.ReclaimUsernameAndLinkJob
 import org.thoughtcrime.securesms.jobs.RefreshOwnProfileJob
 import org.thoughtcrime.securesms.jobs.RotateCertificateJob
 import org.thoughtcrime.securesms.keyvalue.SignalStore
@@ -296,6 +297,67 @@ class AppRegistrationStorageControllerTest {
     assertThat(SignalStore.backup.backupSecretRestoreRequired).isFalse()
   }
 
+  @Test
+  fun `commit - re-registration - flags that the username needs to be reclaimed`() = runBlocking<Unit> {
+    seedInProgressData(
+      RegistrationData(
+        accountData = accountData(reRegistration = true),
+        accountEntropyPool = aep.value
+      )
+    )
+
+    controller.commitRegistrationData()
+
+    assertThat(SignalStore.misc.needsUsernameRestore).isTrue()
+  }
+
+  @Test
+  fun `commit - new account - does not flag that the username needs to be reclaimed`() = runBlocking<Unit> {
+    seedInProgressData(
+      RegistrationData(
+        accountData = accountData(reRegistration = false),
+        accountEntropyPool = aep.value
+      )
+    )
+
+    controller.commitRegistrationData()
+
+    assertThat(SignalStore.misc.needsUsernameRestore).isFalse()
+  }
+
+  @Test
+  fun `onRegistrationFlowFinished - username reclaim pending - enqueues reclaim job`() = runBlocking<Unit> {
+    SignalStore.misc.needsUsernameRestore = true
+
+    controller.onRegistrationFlowFinished()
+
+    verify { AppDependencies.jobManager.add(ofType<ReclaimUsernameAndLinkJob>()) }
+  }
+
+  @Test
+  fun `onRegistrationFlowFinished - no username reclaim pending - does not enqueue reclaim job`() = runBlocking<Unit> {
+    SignalStore.misc.needsUsernameRestore = false
+
+    controller.onRegistrationFlowFinished()
+
+    verify(exactly = 0) { AppDependencies.jobManager.add(ofType<ReclaimUsernameAndLinkJob>()) }
+  }
+
+  @Test
+  fun `re-registration - commit then flow finished - enqueues reclaim job`() = runBlocking<Unit> {
+    seedInProgressData(
+      RegistrationData(
+        accountData = accountData(reRegistration = true),
+        accountEntropyPool = aep.value
+      )
+    )
+
+    controller.commitRegistrationData()
+    controller.onRegistrationFlowFinished()
+
+    verify { AppDependencies.jobManager.add(ofType<ReclaimUsernameAndLinkJob>()) }
+  }
+
   @Test
   fun `commit - called twice - only applies account data once`() = runBlocking<Unit> {
     seedInProgressData(

feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt

@@ -142,6 +142,7 @@ class RegistrationEndToEndTest {
     assert(committed.pin == PIN) { "Expected committed pin $PIN but was ${committed.pin}" }
     assert(committed.accountEntropyPool.isNotEmpty()) { "Expected committed AEP to be populated" }
     assert(committed.accountData?.reRegistration == false) { "Expected a new registration to not be flagged as a re-registration" }
+    assert(storageController.registrationFlowFinishedCount == 1) { "Expected the flow-finished hook to fire exactly once but fired ${storageController.registrationFlowFinishedCount} times" }
 
     assert(networkController.lastCreateSessionE164 == E164) { "Expected a session for $E164 but was ${networkController.lastCreateSessionE164}" }
     assert(networkController.lastRegisterAccountRequest?.e164 == E164) { "Expected registration for $E164 but was ${networkController.lastRegisterAccountRequest}" }
@@ -645,6 +646,10 @@ class RegistrationEndToEndTest {
     assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" }
     assert(committed.pin == PIN) { "Expected committed pin $PIN but was ${committed.pin}" }
     assert(committed.accountData?.reRegistration == true) { "Expected the committed account data to be flagged as a re-registration" }
+
+    // The re-registration flag is what tells the app to reclaim the username we just released, and the flow-finished
+    // hook is where it enqueues the job that does it. See AppRegistrationStorageController.
+    assert(storageController.registrationFlowFinishedCount == 1) { "Expected the flow-finished hook to fire exactly once but fired ${storageController.registrationFlowFinishedCount} times" }
   }
 
   @Test

feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt

@@ -39,6 +39,10 @@ class FakeStorageController : StorageController {
   var restoreDecision: RestoreDecision? = null
     private set
 
+  /** How many times [onRegistrationFlowFinished] has been called. The app hangs post-registration bookkeeping off of this. */
+  var registrationFlowFinishedCount: Int = 0
+    private set
+
   /** Simulates a previously-registered device, which the flow will try to re-register via recovery password. */
   var preExistingRegistrationData: PreExistingRegistrationData? = null
 
@@ -102,7 +106,9 @@ class FakeStorageController : StorageController {
     }
   }
 
-  override suspend fun onRegistrationFlowFinished() = Unit
+  override suspend fun onRegistrationFlowFinished() {
+    registrationFlowFinishedCount++
+  }
 
   override suspend fun setRestoreDecision(decision: RestoreDecision) {
     // Mirrors the real controller: only the first decision sticks, later ones are ignored