Flesh out verification challenge support for registration v2.

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

@@ -72,6 +72,7 @@ import org.whispersystems.signalservice.internal.push.RegistrationSessionMetadat
 import java.io.IOException
 import java.nio.charset.StandardCharsets
 import java.util.Locale
+import java.util.Optional
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import kotlin.time.Duration.Companion.seconds
@@ -277,6 +278,7 @@ object RegistrationRepository {
    */
   suspend fun createSession(context: Context, e164: String, password: String, mcc: String?, mnc: String?): RegistrationSessionCreationResult =
     withContext(Dispatchers.IO) {
+      Log.d(TAG, "About to create a registration session…")
       val fcmToken: String? = FcmUtil.getToken(context).orElse(null)
       val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password).registrationApi
 
@@ -301,14 +303,21 @@ object RegistrationRepository {
    * Validates an existing session, if its ID is provided. If the session is expired/invalid, or none is provided, it will attempt to initiate a new session.
    */
   suspend fun createOrValidateSession(context: Context, sessionId: String?, e164: String, password: String, mcc: String?, mnc: String?): RegistrationSessionResult {
-    if (sessionId != null) {
+    val savedSessionId = if (sessionId == null && e164 == SignalStore.registrationValues().sessionE164) {
+      SignalStore.registrationValues().sessionId
+    } else {
+      sessionId
+    }
+
+    if (savedSessionId != null) {
       Log.d(TAG, "Validating existing registration session.")
-      val sessionValidationResult = validateSession(context, sessionId, e164, password)
+      val sessionValidationResult = validateSession(context, savedSessionId, e164, password)
       when (sessionValidationResult) {
         is RegistrationSessionCheckResult.Success -> {
           Log.d(TAG, "Existing registration session is valid.")
           return sessionValidationResult
         }
+
         is RegistrationSessionCheckResult.UnknownError -> {
           Log.w(TAG, "Encountered error when validating existing session.", sessionValidationResult.getCause())
           return sessionValidationResult
@@ -338,9 +347,9 @@ object RegistrationRepository {
   /**
    * Submits the user-entered verification code to the service.
    */
-  suspend fun submitVerificationCode(context: Context, e164: String, password: String, sessionId: String, registrationData: RegistrationData): VerificationCodeRequestResult =
+  suspend fun submitVerificationCode(context: Context, sessionId: String, registrationData: RegistrationData): VerificationCodeRequestResult =
     withContext(Dispatchers.IO) {
-      val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password).registrationApi
+      val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, registrationData.e164, SignalServiceAddress.DEFAULT_DEVICE_ID, registrationData.password).registrationApi
       val result = api.verifyAccount(sessionId = sessionId, verificationCode = registrationData.code)
       return@withContext VerificationCodeRequestResult.from(result)
     }
@@ -355,6 +364,15 @@ object RegistrationRepository {
       return@withContext VerificationCodeRequestResult.from(captchaSubmissionResult)
     }
 
+  suspend fun requestAndVerifyPushToken(context: Context, sessionId: String, e164: String, password: String) =
+    withContext(Dispatchers.IO) {
+      val fcmToken = getFcmToken(context)
+      val accountManager = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password)
+      val pushChallenge = PushChallengeRequest.getPushChallengeBlocking(accountManager, sessionId, Optional.ofNullable(fcmToken), PUSH_REQUEST_TIMEOUT).orElse(null)
+      val pushSubmissionResult = accountManager.registrationApi.submitPushChallengeToken(sessionId = sessionId, pushChallengeToken = pushChallenge)
+      return@withContext VerificationCodeRequestResult.from(pushSubmissionResult)
+    }
+
   /**
    * Submit the necessary assets as a verified account so that the user can actually use the service.
    */
@@ -439,7 +457,7 @@ object RegistrationRepository {
         Log.i(TAG, "Push challenge unsuccessful. Updating registration state accordingly.")
         return@withContext NetworkResult.ApplicationError<RegistrationSessionMetadataResponse>(NullPointerException())
       } catch (ex: Exception) {
-        Log.w(TAG, "Exception caught, but the earlier try block should have caught it?", ex) // TODO [regv2]: figure out why this exception is not caught
+        Log.w(TAG, "Exception caught, but the earlier try block should have caught it?", ex)
         return@withContext NetworkResult.ApplicationError<RegistrationSessionMetadataResponse>(ex)
       }
     }
@@ -529,8 +547,7 @@ object RegistrationRepository {
   enum class Mode(val isSmsRetrieverSupported: Boolean, val transport: PushServiceSocket.VerificationCodeTransport) {
     SMS_WITH_LISTENER(true, PushServiceSocket.VerificationCodeTransport.SMS),
     SMS_WITHOUT_LISTENER(false, PushServiceSocket.VerificationCodeTransport.SMS),
-    PHONE_CALL(false, PushServiceSocket.VerificationCodeTransport.VOICE),
+    PHONE_CALL(false, PushServiceSocket.VerificationCodeTransport.VOICE)
-    NONE(false, PushServiceSocket.VerificationCodeTransport.SMS)
   }
 
   private class PushTokenChallengeSubscriber {

app/src/main/java/org/thoughtcrime/securesms/registration/v2/data/network/Challenge.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.registration.v2.data.network
+
+import org.signal.core.util.logging.Log
+
+enum class Challenge(val key: String) {
+  CAPTCHA("captcha"),
+  PUSH("pushChallenge");
+
+  companion object {
+    private val TAG = Log.tag(Challenge::class)
+
+    fun parse(strings: List<String>): List<Challenge> {
+      return strings.mapNotNull {
+        when (it) {
+          CAPTCHA.key -> CAPTCHA
+          PUSH.key -> PUSH
+          else -> {
+            Log.i(TAG, "Encountered unknown challenge type: $it")
+            null
+          }
+        }
+      }
+    }
+  }
+
+  fun stringify(challenges: List<Challenge>): String {
+    return challenges.joinToString { it.key }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/registration/v2/data/network/RegistrationSessionResult.kt

@@ -8,6 +8,7 @@ package org.thoughtcrime.securesms.registration.v2.data.network
 import org.signal.core.util.logging.Log
 import org.whispersystems.signalservice.api.NetworkResult
 import org.whispersystems.signalservice.api.push.exceptions.MalformedRequestException
+import org.whispersystems.signalservice.api.push.exceptions.NoSuchSessionException
 import org.whispersystems.signalservice.api.push.exceptions.NotFoundException
 import org.whispersystems.signalservice.api.push.exceptions.RateLimitException
 import org.whispersystems.signalservice.internal.push.RegistrationSessionMetadataResponse
@@ -29,11 +30,12 @@ sealed class RegistrationSessionCreationResult(cause: Throwable?) : Registration
         is NetworkResult.Success -> {
           Success(networkResult.result)
         }
+
         is NetworkResult.ApplicationError -> UnknownError(networkResult.throwable)
         is NetworkResult.NetworkError -> UnknownError(networkResult.exception)
         is NetworkResult.StatusCodeError -> {
           when (val cause = networkResult.exception) {
-            is RateLimitException -> RateLimited(cause)
+            is RateLimitException -> createRateLimitProcessor(cause)
             is MalformedRequestException -> MalformedRequest(cause)
             else -> if (networkResult.code == 422) {
               ServerUnableToParse(cause)
@@ -44,6 +46,14 @@ sealed class RegistrationSessionCreationResult(cause: Throwable?) : Registration
         }
       }
     }
+
+    private fun createRateLimitProcessor(exception: RateLimitException): RegistrationSessionCreationResult {
+      return if (exception.retryAfterMilliseconds.isPresent) {
+        RateLimited(exception, exception.retryAfterMilliseconds.get())
+      } else {
+        AttemptsExhausted(exception)
+      }
+    }
   }
 
   class Success(private val metadata: RegistrationSessionMetadataResponse) : RegistrationSessionCreationResult(null), SessionMetadataHolder {
@@ -52,7 +62,8 @@ sealed class RegistrationSessionCreationResult(cause: Throwable?) : Registration
     }
   }
 
-  class RateLimited(cause: Throwable) : RegistrationSessionCreationResult(cause)
+  class RateLimited(cause: Throwable, val timeRemaining: Long) : RegistrationSessionCreationResult(cause)
+  class AttemptsExhausted(cause: Throwable) : RegistrationSessionCreationResult(cause)
   class ServerUnableToParse(cause: Throwable) : RegistrationSessionCreationResult(cause)
   class MalformedRequest(cause: Throwable) : RegistrationSessionCreationResult(cause)
   class UnknownError(cause: Throwable) : RegistrationSessionCreationResult(cause)
@@ -70,7 +81,7 @@ sealed class RegistrationSessionCheckResult(cause: Throwable?) : RegistrationSes
         is NetworkResult.NetworkError -> UnknownError(networkResult.exception)
         is NetworkResult.StatusCodeError -> {
           when (val cause = networkResult.exception) {
-            is NotFoundException -> SessionNotFound(cause)
+            is NoSuchSessionException, is NotFoundException -> SessionNotFound(cause)
             else -> UnknownError(cause)
           }
         }

app/src/main/java/org/thoughtcrime/securesms/registration/v2/data/network/VerificationCodeRequestResult.kt

@@ -9,11 +9,13 @@ import okio.IOException
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
 import org.whispersystems.signalservice.api.NetworkResult
+import org.whispersystems.signalservice.api.push.exceptions.AlreadyVerifiedException
 import org.whispersystems.signalservice.api.push.exceptions.CaptchaRequiredException
 import org.whispersystems.signalservice.api.push.exceptions.ExternalServiceFailureException
 import org.whispersystems.signalservice.api.push.exceptions.ImpossiblePhoneNumberException
 import org.whispersystems.signalservice.api.push.exceptions.InvalidTransportModeException
 import org.whispersystems.signalservice.api.push.exceptions.MalformedRequestException
+import org.whispersystems.signalservice.api.push.exceptions.NoSuchSessionException
 import org.whispersystems.signalservice.api.push.exceptions.NonNormalizedPhoneNumberException
 import org.whispersystems.signalservice.api.push.exceptions.PushChallengeRequiredException
 import org.whispersystems.signalservice.api.push.exceptions.RateLimitException
@@ -36,15 +38,17 @@ sealed class VerificationCodeRequestResult(cause: Throwable?) : RegistrationResu
     fun from(networkResult: NetworkResult<RegistrationSessionMetadataResponse>): VerificationCodeRequestResult {
       return when (networkResult) {
         is NetworkResult.Success -> {
-          val challenges = networkResult.result.body.requestedInformation
+          val challenges = Challenge.parse(networkResult.result.body.requestedInformation)
           if (challenges.isNotEmpty()) {
+            Log.d(TAG, "Received \"successful\" response that contains challenges: ${challenges.joinToString { it.key }}")
             ChallengeRequired(challenges)
           } else {
             Success(
               sessionId = networkResult.result.body.id,
               allowedToRequestCode = networkResult.result.body.allowedToRequestCode,
               nextSmsTimestamp = RegistrationRepository.deriveTimestamp(networkResult.result.headers, networkResult.result.body.nextSms),
-              nextCallTimestamp = RegistrationRepository.deriveTimestamp(networkResult.result.headers, networkResult.result.body.nextCall)
+              nextCallTimestamp = RegistrationRepository.deriveTimestamp(networkResult.result.headers, networkResult.result.body.nextCall),
+              verified = networkResult.result.body.verified
             )
           }
         }
@@ -64,6 +68,8 @@ sealed class VerificationCodeRequestResult(cause: Throwable?) : RegistrationResu
             is MalformedRequestException -> MalformedRequest(cause)
             is RegistrationRetryException -> MustRetry(cause)
             is LockedException -> RegistrationLocked(cause = cause, timeRemaining = cause.timeRemaining)
+            is NoSuchSessionException -> NoSuchSession(cause)
+            is AlreadyVerifiedException -> AlreadyVerified(cause)
             else -> UnknownError(cause)
           }
         }
@@ -78,7 +84,7 @@ sealed class VerificationCodeRequestResult(cause: Throwable?) : RegistrationResu
 
       try {
         val response = JsonUtil.fromJson(errorResult.body, RegistrationSessionMetadataJson::class.java)
-        return ChallengeRequired(response.requestedInformation)
+        return ChallengeRequired(Challenge.parse(response.requestedInformation))
       } catch (parseException: IOException) {
         Log.w(TAG, "Attempted to parse error body for list of requested information, but encountered exception.", parseException)
         return UnknownError(parseException)
@@ -94,9 +100,9 @@ sealed class VerificationCodeRequestResult(cause: Throwable?) : RegistrationResu
     }
   }
 
-  class Success(val sessionId: String, val allowedToRequestCode: Boolean, val nextSmsTimestamp: Long, val nextCallTimestamp: Long) : VerificationCodeRequestResult(null)
+  class Success(val sessionId: String, val allowedToRequestCode: Boolean, val nextSmsTimestamp: Long, val nextCallTimestamp: Long, val verified: Boolean) : VerificationCodeRequestResult(null)
 
-  class ChallengeRequired(val challenges: List<String>) : VerificationCodeRequestResult(null)
+  class ChallengeRequired(val challenges: List<Challenge>) : VerificationCodeRequestResult(null)
 
   class RateLimited(cause: Throwable, val timeRemaining: Long) : VerificationCodeRequestResult(cause)
 
@@ -118,5 +124,9 @@ sealed class VerificationCodeRequestResult(cause: Throwable?) : RegistrationResu
 
   class RegistrationLocked(cause: Throwable, val timeRemaining: Long) : VerificationCodeRequestResult(cause)
 
+  class NoSuchSession(cause: Throwable) : VerificationCodeRequestResult(cause)
+
+  class AlreadyVerified(cause: Throwable) : VerificationCodeRequestResult(cause)
+
   class UnknownError(cause: Throwable) : VerificationCodeRequestResult(cause)
 }

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

@@ -7,6 +7,7 @@ package org.thoughtcrime.securesms.registration.v2.ui
 
 import com.google.i18n.phonenumbers.Phonenumber
 import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.registration.v2.data.network.Challenge
 import org.whispersystems.signalservice.internal.push.AuthCredentials
 
 /**
@@ -14,6 +15,7 @@ import org.whispersystems.signalservice.internal.push.AuthCredentials
  */
 data class RegistrationV2State(
   val sessionId: String? = null,
+  val enteredCode: String? = null,
   val phoneNumber: Phonenumber.PhoneNumber? = null,
   val inProgress: Boolean = false,
   val isReRegister: Boolean = false,
@@ -22,13 +24,19 @@ data class RegistrationV2State(
   val svrAuthCredentials: AuthCredentials? = null,
   val svrTriesRemaining: Int = 10,
   val isRegistrationLockEnabled: Boolean = false,
+  val lockedTimeRemaining: Long = 0L,
   val userSkippedReregistration: Boolean = false,
   val isFcmSupported: Boolean = false,
+  val isAllowedToRequestCode: Boolean = false,
   val fcmToken: String? = null,
+  val challengesRequested: List<Challenge> = emptyList(),
+  val challengesPresented: Set<Challenge> = emptySet(),
   val captchaToken: String? = null,
-  val nextSms: Long = 0L,
+  val nextSmsTimestamp: Long = 0L,
-  val nextCall: Long = 0L,
+  val nextCallTimestamp: Long = 0L,
   val smsListenerTimeout: Long = 0L,
   val registrationCheckpoint: RegistrationCheckpoint = RegistrationCheckpoint.INITIALIZATION,
   val networkError: Throwable? = null
-)
+) {
+  val challengesRemaining: List<Challenge> = challengesRequested.filterNot { it in challengesPresented }
+}

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

@@ -15,10 +15,10 @@ import com.google.i18n.phonenumbers.Phonenumber
 import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
-import org.signal.core.util.isNotNullOrBlank
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.jobs.MultiDeviceProfileContentUpdateJob
@@ -30,11 +30,13 @@ import org.thoughtcrime.securesms.registration.RegistrationData
 import org.thoughtcrime.securesms.registration.RegistrationUtil
 import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
 import org.thoughtcrime.securesms.registration.v2.data.network.BackupAuthCheckResult
+import org.thoughtcrime.securesms.registration.v2.data.network.Challenge
 import org.thoughtcrime.securesms.registration.v2.data.network.RegisterAccountResult
+import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationResult
 import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationSessionCheckResult
 import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationSessionCreationResult
 import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationSessionResult
-import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult
+import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.AlreadyVerified
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.AttemptsExhausted
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.ChallengeRequired
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.ExternalServiceFailure
@@ -42,6 +44,7 @@ import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeR
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.InvalidTransportModeFailure
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.MalformedRequest
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.MustRetry
+import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.NoSuchSession
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.NonNormalizedNumber
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.RateLimited
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.RegistrationLocked
@@ -77,6 +80,16 @@ class RegistrationV2ViewModel : ViewModel() {
 
   val uiState = store.asLiveData()
 
+  val checkpoint = store.map { it.registrationCheckpoint }.asLiveData()
+
+  val lockedTimeRemaining = store.map { it.lockedTimeRemaining }.asLiveData()
+
+  val svrTriesRemaining: Int
+    get() = store.value.svrTriesRemaining
+
+  val isReregister: Boolean
+    get() = store.value.isReRegister
+
   init {
     val existingE164 = SignalStore.registrationValues().sessionE164
     if (existingE164 != null) {
@@ -118,6 +131,18 @@ class RegistrationV2ViewModel : ViewModel() {
     }
   }
 
+  fun addPresentedChallenge(challenge: Challenge) {
+    store.update {
+      it.copy(challengesPresented = it.challengesPresented.plus(challenge))
+    }
+  }
+
+  fun removePresentedChallenge(challenge: Challenge) {
+    store.update {
+      it.copy(challengesPresented = it.challengesPresented.minus(challenge))
+    }
+  }
+
   fun fetchFcmToken(context: Context) {
     viewModelScope.launch(context = coroutineExceptionHandler) {
       val fcmToken = RegistrationRepository.getFcmToken(context)
@@ -144,7 +169,7 @@ class RegistrationV2ViewModel : ViewModel() {
     }
   }
 
-  fun onUserConfirmedPhoneNumber(context: Context, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit) {
+  fun onUserConfirmedPhoneNumber(context: Context, errorHandler: (RegistrationResult) -> Unit) {
     setRegistrationCheckpoint(RegistrationCheckpoint.PHONE_NUMBER_CONFIRMED)
     val state = store.value
     if (state.phoneNumber == null) {
@@ -188,12 +213,12 @@ class RegistrationV2ViewModel : ViewModel() {
         }
       }
 
-      val validSession = getOrCreateValidSession(context) ?: return@launch
+      val validSession = getOrCreateValidSession(context, errorHandler) ?: return@launch
 
       if (!validSession.body.allowedToRequestCode) {
         val challenges = validSession.body.requestedInformation.joinToString()
         Log.i(TAG, "Not allowed to request code! Remaining challenges: $challenges")
-        handleSessionStateResult(context, ChallengeRequired(validSession.body.requestedInformation), RegistrationRepository.Mode.NONE, errorHandler)
+        handleSessionStateResult(context, ChallengeRequired(Challenge.parse(validSession.body.requestedInformation)), errorHandler)
         return@launch
       }
 
@@ -201,7 +226,7 @@ class RegistrationV2ViewModel : ViewModel() {
     }
   }
 
-  fun requestSmsCode(context: Context, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit) {
+  fun requestSmsCode(context: Context, errorHandler: (RegistrationResult) -> Unit) {
     val e164 = getCurrentE164()
 
     if (e164 == null) {
@@ -211,12 +236,12 @@ class RegistrationV2ViewModel : ViewModel() {
     }
 
     viewModelScope.launch {
-      val validSession = getOrCreateValidSession(context) ?: return@launch
+      val validSession = getOrCreateValidSession(context, errorHandler) ?: return@launch
       requestSmsCodeInternal(context, validSession.body.id, e164, errorHandler)
     }
   }
 
-  fun requestVerificationCall(context: Context, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit) {
+  fun requestVerificationCall(context: Context, errorHandler: (RegistrationResult) -> Unit) {
     val e164 = getCurrentE164()
 
     if (e164 == null) {
@@ -226,7 +251,7 @@ class RegistrationV2ViewModel : ViewModel() {
     }
 
     viewModelScope.launch {
-      val validSession = getOrCreateValidSession(context) ?: return@launch
+      val validSession = getOrCreateValidSession(context, errorHandler) ?: return@launch
       Log.d(TAG, "Requesting voice call code…")
       val codeRequestResponse = RegistrationRepository.requestSmsCode(
         context = context,
@@ -237,13 +262,13 @@ class RegistrationV2ViewModel : ViewModel() {
       )
       Log.d(TAG, "Voice call code request submitted.")
 
-      handleSessionStateResult(context, codeRequestResponse, RegistrationRepository.Mode.PHONE_CALL, errorHandler)
+      handleSessionStateResult(context, codeRequestResponse, errorHandler)
     }
   }
 
-  private suspend fun requestSmsCodeInternal(context: Context, sessionId: String, e164: String, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit) {
+  private suspend fun requestSmsCodeInternal(context: Context, sessionId: String, e164: String, errorHandler: (RegistrationResult) -> Unit) {
     var smsListenerReady = false
-    Log.d(TAG, "Captcha token submitted.")
+    Log.d(TAG, "Initializing SMS listener.")
     if (store.value.smsListenerTimeout < System.currentTimeMillis()) {
       smsListenerReady = store.value.isFcmSupported && RegistrationRepository.registerSmsListener(context)
 
@@ -265,97 +290,92 @@ class RegistrationV2ViewModel : ViewModel() {
       password = password,
       mode = transportMode
     )
-    Log.d(TAG, "SMS code request submitted.")
+    Log.d(TAG, "SMS code request network call completed.")
 
-    handleSessionStateResult(context, codeRequestResponse, transportMode, errorHandler)
+    handleSessionStateResult(context, codeRequestResponse, errorHandler)
+
+    if (codeRequestResponse is Success) {
+      Log.d(TAG, "SMS code request was successful.")
+      store.update {
+        it.copy(
+          registrationCheckpoint = RegistrationCheckpoint.VERIFICATION_CODE_REQUESTED
+        )
+      }
+    }
   }
 
-  private suspend fun getOrCreateValidSession(context: Context): RegistrationSessionMetadataResponse? {
+  private suspend fun getOrCreateValidSession(context: Context, errorHandler: (RegistrationResult) -> Unit): RegistrationSessionMetadataResponse? {
     val e164 = getCurrentE164() ?: throw IllegalStateException("E164 required to create session!")
-    Log.d(TAG, "Validating/creating a registration session.")
     val mccMncProducer = MccMncProducer(context)
 
     val existingSessionId = store.value.sessionId
-    val sessionResult: RegistrationSessionResult = RegistrationRepository.createOrValidateSession(context, existingSessionId, e164, password, mccMncProducer.mcc, mccMncProducer.mnc)
+    return getOrCreateValidSession(
-    when (sessionResult) {
+      context = context,
-      is RegistrationSessionCheckResult.Success -> {
+      existingSessionId = existingSessionId,
-        val metadata = sessionResult.getMetadata()
+      e164 = e164,
-        val newSessionId = metadata.body.id
+      password = password,
-        if (newSessionId.isNotNullOrBlank() && newSessionId != existingSessionId) {
+      mcc = mccMncProducer.mcc,
+      mnc = mccMncProducer.mnc,
+      successListener = { freshSession ->
+        val freshSessionId = freshSession.body.id
+        if (freshSessionId != existingSessionId) {
           store.update {
-            it.copy(
+            it.copy(sessionId = freshSessionId)
-              sessionId = newSessionId
+          }
+        }
+      },
+      errorHandler = errorHandler
     )
   }
-        }
-        Log.d(TAG, "Registration session validated.")
-        return metadata
-      }
 
-      is RegistrationSessionCreationResult.Success -> {
+  fun submitCaptchaToken(context: Context, errorHandler: (RegistrationResult) -> Unit) {
-        val metadata = sessionResult.getMetadata()
-        val newSessionId = metadata.body.id
-        if (newSessionId.isNotNullOrBlank() && newSessionId != existingSessionId) {
-          store.update {
-            it.copy(
-              sessionId = newSessionId
-            )
-          }
-        }
-        Log.d(TAG, "Registration session created.")
-        return metadata
-      }
-
-      is RegistrationSessionCheckResult.SessionNotFound -> {
-        Log.w(TAG, "This should be impossible to reach at this stage; it should have been handled in RegistrationRepository.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-
-      is RegistrationSessionCheckResult.UnknownError -> {
-        Log.i(TAG, "Unknown error occurred while checking registration session.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-
-      is RegistrationSessionCreationResult.MalformedRequest -> {
-        Log.i(TAG, "Malformed request error occurred while creating registration session.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-
-      is RegistrationSessionCreationResult.RateLimited -> {
-        Log.i(TAG, "Rate limit occurred while creating registration session.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-
-      is RegistrationSessionCreationResult.ServerUnableToParse -> {
-        Log.i(TAG, "Server unable to parse request for creating registration session.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-
-      is RegistrationSessionCreationResult.UnknownError -> {
-        Log.i(TAG, "Unknown error occurred while checking registration session.", sessionResult.getCause())
-        handleGenericError(sessionResult.getCause())
-      }
-    }
-    return null
-  }
-
-  fun submitCaptchaToken(context: Context, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit) {
     val e164 = getCurrentE164() ?: throw IllegalStateException("Can't submit captcha token if no phone number is set!")
     val captchaToken = store.value.captchaToken ?: throw IllegalStateException("Can't submit captcha token if no captcha token is set!")
 
     viewModelScope.launch {
-      val session = getOrCreateValidSession(context) ?: return@launch
+      val session = getOrCreateValidSession(context, errorHandler) ?: return@launch
       Log.d(TAG, "Submitting captcha token…")
       val captchaSubmissionResult = RegistrationRepository.submitCaptchaToken(context, e164, password, session.body.id, captchaToken)
       Log.d(TAG, "Captcha token submitted.")
-      handleSessionStateResult(context, captchaSubmissionResult, RegistrationRepository.Mode.NONE, errorHandler)
+      store.update {
+        it.copy(captchaToken = null)
+      }
+      handleSessionStateResult(context, captchaSubmissionResult, errorHandler)
+    }
+  }
+
+  fun requestAndSubmitPushToken(context: Context, errorHandler: (RegistrationResult) -> Unit) {
+    Log.v(TAG, "validatePushToken()")
+
+    addPresentedChallenge(Challenge.PUSH)
+
+    val e164 = getCurrentE164() ?: throw IllegalStateException("Can't submit captcha token if no phone number is set!")
+
+    viewModelScope.launch {
+      Log.d(TAG, "Getting session in order to perform push token verification…")
+      val session = getOrCreateValidSession(context, errorHandler) ?: return@launch
+
+      if (!Challenge.parse(session.body.requestedInformation).contains(Challenge.PUSH)) {
+        Log.d(TAG, "Push submission no longer necessary, bailing.")
+        store.update {
+          it.copy(
+            inProgress = false
+          )
+        }
+        return@launch
+      }
+
+      Log.d(TAG, "Requesting push challenge token…")
+      val pushSubmissionResult = RegistrationRepository.requestAndVerifyPushToken(context, session.body.id, e164, password)
+      Log.d(TAG, "Push challenge token submitted.")
+      handleSessionStateResult(context, pushSubmissionResult, errorHandler)
     }
   }
 
   /**
    * @return whether the request was successful and execution should continue
    */
-  private suspend fun handleSessionStateResult(context: Context, sessionResult: VerificationCodeRequestResult, requestedTransport: RegistrationRepository.Mode, errorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit): Boolean {
+  private suspend fun handleSessionStateResult(context: Context, sessionResult: RegistrationResult, errorHandler: (RegistrationResult) -> Unit): Boolean {
     when (sessionResult) {
       is UnknownError -> {
         handleGenericError(sessionResult.getCause())
@@ -368,9 +388,10 @@ class RegistrationV2ViewModel : ViewModel() {
         store.update {
           it.copy(
             sessionId = sessionResult.sessionId,
-            nextSms = sessionResult.nextSmsTimestamp,
+            nextSmsTimestamp = sessionResult.nextSmsTimestamp,
-            nextCall = sessionResult.nextCallTimestamp,
+            nextCallTimestamp = sessionResult.nextCallTimestamp,
-            registrationCheckpoint = RegistrationCheckpoint.VERIFICATION_CODE_REQUESTED,
+            isAllowedToRequestCode = sessionResult.allowedToRequestCode,
+            challengesRequested = emptyList(),
             inProgress = false
           )
         }
@@ -382,6 +403,7 @@ class RegistrationV2ViewModel : ViewModel() {
         store.update {
           it.copy(
             registrationCheckpoint = RegistrationCheckpoint.CHALLENGE_RECEIVED,
+            challengesRequested = sessionResult.challenges,
             inProgress = false
           )
         }
@@ -407,9 +429,13 @@ class RegistrationV2ViewModel : ViewModel() {
       is TokenNotAccepted -> Log.i(TAG, "Received TokenNotAccepted.", sessionResult.getCause())
 
       is RegistrationLocked -> Log.i(TAG, "Received RegistrationLocked.", sessionResult.getCause())
+
+      is NoSuchSession -> Log.i(TAG, "Received NoSuchSession.", sessionResult.getCause())
+
+      is AlreadyVerified -> Log.i(TAG, "Received AlreadyVerified", sessionResult.getCause())
     }
     setInProgress(false)
-    errorHandler(sessionResult, requestedTransport)
+    errorHandler(sessionResult)
     return false
   }
 
@@ -417,18 +443,23 @@ class RegistrationV2ViewModel : ViewModel() {
    * @return whether the request was successful and execution should continue
    */
   private suspend fun handleRegistrationResult(context: Context, registrationData: RegistrationData, registrationResult: RegisterAccountResult, reglockEnabled: Boolean, errorHandler: (RegisterAccountResult) -> Unit): Boolean {
+    Log.v(TAG, "handleRegistrationResult()")
     when (registrationResult) {
       is RegisterAccountResult.Success -> {
+        Log.i(TAG, "Register account result: Success!")
         onSuccessfulRegistration(context, registrationData, registrationResult.accountRegistrationResult, reglockEnabled)
         return true
       }
+
       is RegisterAccountResult.IncorrectRecoveryPassword -> {
         Log.i(TAG, "Registration recovery password was incorrect, falling back to SMS verification.", registrationResult.getCause())
         setUserSkippedReRegisterFlow(true)
       }
+
       is RegisterAccountResult.RegistrationLocked -> {
         Log.i(TAG, "Account is registration locked!", registrationResult.getCause())
       }
+
       is RegisterAccountResult.AttemptsExhausted,
       is RegisterAccountResult.RateLimited,
       is RegisterAccountResult.AuthorizationFailed,
@@ -501,6 +532,7 @@ class RegistrationV2ViewModel : ViewModel() {
           wrongPinHandler()
         } catch (noData: SvrNoDataException) {
           Log.w(TAG, "SVR has no data for these credentials. Aborting skip SMS flow.", noData)
+          updateSvrTriesRemaining(0)
           setUserSkippedReRegisterFlow(true)
         }
         setInProgress(false)
@@ -509,7 +541,6 @@ class RegistrationV2ViewModel : ViewModel() {
     }
 
     Log.w(TAG, "Could not get credentials to skip SMS registration, aborting!")
-    // TODO [regv2]: Investigate why in v1, this case throws a [IncorrectRegistrationRecoveryPasswordException], which seems weird.
     store.update {
       it.copy(canSkipSms = false, inProgress = false)
     }
@@ -518,7 +549,7 @@ class RegistrationV2ViewModel : ViewModel() {
   private suspend fun verifyReRegisterInternal(context: Context, pin: String, masterKey: MasterKey, registrationErrorHandler: (RegisterAccountResult) -> Unit) {
     updateFcmToken(context)
 
-    val registrationData = getRegistrationData("")
+    val registrationData = getRegistrationData()
 
     val resultAndRegLockStatus = registerAccountInternal(context, null, registrationData, pin, masterKey)
     val result = resultAndRegLockStatus.first
@@ -548,37 +579,42 @@ class RegistrationV2ViewModel : ViewModel() {
     return Pair(RegistrationRepository.registerAccount(context = context, sessionId = sessionId, registrationData = registrationData, pin = pin) { masterKey }, true)
   }
 
-  fun verifyCodeWithoutRegistrationLock(context: Context, code: String, submissionErrorHandler: (VerificationCodeRequestResult, RegistrationRepository.Mode) -> Unit, registrationErrorHandler: (RegisterAccountResult) -> Unit) {
+  fun verifyCodeWithoutRegistrationLock(context: Context, code: String, submissionErrorHandler: (RegistrationResult) -> Unit, registrationErrorHandler: (RegisterAccountResult) -> Unit) {
+    Log.v(TAG, "verifyCodeWithoutRegistrationLock()")
     store.update {
-      it.copy(inProgress = true, registrationCheckpoint = RegistrationCheckpoint.VERIFICATION_CODE_ENTERED)
+      it.copy(
+        inProgress = true,
+        enteredCode = code,
+        registrationCheckpoint = RegistrationCheckpoint.VERIFICATION_CODE_ENTERED
+      )
     }
 
-    val sessionId = store.value.sessionId
-    if (sessionId == null) {
-      Log.w(TAG, "Session ID was null. TODO: handle this better in the UI.")
-      return
-    }
-
-    val e164: String = getCurrentE164() ?: throw IllegalStateException()
-
     viewModelScope.launch(context = coroutineExceptionHandler) {
-      val registrationData = getRegistrationData(code)
+      verifyCodeInternal(
-      val verificationResponse = RegistrationRepository.submitVerificationCode(context, e164, password, sessionId, registrationData)
+        context = context,
-
+        pin = null,
-      if (!verificationResponse.isSuccess()) {
+        reglockEnabled = false,
-        Log.w(TAG, "Could not verify code!")
+        submissionErrorHandler = submissionErrorHandler,
-        handleSessionStateResult(context, verificationResponse, RegistrationRepository.Mode.NONE, submissionErrorHandler)
+        registrationErrorHandler = registrationErrorHandler
-        return@launch
+      )
+    }
   }
 
-      setRegistrationCheckpoint(RegistrationCheckpoint.VERIFICATION_CODE_VALIDATED)
+  private suspend fun verifyCodeInternal(context: Context, reglockEnabled: Boolean, pin: String?, submissionErrorHandler: (RegistrationResult) -> Unit, registrationErrorHandler: (RegisterAccountResult) -> Unit) {
+    val sessionId = getOrCreateValidSession(context, submissionErrorHandler)?.body?.id ?: return
+    val registrationData = getRegistrationData()
 
-      val registrationResponse: RegisterAccountResult = RegistrationRepository.registerAccount(context, sessionId, registrationData)
+    val registrationResponse = verifyCode(context, sessionId, registrationData, pin) {
-      handleRegistrationResult(context, registrationData, registrationResponse, false, registrationErrorHandler)
+      viewModelScope.launch { // TODO: validate the scopes are correct here
+        handleSessionStateResult(context, it, submissionErrorHandler)
       }
+    } ?: return
+
+    handleRegistrationResult(context, registrationData, registrationResponse, reglockEnabled, registrationErrorHandler)
   }
 
   private suspend fun onSuccessfulRegistration(context: Context, registrationData: RegistrationData, remoteResult: RegistrationRepository.AccountRegistrationResult, reglockEnabled: Boolean) {
+    Log.v(TAG, "onSuccessfulRegistration()")
     RegistrationRepository.registerAccountLocally(context, registrationData, remoteResult, reglockEnabled)
 
     refreshFeatureFlags()
@@ -622,9 +658,10 @@ class RegistrationV2ViewModel : ViewModel() {
     return store.value.phoneNumber?.toE164()
   }
 
-  private suspend fun getRegistrationData(code: String): RegistrationData {
+  private suspend fun getRegistrationData(): RegistrationData {
     val currentState = store.value
-    val e164: String = currentState.phoneNumber?.toE164() ?: throw IllegalStateException()
+    val code = currentState.enteredCode ?: throw IllegalStateException("Can't construct registration data without entered code!")
+    val e164: String = currentState.phoneNumber?.toE164() ?: throw IllegalStateException("Can't construct registration data without E164!")
     val recoveryPassword = if (currentState.sessionId == null) SignalStore.svr().getRecoveryPassword() else null
     return RegistrationData(code, e164, password, RegistrationRepository.getRegistrationId(), RegistrationRepository.getProfileKey(e164), currentState.fcmToken, RegistrationRepository.getPniRegistrationId(), recoveryPassword)
   }
@@ -649,5 +686,61 @@ class RegistrationV2ViewModel : ViewModel() {
         Log.w(TAG, "Failed to refresh flags after " + (System.currentTimeMillis() - startTime) + " ms.", e)
       }
     }
+
+    suspend fun getOrCreateValidSession(
+      context: Context,
+      existingSessionId: String?,
+      e164: String,
+      password: String,
+      mcc: String?,
+      mnc: String?,
+      successListener: (RegistrationSessionMetadataResponse) -> Unit,
+      errorHandler: (RegistrationSessionResult) -> Unit
+    ): RegistrationSessionMetadataResponse? {
+      Log.d(TAG, "Validating/creating a registration session.")
+      val sessionResult: RegistrationSessionResult = RegistrationRepository.createOrValidateSession(context, existingSessionId, e164, password, mcc, mnc)
+      when (sessionResult) {
+        is RegistrationSessionCheckResult.Success -> {
+          val metadata = sessionResult.getMetadata()
+          successListener(metadata)
+          Log.d(TAG, "Registration session validated.")
+          return metadata
+        }
+
+        is RegistrationSessionCreationResult.Success -> {
+          val metadata = sessionResult.getMetadata()
+          successListener(metadata)
+          Log.d(TAG, "Registration session created.")
+          return metadata
+        }
+
+        else -> errorHandler(sessionResult)
+      }
+      return null
+    }
+
+    suspend fun verifyCode(context: Context, sessionId: String, registrationData: RegistrationData, pin: String?, submissionErrorHandler: (RegistrationResult) -> Unit): RegisterAccountResult? {
+      Log.d(TAG, "Getting valid session in order to submit verification code.")
+
+      Log.d(TAG, "Submitting verification code…")
+
+      val verificationResponse = RegistrationRepository.submitVerificationCode(context, sessionId, registrationData)
+
+      val submissionSuccessful = verificationResponse is Success
+      val alreadyVerified = verificationResponse is AlreadyVerified
+
+      Log.i(TAG, "Verification code submission network call completed. Submission successful? $submissionSuccessful Account already verified? $alreadyVerified")
+
+      if (!submissionSuccessful && !alreadyVerified) {
+        submissionErrorHandler(verificationResponse)
+        return null
+      }
+
+      Log.d(TAG, "Submitting registration…")
+
+      val registrationResponse: RegisterAccountResult = RegistrationRepository.registerAccount(context, sessionId, registrationData, pin)
+      Log.d(TAG, "Registration network call completed.")
+      return registrationResponse
+    }
   }
 }

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/accountlocked/AccountLockedV2Fragment.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.registration.v2.ui.accountlocked
+
+import android.content.Intent
+import android.net.Uri
+import android.os.Bundle
+import android.view.View
+import android.widget.TextView
+import androidx.activity.OnBackPressedCallback
+import androidx.fragment.app.activityViewModels
+import org.thoughtcrime.securesms.LoggingFragment
+import org.thoughtcrime.securesms.R
+import org.thoughtcrime.securesms.registration.fragments.RegistrationViewDelegate.setDebugLogSubmitMultiTapView
+import org.thoughtcrime.securesms.registration.v2.ui.RegistrationV2ViewModel
+import kotlin.time.Duration.Companion.milliseconds
+
+/**
+ * Screen educating the user that they need to wait some number of days to register.
+ */
+class AccountLockedV2Fragment : LoggingFragment(R.layout.account_locked_fragment) {
+  private val viewModel by activityViewModels<RegistrationV2ViewModel>()
+  override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
+    super.onViewCreated(view, savedInstanceState)
+
+    setDebugLogSubmitMultiTapView(view.findViewById(R.id.account_locked_title))
+
+    val description = view.findViewById<TextView>(R.id.account_locked_description)
+
+    viewModel.lockedTimeRemaining.observe(
+      viewLifecycleOwner
+    ) { t: Long? -> description.text = getString(R.string.AccountLockedFragment__your_account_has_been_locked_to_protect_your_privacy, durationToDays(t!!)) }
+
+    view.findViewById<View>(R.id.account_locked_next).setOnClickListener { v: View? -> onNext() }
+    view.findViewById<View>(R.id.account_locked_learn_more).setOnClickListener { v: View? -> learnMore() }
+
+    requireActivity().onBackPressedDispatcher.addCallback(
+      viewLifecycleOwner,
+      object : OnBackPressedCallback(true) {
+        override fun handleOnBackPressed() {
+          onNext()
+        }
+      }
+    )
+  }
+
+  private fun learnMore() {
+    val intent = Intent(Intent.ACTION_VIEW)
+    intent.setData(Uri.parse(getString(R.string.AccountLockedFragment__learn_more_url)))
+    startActivity(intent)
+  }
+
+  fun onNext() {
+    requireActivity().finish()
+  }
+
+  private fun durationToDays(duration: Long): Long {
+    return if (duration != 0L) getLockoutDays(duration).toLong() else 7
+  }
+
+  private fun getLockoutDays(timeRemainingMs: Long): Int {
+    return timeRemainingMs.milliseconds.inWholeDays.toInt() + 1
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/captcha/CaptchaV2Fragment.kt

@@ -10,7 +10,8 @@ import android.os.Bundle
 import android.view.View
 import android.webkit.WebView
 import android.webkit.WebViewClient
-import androidx.fragment.app.activityViewModels
+import androidx.activity.OnBackPressedCallback
+import androidx.activity.addCallback
 import androidx.navigation.fragment.findNavController
 import org.thoughtcrime.securesms.BuildConfig
 import org.thoughtcrime.securesms.LoggingFragment
@@ -18,13 +19,17 @@ import org.thoughtcrime.securesms.R
 import org.thoughtcrime.securesms.components.ViewBinderDelegate
 import org.thoughtcrime.securesms.databinding.FragmentRegistrationCaptchaV2Binding
 import org.thoughtcrime.securesms.registration.fragments.RegistrationConstants
-import org.thoughtcrime.securesms.registration.v2.ui.RegistrationV2ViewModel
 
-class CaptchaFragment : LoggingFragment(R.layout.fragment_registration_captcha_v2) {
+abstract class CaptchaV2Fragment : LoggingFragment(R.layout.fragment_registration_captcha_v2) {
 
-  private val sharedViewModel by activityViewModels<RegistrationV2ViewModel>()
   private val binding: FragmentRegistrationCaptchaV2Binding by ViewBinderDelegate(FragmentRegistrationCaptchaV2Binding::bind)
 
+  private val backListener = object : OnBackPressedCallback(true) {
+    override fun handleOnBackPressed() {
+      handleUserExit()
+    }
+  }
+
   @SuppressLint("SetJavaScriptEnabled")
   override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
     super.onViewCreated(view, savedInstanceState)
@@ -36,14 +41,21 @@ class CaptchaFragment : LoggingFragment(R.layout.fragment_registration_captcha_v
       override fun shouldOverrideUrlLoading(view: WebView, url: String): Boolean {
         if (url.startsWith(RegistrationConstants.SIGNAL_CAPTCHA_SCHEME)) {
           val token = url.substring(RegistrationConstants.SIGNAL_CAPTCHA_SCHEME.length)
-          sharedViewModel.setCaptchaResponse(token)
+          handleCaptchaToken(token)
+          backListener.isEnabled = false
           findNavController().navigateUp()
           return true
         }
         return false
       }
     }
-
+    requireActivity().onBackPressedDispatcher.addCallback(viewLifecycleOwner) {
+      handleUserExit()
+    }
     binding.registrationCaptchaWebView.loadUrl(BuildConfig.SIGNAL_CAPTCHA_URL)
   }
+
+  abstract fun handleCaptchaToken(token: String)
+
+  abstract fun handleUserExit()
 }

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/captcha/RegistrationCaptchaV2Fragment.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.registration.v2.ui.captcha
+
+import android.os.Bundle
+import android.view.View
+import androidx.fragment.app.activityViewModels
+import org.thoughtcrime.securesms.registration.v2.data.network.Challenge
+import org.thoughtcrime.securesms.registration.v2.ui.RegistrationV2ViewModel
+
+/**
+ * Screen that displays a captcha as part of the registration flow.
+ * This subclass plugs in [RegistrationV2ViewModel] to the shared super class.
+ *
+ * @see CaptchaV2Fragment
+ */
+class RegistrationCaptchaV2Fragment : CaptchaV2Fragment() {
+  private val sharedViewModel by activityViewModels<RegistrationV2ViewModel>()
+  override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
+    super.onViewCreated(view, savedInstanceState)
+    sharedViewModel.addPresentedChallenge(Challenge.CAPTCHA)
+  }
+
+  override fun handleCaptchaToken(token: String) {
+    sharedViewModel.setCaptchaResponse(token)
+  }
+
+  override fun handleUserExit() {
+    sharedViewModel.removePresentedChallenge(Challenge.CAPTCHA)
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/entercode/EnterCodeV2Fragment.kt

@@ -21,8 +21,8 @@ import org.thoughtcrime.securesms.databinding.FragmentRegistrationEnterCodeV2Bin
 import org.thoughtcrime.securesms.registration.fragments.ContactSupportBottomSheetFragment
 import org.thoughtcrime.securesms.registration.fragments.RegistrationViewDelegate.setDebugLogSubmitMultiTapView
 import org.thoughtcrime.securesms.registration.fragments.SignalStrengthPhoneStateListener
-import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
 import org.thoughtcrime.securesms.registration.v2.data.network.RegisterAccountResult
+import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationResult
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult
 import org.thoughtcrime.securesms.registration.v2.ui.RegistrationCheckpoint
 import org.thoughtcrime.securesms.registration.v2.ui.RegistrationV2ViewModel
@@ -102,8 +102,8 @@ class EnterCodeV2Fragment : LoggingFragment(R.layout.fragment_registration_enter
     }
 
     sharedViewModel.uiState.observe(viewLifecycleOwner) {
-      binding.resendSmsCountDown.startCountDownTo(it.nextSms)
+      binding.resendSmsCountDown.startCountDownTo(it.nextSmsTimestamp)
-      binding.callMeCountDown.startCountDownTo(it.nextCall)
+      binding.callMeCountDown.startCountDownTo(it.nextCallTimestamp)
       if (it.inProgress) {
         binding.keyboard.displayProgress()
       } else {
@@ -112,66 +112,35 @@ class EnterCodeV2Fragment : LoggingFragment(R.layout.fragment_registration_enter
     }
   }
 
-  private fun handleSessionErrorResponse(requestResult: VerificationCodeRequestResult, mode: RegistrationRepository.Mode) {
+  private fun handleSessionErrorResponse(result: RegistrationResult) {
-    when (requestResult) {
+    when (result) {
       is VerificationCodeRequestResult.Success -> binding.keyboard.displaySuccess()
-      is VerificationCodeRequestResult.RateLimited -> {
+      is VerificationCodeRequestResult.RateLimited -> presentRateLimitedDialog()
-        binding.keyboard.displayFailure().addListener(
+      is VerificationCodeRequestResult.RegistrationLocked -> presentRegistrationLocked(result.timeRemaining)
-          object : AssertedSuccessListener<Boolean>() {
+      else -> presentGenericError(result)
-            override fun onSuccess(result: Boolean?) {
-              presentRemoteErrorDialog(getString(R.string.RegistrationActivity_you_have_made_too_many_attempts_please_try_again_later)) { _, _ ->
-                binding.code.clear()
-              }
-            }
-          }
-        )
-      }
-
-      is VerificationCodeRequestResult.RegistrationLocked -> {
-        binding.keyboard.displayLocked().addListener(
-          object : AssertedSuccessListener<Boolean>() {
-            override fun onSuccess(result: Boolean?) {
-              findNavController().safeNavigate(EnterCodeV2FragmentDirections.actionRequireKbsLockPin(requestResult.timeRemaining))
-            }
-          }
-        )
-      }
-
-      else -> {
-        binding.keyboard.displayFailure().addListener(
-          object : AssertedSuccessListener<Boolean>() {
-            override fun onSuccess(result: Boolean?) {
-              Log.w(TAG, "Encountered unexpected error!", requestResult.getCause())
-              presentRemoteErrorDialog(getString(R.string.RegistrationActivity_error_connecting_to_service))
-            }
-          }
-        )
-      }
     }
   }
 
   private fun handleRegistrationErrorResponse(result: RegisterAccountResult) {
     when (result) {
-      is RegisterAccountResult.Success -> Log.d(TAG, "Register account was successful.")
+      is RegisterAccountResult.Success -> binding.keyboard.displaySuccess()
-      is RegisterAccountResult.AuthorizationFailed -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_error_connecting_to_service))
+      is RegisterAccountResult.RegistrationLocked -> presentRegistrationLocked(result.timeRemaining)
-      is RegisterAccountResult.MalformedRequest -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_error_connecting_to_service))
-      is RegisterAccountResult.RegistrationLocked -> {
-        Log.w(TAG, "Account is registration locked, cannot register.")
-        findNavController().safeNavigate(EnterCodeV2FragmentDirections.actionRequireKbsLockPin(result.timeRemaining))
-      }
-      is RegisterAccountResult.UnknownError -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_error_connecting_to_service))
-      is RegisterAccountResult.ValidationError -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_error_connecting_to_service))
-      is RegisterAccountResult.IncorrectRecoveryPassword -> {
-        Log.w(TAG, "User somehow got recovery password error while entering code. This is very suspicious!")
-        sharedViewModel.setUserSkippedReRegisterFlow(true)
-        popBackStack()
-      }
-
       is RegisterAccountResult.AttemptsExhausted,
       is RegisterAccountResult.RateLimited -> presentRateLimitedDialog()
+      else -> presentGenericError(result)
     }
   }
 
+  private fun presentRegistrationLocked(timeRemaining: Long) {
+    binding.keyboard.displayLocked().addListener(
+      object : AssertedSuccessListener<Boolean>() {
+        override fun onSuccess(result: Boolean?) {
+          findNavController().safeNavigate(EnterCodeV2FragmentDirections.actionRequireKbsLockPin(timeRemaining))
+        }
+      }
+    )
+  }
+
   private fun presentRateLimitedDialog() {
     binding.keyboard.displayFailure().addListener(
       object : AssertedSuccessListener<Boolean?>() {
@@ -193,16 +162,23 @@ class EnterCodeV2Fragment : LoggingFragment(R.layout.fragment_registration_enter
     )
   }
 
-  private fun presentRemoteErrorDialog(message: String, title: String? = null, positiveButtonListener: DialogInterface.OnClickListener? = null) {
+  private fun presentGenericError(requestResult: RegistrationResult) {
+    binding.keyboard.displayFailure().addListener(
+      object : AssertedSuccessListener<Boolean>() {
+        override fun onSuccess(result: Boolean?) {
+          Log.w(TAG, "Encountered unexpected error!", requestResult.getCause())
           MaterialAlertDialogBuilder(requireContext()).apply {
-      title?.let {
+            null?.let<String, MaterialAlertDialogBuilder> {
               setTitle(it)
             }
-      setMessage(message)
+            setMessage(getString(R.string.RegistrationActivity_error_connecting_to_service))
-      setPositiveButton(android.R.string.ok, positiveButtonListener)
+            setPositiveButton(android.R.string.ok, null)
             show()
           }
         }
+      }
+    )
+  }
 
   private fun popBackStack() {
     sharedViewModel.setRegistrationCheckpoint(RegistrationCheckpoint.PUSH_NETWORK_AUDITED)

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/phonenumber/EnterPhoneNumberV2Fragment.kt

@@ -34,6 +34,7 @@ import com.google.android.material.textfield.TextInputEditText
 import com.google.i18n.phonenumbers.NumberParseException
 import com.google.i18n.phonenumbers.PhoneNumberUtil
 import com.google.i18n.phonenumbers.Phonenumber.PhoneNumber
+import org.signal.core.util.isNotNullOrBlank
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.LoggingFragment
 import org.thoughtcrime.securesms.R
@@ -43,6 +44,9 @@ import org.thoughtcrime.securesms.phonenumbers.PhoneNumberFormatter
 import org.thoughtcrime.securesms.registration.fragments.RegistrationViewDelegate.setDebugLogSubmitMultiTapView
 import org.thoughtcrime.securesms.registration.util.CountryPrefix
 import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
+import org.thoughtcrime.securesms.registration.v2.data.network.Challenge
+import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationResult
+import org.thoughtcrime.securesms.registration.v2.data.network.RegistrationSessionCreationResult
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult
 import org.thoughtcrime.securesms.registration.v2.ui.RegistrationCheckpoint
 import org.thoughtcrime.securesms.registration.v2.ui.RegistrationV2State
@@ -111,12 +115,14 @@ class EnterPhoneNumberV2Fragment : LoggingFragment(R.layout.fragment_registratio
         presentNetworkError(it)
       }
 
-      if (sharedState.registrationCheckpoint >= RegistrationCheckpoint.PHONE_NUMBER_CONFIRMED && sharedState.canSkipSms) {
+      if (sharedState.challengesRequested.contains(Challenge.CAPTCHA) && sharedState.captchaToken.isNotNullOrBlank()) {
+        sharedViewModel.submitCaptchaToken(requireContext(), ::handleErrorResponse)
+      } else if (sharedState.challengesRemaining.isNotEmpty()) {
+        handleChallenges(sharedState.challengesRemaining)
+      } else if (sharedState.registrationCheckpoint >= RegistrationCheckpoint.PHONE_NUMBER_CONFIRMED && sharedState.canSkipSms) {
         moveToEnterPinScreen()
       } else if (sharedState.registrationCheckpoint >= RegistrationCheckpoint.VERIFICATION_CODE_REQUESTED) {
         moveToVerificationEntryScreen()
-      } else if (sharedState.registrationCheckpoint >= RegistrationCheckpoint.CHALLENGE_COMPLETED) {
-        sharedViewModel.submitCaptchaToken(requireContext(), ::handleErrorResponse)
       }
     }
 
@@ -160,6 +166,17 @@ class EnterPhoneNumberV2Fragment : LoggingFragment(R.layout.fragment_registratio
     ViewUtil.focusAndShowKeyboard(phoneNumberInputLayout)
   }
 
+  private fun handleChallenges(remainingChallenges: List<Challenge>) {
+    when (remainingChallenges.first()) {
+      Challenge.CAPTCHA -> moveToCaptcha()
+      Challenge.PUSH -> performPushChallenge()
+    }
+  }
+
+  private fun performPushChallenge() {
+    sharedViewModel.requestAndSubmitPushToken(requireContext(), ::handleErrorResponse)
+  }
+
   private fun initializeInputFields() {
     binding.countryCode.editText?.addTextChangedListener { s ->
       val countryCode: Int = s.toString().toInt()
@@ -263,11 +280,18 @@ class EnterPhoneNumberV2Fragment : LoggingFragment(R.layout.fragment_registratio
     }
   }
 
-  private fun handleErrorResponse(result: VerificationCodeRequestResult, mode: RegistrationRepository.Mode) {
+  private fun handleErrorResponse(result: RegistrationResult) {
+    if (!result.isSuccess()) {
+      Log.i(TAG, "Handling error response.", result.getCause())
+    }
     when (result) {
+      is RegistrationSessionCreationResult.Success,
       is VerificationCodeRequestResult.Success -> Unit
+      is RegistrationSessionCreationResult.AttemptsExhausted,
       is VerificationCodeRequestResult.AttemptsExhausted -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_rate_limited_to_service))
-      is VerificationCodeRequestResult.ChallengeRequired -> findNavController().safeNavigate(EnterPhoneNumberV2FragmentDirections.actionRequestCaptcha())
+      is VerificationCodeRequestResult.ChallengeRequired -> {
+        handleChallenges(result.challenges)
+      }
       is VerificationCodeRequestResult.ExternalServiceFailure -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_unable_to_connect_to_service), skipToNextScreen)
       is VerificationCodeRequestResult.ImpossibleNumber -> {
         MaterialAlertDialogBuilder(requireContext()).apply {
@@ -286,15 +310,21 @@ class EnterPhoneNumberV2Fragment : LoggingFragment(R.layout.fragment_registratio
           show()
         }
       }
+      is RegistrationSessionCreationResult.MalformedRequest,
       is VerificationCodeRequestResult.MalformedRequest -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_unable_to_connect_to_service), skipToNextScreen)
       is VerificationCodeRequestResult.MustRetry -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_unable_to_connect_to_service), skipToNextScreen)
-      is VerificationCodeRequestResult.NonNormalizedNumber -> handleNonNormalizedNumberError(result.originalNumber, result.normalizedNumber, mode)
+      is VerificationCodeRequestResult.NonNormalizedNumber -> handleNonNormalizedNumberError(result.originalNumber, result.normalizedNumber, fragmentViewModel.mode)
+      is RegistrationSessionCreationResult.RateLimited -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_rate_limited_to_try_again, result.timeRemaining.milliseconds.toString()))
       is VerificationCodeRequestResult.RateLimited -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_rate_limited_to_try_again, result.timeRemaining.milliseconds.toString()))
-      is VerificationCodeRequestResult.TokenNotAccepted -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_we_need_to_verify_that_youre_human))
+      is VerificationCodeRequestResult.TokenNotAccepted -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_we_need_to_verify_that_youre_human)) { _, _ -> moveToCaptcha() }
       else -> presentRemoteErrorDialog(getString(R.string.RegistrationActivity_unable_to_connect_to_service))
     }
   }
 
+  private fun moveToCaptcha() {
+    findNavController().safeNavigate(EnterPhoneNumberV2FragmentDirections.actionRequestCaptcha())
+  }
+
   private fun presentRemoteErrorDialog(message: String, positiveButtonListener: DialogInterface.OnClickListener? = null) {
     MaterialAlertDialogBuilder(requireContext()).apply {
       setMessage(message)
@@ -325,7 +355,6 @@ class EnterPhoneNumberV2Fragment : LoggingFragment(R.layout.fragment_registratio
             RegistrationRepository.Mode.SMS_WITH_LISTENER,
             RegistrationRepository.Mode.SMS_WITHOUT_LISTENER -> sharedViewModel.requestSmsCode(requireContext(), ::handleErrorResponse)
             RegistrationRepository.Mode.PHONE_CALL -> sharedViewModel.requestVerificationCall(requireContext(), ::handleErrorResponse)
-            RegistrationRepository.Mode.NONE -> Log.w(TAG, "Somehow got a non normalized number exception even though we didn't request a code.")
           }
           dialogInterface.dismiss()
         }

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/phonenumber/EnterPhoneNumberV2State.kt

@@ -6,17 +6,12 @@
 package org.thoughtcrime.securesms.registration.v2.ui.phonenumber
 
 import android.text.TextWatcher
+import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
 
 /**
  * State holder for the phone number entry screen, including phone number and Play Services errors.
  */
-data class EnterPhoneNumberV2State(val countryPrefixIndex: Int, val phoneNumber: String, val phoneNumberFormatter: TextWatcher? = null, val error: Error = Error.NONE) {
+data class EnterPhoneNumberV2State(val countryPrefixIndex: Int = 1, val phoneNumber: String = "", val phoneNumberFormatter: TextWatcher? = null, val mode: RegistrationRepository.Mode = RegistrationRepository.Mode.SMS_WITHOUT_LISTENER, val error: Error = Error.NONE) {
-
-  companion object {
-    @JvmStatic
-    val INIT = EnterPhoneNumberV2State(1, "")
-  }
-
   enum class Error {
     NONE,
     INVALID_PHONE_NUMBER,

app/src/main/java/org/thoughtcrime/securesms/registration/v2/ui/phonenumber/EnterPhoneNumberV2ViewModel.kt

@@ -19,6 +19,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.registration.util.CountryPrefix
+import org.thoughtcrime.securesms.registration.v2.data.RegistrationRepository
 
 /**
  * ViewModel for the phone number entry screen.
@@ -27,13 +28,19 @@ class EnterPhoneNumberV2ViewModel : ViewModel() {
 
   private val TAG = Log.tag(EnterPhoneNumberV2ViewModel::class.java)
 
-  private val store = MutableStateFlow(EnterPhoneNumberV2State.INIT)
+  private val store = MutableStateFlow(EnterPhoneNumberV2State())
   val uiState = store.asLiveData()
 
   val supportedCountryPrefixes: List<CountryPrefix> = PhoneNumberUtil.getInstance().supportedCallingCodes
     .map { CountryPrefix(it, PhoneNumberUtil.getInstance().getRegionCodeForCountryCode(it)) }
     .sortedBy { it.digits }
 
+  var mode: RegistrationRepository.Mode
+    get() = store.value.mode
+    set(value) = store.update {
+      it.copy(mode = value)
+    }
+
   fun countryPrefix(): CountryPrefix {
     return supportedCountryPrefixes[store.value.countryPrefixIndex]
   }