Fleshed out session management in registration v2.

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

@@ -42,6 +42,9 @@ import org.thoughtcrime.securesms.registration.PushChallengeRequest
 import org.thoughtcrime.securesms.registration.RegistrationData
 import org.thoughtcrime.securesms.registration.VerifyAccountRepository
 import org.thoughtcrime.securesms.registration.v2.data.network.BackupAuthCheckResult
+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.viewmodel.SvrAuthCredentialSet
 import org.thoughtcrime.securesms.service.DirectoryRefreshListener
@@ -254,13 +257,19 @@ object RegistrationRepository {
     }
 
   /**
-   * Asks the service to send a verification code through one of our supported channels (SMS, phone call).
-   * This requires two or more network calls:
-   * 1. Create (or reuse) a session.
-   * 2. (Optional) If the session has any proof requirements ("challenges"), the user must solve them and submit the proof.
-   * 3. Once the service responds we are allowed to, we request the verification code.
+   * Validates a session ID.
    */
-  suspend fun requestSmsCode(context: Context, e164: String, password: String, mcc: String?, mnc: String?, mode: Mode = Mode.SMS_WITHOUT_LISTENER): VerificationCodeRequestResult =
+  suspend fun validateSession(context: Context, sessionId: String, e164: String, password: String): RegistrationSessionCheckResult =
+    withContext(Dispatchers.IO) {
+      val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password).registrationApi
+      val registrationSessionResult = api.getRegistrationSessionStatus(sessionId)
+      return@withContext RegistrationSessionCheckResult.from(registrationSessionResult)
+    }
+
+  /**
+   * Initiates a new registration session on the service.
+   */
+  suspend fun createSession(context: Context, e164: String, password: String, mcc: String?, mnc: String?): RegistrationSessionCreationResult =
     withContext(Dispatchers.IO) {
       val fcmToken: String? = FcmUtil.getToken(context).orElse(null)
       val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password).registrationApi
@@ -270,17 +279,46 @@ object RegistrationRepository {
       } else {
         createSessionAndBlockForPushChallenge(api, fcmToken, mcc, mnc)
       }
-      val session = registrationSessionResult.successOrThrow()
-      val sessionId = session.body.id
-      SignalStore.registrationValues().sessionId = sessionId
-      SignalStore.registrationValues().sessionE164 = e164
-      if (!session.body.allowedToRequestCode) {
-        val challenges = session.body.requestedInformation.joinToString()
-        Log.w(TAG, "Not allowed to request code! Remaining challenges: $challenges")
-        return@withContext VerificationCodeRequestResult.from(registrationSessionResult)
+      val result = RegistrationSessionCreationResult.from(registrationSessionResult)
+      if (result is RegistrationSessionCreationResult.Success) {
+        SignalStore.registrationValues().sessionId = result.getMetadata().body.id
+        SignalStore.registrationValues().sessionE164 = e164
       }
+
+      return@withContext result
+    }
+
+  /**
+   * 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 sessionValidationResult = validateSession(context, sessionId, e164, password)
+      when (sessionValidationResult) {
+        is RegistrationSessionCheckResult.Success -> return sessionValidationResult
+        is RegistrationSessionCheckResult.UnknownError -> {
+          Log.w(TAG, "Encountered error when validating existing session.", sessionValidationResult.getCause())
+          return sessionValidationResult
+        }
+
+        is RegistrationSessionCheckResult.SessionNotFound -> {
+          Log.i(TAG, "Current session is invalid or has expired. Must create new one.")
+          // fall through to creation
+        }
+      }
+    }
+    return createSession(context, e164, password, mcc, mnc)
+  }
+
+  /**
+   * Asks the service to send a verification code through one of our supported channels (SMS, phone call).
+   */
+  suspend fun requestSmsCode(context: Context, sessionId: String, e164: String, password: String, mode: Mode = Mode.SMS_WITHOUT_LISTENER): VerificationCodeRequestResult =
+    withContext(Dispatchers.IO) {
+      val api: RegistrationApi = AccountManagerFactory.getInstance().createUnauthenticated(context, e164, SignalServiceAddress.DEFAULT_DEVICE_ID, password).registrationApi
+
       // TODO [regv2]: support other verification code [Mode] options
-      if (mode == Mode.PHONE_CALL) {
+      val codeRequestResult = if (mode == Mode.PHONE_CALL) {
         // TODO [regv2]
         val notImplementedError = NotImplementedError()
         Log.w(TAG, "Not yet implemented!", notImplementedError)
@@ -289,7 +327,7 @@ object RegistrationRepository {
         api.requestSmsVerificationCode(sessionId, Locale.getDefault(), mode.isSmsRetrieverSupported)
       }
 
-      return@withContext VerificationCodeRequestResult.from(registrationSessionResult)
+      return@withContext VerificationCodeRequestResult.from(codeRequestResult)
     }
 
   /**
@@ -362,7 +400,7 @@ object RegistrationRepository {
         }
     }
 
-  private suspend fun createSessionAndBlockForPushChallenge(accountManager: RegistrationApi, fcmToken: String, mcc: String?, mnc: String?): NetworkResult<RegistrationSessionMetadataResponse> =
+  suspend fun createSessionAndBlockForPushChallenge(accountManager: RegistrationApi, fcmToken: String, mcc: String?, mnc: String?): NetworkResult<RegistrationSessionMetadataResponse> =
     withContext(Dispatchers.IO) {
       // TODO [regv2]: do not use event bus nor latch
       val subscriber = PushTokenChallengeSubscriber()

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

@@ -0,0 +1,89 @@
+/*
+ * 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
+import org.whispersystems.signalservice.api.NetworkResult
+import org.whispersystems.signalservice.api.push.exceptions.MalformedRequestException
+import org.whispersystems.signalservice.api.push.exceptions.NotFoundException
+import org.whispersystems.signalservice.api.push.exceptions.RateLimitException
+import org.whispersystems.signalservice.internal.push.RegistrationSessionMetadataResponse
+
+sealed class RegistrationSessionResult(cause: Throwable?) : RegistrationResult(cause)
+
+interface SessionMetadataHolder {
+  fun getMetadata(): RegistrationSessionMetadataResponse
+}
+
+sealed class RegistrationSessionCreationResult(cause: Throwable?) : RegistrationSessionResult(cause) {
+  companion object {
+
+    private val TAG = Log.tag(RegistrationSessionResult::class.java)
+
+    @JvmStatic
+    fun from(networkResult: NetworkResult<RegistrationSessionMetadataResponse>): RegistrationSessionCreationResult {
+      return when (networkResult) {
+        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 MalformedRequestException -> MalformedRequest(cause)
+            else -> if (networkResult.code == 422) {
+              ServerUnableToParse(cause)
+            } else {
+              UnknownError(cause)
+            }
+          }
+        }
+      }
+    }
+  }
+
+  class Success(private val metadata: RegistrationSessionMetadataResponse) : RegistrationSessionCreationResult(null), SessionMetadataHolder {
+    override fun getMetadata(): RegistrationSessionMetadataResponse {
+      return metadata
+    }
+  }
+
+  class RateLimited(cause: Throwable) : RegistrationSessionCreationResult(cause)
+  class ServerUnableToParse(cause: Throwable) : RegistrationSessionCreationResult(cause)
+  class MalformedRequest(cause: Throwable) : RegistrationSessionCreationResult(cause)
+  class UnknownError(cause: Throwable) : RegistrationSessionCreationResult(cause)
+}
+
+sealed class RegistrationSessionCheckResult(cause: Throwable?) : RegistrationSessionResult(cause) {
+  companion object {
+    fun from(networkResult: NetworkResult<RegistrationSessionMetadataResponse>): RegistrationSessionCheckResult {
+      return when (networkResult) {
+        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 NotFoundException -> SessionNotFound(cause)
+            else -> UnknownError(cause)
+          }
+        }
+      }
+    }
+  }
+
+  class Success(private val metadata: RegistrationSessionMetadataResponse) : RegistrationSessionCheckResult(null), SessionMetadataHolder {
+    override fun getMetadata(): RegistrationSessionMetadataResponse {
+      return metadata
+    }
+  }
+
+  class SessionNotFound(cause: Throwable) : RegistrationSessionCheckResult(cause)
+  class UnknownError(cause: Throwable) : RegistrationSessionCheckResult(cause)
+}

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

@@ -18,6 +18,7 @@ import kotlinx.coroutines.flow.MutableStateFlow
 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.ApplicationDependencies
 import org.thoughtcrime.securesms.jobs.MultiDeviceProfileContentUpdateJob
@@ -29,6 +30,9 @@ 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.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.AttemptsExhausted
 import org.thoughtcrime.securesms.registration.v2.data.network.VerificationCodeRequestResult.ChallengeRequired
@@ -49,6 +53,7 @@ import org.whispersystems.signalservice.api.NetworkResult
 import org.whispersystems.signalservice.api.SvrNoDataException
 import org.whispersystems.signalservice.api.kbs.MasterKey
 import org.whispersystems.signalservice.internal.push.LockedException
+import org.whispersystems.signalservice.internal.push.RegistrationSessionMetadataResponse
 import java.io.IOException
 
 /**
@@ -173,12 +178,63 @@ class RegistrationV2ViewModel : ViewModel() {
         is BackupAuthCheckResult.SuccessWithoutCredentials -> Log.d(TAG, "No local SVR auth credentials could be found and/or validated.")
       }
 
-      val codeRequestResponse = RegistrationRepository.requestSmsCode(context, e164, password, mccMncProducer.mcc, mccMncProducer.mnc)
+      val validSession = getOrCreateValidSession(context) ?: 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))
+        return@launch
+      }
+
+      val codeRequestResponse = RegistrationRepository.requestSmsCode(context, validSession.body.id, e164, password)
 
       handleSessionStateResult(context, codeRequestResponse)
     }
   }
 
+  private suspend fun getOrCreateValidSession(context: Context): RegistrationSessionMetadataResponse? {
+    val e164 = getCurrentE164() ?: throw IllegalStateException("E164 required to create session!")
+    val mccMncProducer = MccMncProducer(context)
+
+    val existingSessionId = store.value.sessionId
+    val sessionResult: RegistrationSessionResult = RegistrationRepository.createOrValidateSession(context, existingSessionId, e164, password, mccMncProducer.mcc, mccMncProducer.mnc)
+    when (sessionResult) {
+      is RegistrationSessionCheckResult.Success -> {
+        val metadata = sessionResult.getMetadata()
+        val newSessionId = metadata.body.id
+        if (newSessionId.isNotNullOrBlank() && newSessionId != existingSessionId) {
+          store.update {
+            it.copy(
+              sessionId = newSessionId
+            )
+          }
+        }
+        return metadata
+      }
+      is RegistrationSessionCreationResult.Success -> {
+        val metadata = sessionResult.getMetadata()
+        val newSessionId = metadata.body.id
+        if (newSessionId.isNotNullOrBlank() && newSessionId != existingSessionId) {
+          store.update {
+            it.copy(
+              sessionId = newSessionId
+            )
+          }
+        }
+        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())
+      is RegistrationSessionCheckResult.UnknownError -> Log.i(TAG, "Unknown error occurred while checking registration session.", sessionResult.getCause())
+      is RegistrationSessionCreationResult.MalformedRequest -> Log.i(TAG, "Malformed request error occurred while creating registration session.", sessionResult.getCause())
+      is RegistrationSessionCreationResult.RateLimited -> Log.i(TAG, "Rate limit occurred while creating registration session.", sessionResult.getCause())
+      is RegistrationSessionCreationResult.ServerUnableToParse -> Log.i(TAG, "Server unable to parse request for creating registration session.", sessionResult.getCause())
+      is RegistrationSessionCreationResult.UnknownError -> Log.i(TAG, "Unknown error occurred while checking registration session.", sessionResult.getCause())
+    }
+    setInProgress(false)
+    return null
+  }
+
   fun submitCaptchaToken(context: Context) {
     val e164 = getCurrentE164() ?: throw IllegalStateException("TODO")
     val sessionId = store.value.sessionId ?: throw IllegalStateException("TODO")
@@ -216,6 +272,7 @@ class RegistrationV2ViewModel : ViewModel() {
 
       is AttemptsExhausted -> Log.w(TAG, "TODO")
       is ChallengeRequired -> store.update {
+        // TODO [regv2] handle push challenge required
         it.copy(
           registrationCheckpoint = RegistrationCheckpoint.CHALLENGE_RECEIVED
         )

libsignal-service/src/main/java/org/whispersystems/signalservice/api/registration/RegistrationApi.kt

@@ -30,6 +30,15 @@ class RegistrationApi(
     }
   }
 
+  /**
+   * Retrieve current status of a registration session.
+   */
+  fun getRegistrationSessionStatus(sessionId: String): NetworkResult<RegistrationSessionMetadataResponse> {
+    return NetworkResult.fromFetch {
+      pushServiceSocket.getSessionStatus(sessionId)
+    }
+  }
+
   /**
    * Submit an FCM token to the service as proof that this is an honest user attempting to register.
    */