chore(auth): Refactor fetchAuthSession to useCase (#3017)

Author: mattcreaser

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPlugin.kt

@@ -176,7 +176,7 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
             logger
         )
 
-        useCaseFactory = AuthUseCaseFactory(realPlugin, authEnvironment, authStateMachine)
+        useCaseFactory = AuthUseCaseFactory(authEnvironment, authStateMachine)
 
         blockQueueChannelWhileConfiguring()
     }
@@ -291,10 +291,10 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
         options: AuthFetchSessionOptions,
         onSuccess: Consumer<AuthSession>,
         onError: Consumer<AuthException>
-    ) = enqueue(onSuccess, onError) { queueFacade.fetchAuthSession(options) }
+    ) = enqueue(onSuccess, onError) { useCaseFactory.fetchAuthSession().execute(options) }
 
     override fun fetchAuthSession(onSuccess: Consumer<AuthSession>, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.fetchAuthSession() }
+        enqueue(onSuccess, onError) { useCaseFactory.fetchAuthSession().execute() }
 
     override fun rememberDevice(onSuccess: Action, onError: Consumer<AuthException>) =
         enqueue(onSuccess, onError) { useCaseFactory.rememberDevice().execute() }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/AWSCognitoAuthSession.kt

@@ -66,9 +66,7 @@ internal fun AmplifyCredential.isValid(): Boolean = when (this) {
     else -> false
 }
 
-internal fun AmplifyCredential.getCognitoSession(
-    exception: AuthException? = null
-): AWSAuthSessionBehavior<AWSCognitoUserPoolTokens> {
+internal fun AmplifyCredential.getCognitoSession(exception: AuthException? = null): AWSCognitoAuthSession {
     fun getCredentialsResult(
         awsCredentials: CognitoCredentials,
         exception: AuthException?

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/KotlinAuthFacadeInternal.kt

@@ -18,11 +18,9 @@ package com.amplifyframework.auth.cognito
 import android.app.Activity
 import android.content.Intent
 import com.amplifyframework.auth.AuthProvider
-import com.amplifyframework.auth.AuthSession
 import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
 import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
 import com.amplifyframework.auth.options.AuthConfirmSignInOptions
-import com.amplifyframework.auth.options.AuthFetchSessionOptions
 import com.amplifyframework.auth.options.AuthSignInOptions
 import com.amplifyframework.auth.options.AuthSignOutOptions
 import com.amplifyframework.auth.options.AuthWebUISignInOptions
@@ -118,21 +116,6 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
         delegate.handleWebUISignInResponse(intent)
     }
 
-    suspend fun fetchAuthSession(): AuthSession = suspendCoroutine { continuation ->
-        delegate.fetchAuthSession(
-            { continuation.resume(it) },
-            { continuation.resumeWithException(it) }
-        )
-    }
-
-    suspend fun fetchAuthSession(options: AuthFetchSessionOptions): AuthSession = suspendCoroutine { continuation ->
-        delegate.fetchAuthSession(
-            options,
-            { continuation.resume(it) },
-            { continuation.resumeWithException(it) }
-        )
-    }
-
     suspend fun signOut(): AuthSignOutResult = suspendCoroutine { continuation ->
         delegate.signOut { continuation.resume(it) }
     }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt

@@ -29,13 +29,11 @@ import com.amplifyframework.auth.AuthCodeDeliveryDetails
 import com.amplifyframework.auth.AuthException
 import com.amplifyframework.auth.AuthFactorType
 import com.amplifyframework.auth.AuthProvider
-import com.amplifyframework.auth.AuthSession
 import com.amplifyframework.auth.MFAType
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidOauthConfigurationException
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidUserPoolConfigurationException
 import com.amplifyframework.auth.cognito.exceptions.invalidstate.SignedInException
 import com.amplifyframework.auth.cognito.exceptions.service.HostedUISignOutException
-import com.amplifyframework.auth.cognito.exceptions.service.InvalidAccountTypeException
 import com.amplifyframework.auth.cognito.exceptions.service.InvalidParameterException
 import com.amplifyframework.auth.cognito.exceptions.service.UserCancelledException
 import com.amplifyframework.auth.cognito.helpers.HostedUIHelper
@@ -57,15 +55,9 @@ import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
 import com.amplifyframework.auth.cognito.result.GlobalSignOutError
 import com.amplifyframework.auth.cognito.result.HostedUIError
 import com.amplifyframework.auth.cognito.result.RevokeTokenError
-import com.amplifyframework.auth.exceptions.ConfigurationException
 import com.amplifyframework.auth.exceptions.InvalidStateException
-import com.amplifyframework.auth.exceptions.NotAuthorizedException
-import com.amplifyframework.auth.exceptions.ServiceException
-import com.amplifyframework.auth.exceptions.SessionExpiredException
-import com.amplifyframework.auth.exceptions.SignedOutException
 import com.amplifyframework.auth.exceptions.UnknownException
 import com.amplifyframework.auth.options.AuthConfirmSignInOptions
-import com.amplifyframework.auth.options.AuthFetchSessionOptions
 import com.amplifyframework.auth.options.AuthSignInOptions
 import com.amplifyframework.auth.options.AuthSignOutOptions
 import com.amplifyframework.auth.options.AuthWebUISignInOptions
@@ -112,9 +104,6 @@ import java.lang.ref.WeakReference
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
-import kotlin.coroutines.suspendCoroutine
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.takeWhile
 
@@ -990,142 +979,6 @@ internal class RealAWSCognitoAuthPlugin(
         }
     }
 
-    private suspend fun getSession(): AWSCognitoAuthSession = suspendCoroutine { continuation ->
-        fetchAuthSession(
-            { authSession ->
-                if (authSession is AWSCognitoAuthSession) {
-                    continuation.resume(authSession)
-                } else {
-                    continuation.resumeWithException(
-                        UnknownException(
-                            message = "fetchAuthSession did not return a type of AWSCognitoAuthSession"
-                        )
-                    )
-                }
-            },
-            { continuation.resumeWithException(it) }
-        )
-    }
-
-    fun fetchAuthSession(onSuccess: Consumer<AuthSession>, onError: Consumer<AuthException>) {
-        fetchAuthSession(AuthFetchSessionOptions.defaults(), onSuccess, onError)
-    }
-
-    fun fetchAuthSession(
-        options: AuthFetchSessionOptions,
-        onSuccess: Consumer<AuthSession>,
-        onError: Consumer<AuthException>
-    ) {
-        val forceRefresh = options.forceRefresh
-        authStateMachine.getCurrentState { authState ->
-            when (val authZState = authState.authZState) {
-                is AuthorizationState.Configured -> {
-                    authStateMachine.send(AuthorizationEvent(AuthorizationEvent.EventType.FetchUnAuthSession))
-                    _fetchAuthSession(onSuccess)
-                }
-                is AuthorizationState.SessionEstablished -> {
-                    val credential = authZState.amplifyCredential
-                    if (!credential.isValid() || forceRefresh) {
-                        if (credential is AmplifyCredential.IdentityPoolFederated) {
-                            authStateMachine.send(
-                                AuthorizationEvent(
-                                    AuthorizationEvent.EventType.StartFederationToIdentityPool(
-                                        credential.federatedToken,
-                                        credential.identityId,
-                                        credential
-                                    )
-                                )
-                            )
-                        } else {
-                            authStateMachine.send(
-                                AuthorizationEvent(AuthorizationEvent.EventType.RefreshSession(credential))
-                            )
-                        }
-                        _fetchAuthSession(onSuccess)
-                    } else {
-                        onSuccess.accept(credential.getCognitoSession())
-                    }
-                }
-                is AuthorizationState.Error -> {
-                    val error = authZState.exception
-                    if (error is SessionError) {
-                        val amplifyCredential = error.amplifyCredential
-                        if (amplifyCredential is AmplifyCredential.IdentityPoolFederated) {
-                            authStateMachine.send(
-                                AuthorizationEvent(
-                                    AuthorizationEvent.EventType.StartFederationToIdentityPool(
-                                        amplifyCredential.federatedToken,
-                                        amplifyCredential.identityId,
-                                        amplifyCredential
-                                    )
-                                )
-                            )
-                        } else {
-                            authStateMachine.send(
-                                AuthorizationEvent(AuthorizationEvent.EventType.RefreshSession(amplifyCredential))
-                            )
-                        }
-                        _fetchAuthSession(onSuccess)
-                    } else {
-                        onError.accept(InvalidStateException())
-                    }
-                }
-                else -> onError.accept(InvalidStateException())
-            }
-        }
-    }
-
-    private fun _fetchAuthSession(onSuccess: Consumer<AuthSession>) {
-        val token = StateChangeListenerToken()
-        authStateMachine.listen(
-            token,
-            { authState ->
-                when (val authZState = authState.authZState) {
-                    is AuthorizationState.SessionEstablished -> {
-                        authStateMachine.cancel(token)
-                        onSuccess.accept(authZState.amplifyCredential.getCognitoSession())
-                    }
-                    is AuthorizationState.Error -> {
-                        authStateMachine.cancel(token)
-                        when (val error = authZState.exception) {
-                            is SessionError -> {
-                                when (val innerException = error.exception) {
-                                    is SignedOutException -> {
-                                        onSuccess.accept(error.amplifyCredential.getCognitoSession(innerException))
-                                    }
-                                    is SessionExpiredException -> {
-                                        onSuccess.accept(error.amplifyCredential.getCognitoSession(innerException))
-                                        sendHubEvent(AuthChannelEventName.SESSION_EXPIRED.toString())
-                                    }
-                                    is ServiceException -> {
-                                        onSuccess.accept(error.amplifyCredential.getCognitoSession(innerException))
-                                    }
-                                    is NotAuthorizedException -> {
-                                        onSuccess.accept(error.amplifyCredential.getCognitoSession(innerException))
-                                    }
-                                    else -> {
-                                        val errorResult = UnknownException("Fetch auth session failed.", innerException)
-                                        onSuccess.accept(error.amplifyCredential.getCognitoSession(errorResult))
-                                    }
-                                }
-                            }
-                            is ConfigurationException -> {
-                                val errorResult = InvalidAccountTypeException(error)
-                                onSuccess.accept(AmplifyCredential.Empty.getCognitoSession(errorResult))
-                            }
-                            else -> {
-                                val errorResult = UnknownException("Fetch auth session failed.", error)
-                                onSuccess.accept(AmplifyCredential.Empty.getCognitoSession(errorResult))
-                            }
-                        }
-                    }
-                    else -> Unit
-                }
-            },
-            null
-        )
-    }
-
     fun signOut(onComplete: Consumer<AuthSignOutResult>) {
         signOut(AuthSignOutOptions.builder().build(), onComplete)
     }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/AuthUseCaseFactory.kt

@@ -17,19 +17,19 @@ package com.amplifyframework.auth.cognito.usecases
 
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.AuthStateMachine
-import com.amplifyframework.auth.cognito.RealAWSCognitoAuthPlugin
 import com.amplifyframework.auth.cognito.helpers.WebAuthnHelper
 import com.amplifyframework.auth.cognito.requireIdentityProviderClient
 import com.amplifyframework.auth.plugins.core.AuthHubEventEmitter
 
 internal class AuthUseCaseFactory(
-    private val plugin: RealAWSCognitoAuthPlugin,
     private val authEnvironment: AuthEnvironment,
     private val stateMachine: AuthStateMachine,
     private val hubEmitter: AuthHubEventEmitter = AuthHubEventEmitter()
 ) {
 
-    fun fetchAuthSession() = FetchAuthSessionUseCase(plugin)
+    fun fetchAuthSession() = FetchAuthSessionUseCase(
+        stateMachine = stateMachine
+    )
 
     fun associateWebAuthnCredential() = AssociateWebAuthnCredentialUseCase(
         client = authEnvironment.requireIdentityProviderClient(),

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/FetchAuthSessionUseCase.kt

@@ -15,23 +15,110 @@
 
 package com.amplifyframework.auth.cognito.usecases
 
+import com.amplifyframework.auth.AuthChannelEventName
 import com.amplifyframework.auth.cognito.AWSCognitoAuthSession
-import com.amplifyframework.auth.cognito.RealAWSCognitoAuthPlugin
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
-import kotlin.coroutines.suspendCoroutine
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.exceptions.service.InvalidAccountTypeException
+import com.amplifyframework.auth.cognito.getCognitoSession
+import com.amplifyframework.auth.cognito.isValid
+import com.amplifyframework.auth.exceptions.ConfigurationException
+import com.amplifyframework.auth.exceptions.InvalidStateException
+import com.amplifyframework.auth.exceptions.NotAuthorizedException
+import com.amplifyframework.auth.exceptions.ServiceException
+import com.amplifyframework.auth.exceptions.SessionExpiredException
+import com.amplifyframework.auth.exceptions.SignedOutException
+import com.amplifyframework.auth.exceptions.UnknownException
+import com.amplifyframework.auth.options.AuthFetchSessionOptions
+import com.amplifyframework.auth.plugins.core.AuthHubEventEmitter
+import com.amplifyframework.statemachine.codegen.data.AmplifyCredential
+import com.amplifyframework.statemachine.codegen.errors.SessionError
+import com.amplifyframework.statemachine.codegen.events.AuthorizationEvent
+import com.amplifyframework.statemachine.codegen.states.AuthorizationState
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.mapNotNull
 
 internal class FetchAuthSessionUseCase(
-    private val plugin: RealAWSCognitoAuthPlugin
+    private val stateMachine: AuthStateMachine,
+    private val emitter: AuthHubEventEmitter = AuthHubEventEmitter()
 ) {
-    suspend fun execute(): AWSCognitoAuthSession {
-        // TODO - we should migrate the fetch auth session business logic to this class
-        val session = suspendCoroutine { continuation ->
-            plugin.fetchAuthSession(
-                onSuccess = { continuation.resume(it) },
-                onError = { continuation.resumeWithException(it) }
+    suspend fun execute(options: AuthFetchSessionOptions = AuthFetchSessionOptions.defaults()): AWSCognitoAuthSession =
+        when (val authZState = stateMachine.getCurrentState().authZState) {
+            is AuthorizationState.Configured -> {
+                stateMachine.send(AuthorizationEvent(AuthorizationEvent.EventType.FetchUnAuthSession))
+                listenForSessionEstablished()
+            }
+            is AuthorizationState.SessionEstablished -> {
+                val credential = authZState.amplifyCredential
+                if (!credential.isValid() || options.forceRefresh) {
+                    sendRefreshSessionEvent(credential)
+                    listenForSessionEstablished()
+                } else {
+                    credential.getCognitoSession()
+                }
+            }
+            is AuthorizationState.Error -> {
+                val error = authZState.exception
+                if (error is SessionError) {
+                    sendRefreshSessionEvent(error.amplifyCredential)
+                    listenForSessionEstablished()
+                } else {
+                    throw InvalidStateException()
+                }
+            }
+            else -> throw InvalidStateException()
+        }
+
+    private fun sendRefreshSessionEvent(credential: AmplifyCredential) {
+        if (credential is AmplifyCredential.IdentityPoolFederated) {
+            stateMachine.send(
+                AuthorizationEvent(
+                    AuthorizationEvent.EventType.StartFederationToIdentityPool(
+                        credential.federatedToken,
+                        credential.identityId,
+                        credential
+                    )
+                )
             )
+        } else {
+            stateMachine.send(AuthorizationEvent(AuthorizationEvent.EventType.RefreshSession(credential)))
         }
-        return session as AWSCognitoAuthSession
+    }
+
+    private suspend fun listenForSessionEstablished(): AWSCognitoAuthSession {
+        val session = stateMachine.stateTransitions.mapNotNull { authState ->
+            when (val authZState = authState.authZState) {
+                is AuthorizationState.SessionEstablished -> authZState.amplifyCredential.getCognitoSession()
+                is AuthorizationState.Error -> {
+                    when (val error = authZState.exception) {
+                        is SessionError -> {
+                            when (val innerException = error.exception) {
+                                is SignedOutException -> error.amplifyCredential.getCognitoSession(innerException)
+                                is ServiceException -> error.amplifyCredential.getCognitoSession(innerException)
+                                is NotAuthorizedException -> error.amplifyCredential.getCognitoSession(innerException)
+                                is SessionExpiredException -> {
+                                    emitter.sendHubEvent(AuthChannelEventName.SESSION_EXPIRED.toString())
+                                    error.amplifyCredential.getCognitoSession(innerException)
+                                }
+                                else -> {
+                                    val errorResult = UnknownException("Fetch auth session failed.", innerException)
+                                    error.amplifyCredential.getCognitoSession(errorResult)
+                                }
+                            }
+                        }
+                        is ConfigurationException -> {
+                            val errorResult = InvalidAccountTypeException(error)
+                            AmplifyCredential.Empty.getCognitoSession(errorResult)
+                        }
+                        else -> {
+                            val errorResult = UnknownException("Fetch auth session failed.", error)
+                            AmplifyCredential.Empty.getCognitoSession(errorResult)
+                        }
+                    }
+                }
+                else -> null // no-op
+            }
+        }.first()
+
+        return session
     }
 }