Fix for retry sign in when resourceNotFoundException is raised

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

@@ -18,6 +18,7 @@ package com.amplifyframework.auth.cognito.actions
 import aws.sdk.kotlin.services.cognitoidentityprovider.initiateAuth
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AuthFlowType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.ResourceNotFoundException
 import aws.sdk.kotlin.services.cognitoidentityprovider.respondToAuthChallenge
 import com.amplifyframework.AmplifyException
 import com.amplifyframework.auth.cognito.AuthEnvironment
@@ -27,6 +28,7 @@ import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.exceptions.ServiceException
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.actions.SRPActions
+import com.amplifyframework.statemachine.codegen.data.CredentialType
 import com.amplifyframework.statemachine.codegen.events.AuthenticationEvent
 import com.amplifyframework.statemachine.codegen.events.SRPEvent
 import com.amplifyframework.statemachine.codegen.events.SignInEvent
@@ -168,17 +170,20 @@ internal object SRPCognitoActions : SRPActions {
             dispatcher.send(evt)
         }
 
-    override fun verifyPasswordSRPAction(event: SRPEvent.EventType.RespondPasswordVerifier) =
+    override fun verifyPasswordSRPAction(
+        challengeParameters: Map<String, String>,
+        metadata: Map<String, String>,
+        session: String?
+    ) =
         Action<AuthEnvironment>("VerifyPasswordSRP") { id, dispatcher ->
             logger.verbose("$id Starting execution")
             val evt = try {
-                val params = event.challengeParameters
-                val salt = params.getValue(KEY_SALT)
-                val secretBlock = params.getValue(KEY_SECRET_BLOCK)
-                val srpB = params.getValue(KEY_SRP_B)
-                val username = params.getValue(KEY_USERNAME)
-                val userId = params.getValue(KEY_USER_ID_FOR_SRP)
-                val deviceKey = params.getOrDefault(KEY_DEVICE_KEY, "")
+                val salt = challengeParameters.getValue(KEY_SALT)
+                val secretBlock = challengeParameters.getValue(KEY_SECRET_BLOCK)
+                val srpB = challengeParameters.getValue(KEY_SRP_B)
+                val username = challengeParameters.getValue(KEY_USERNAME)
+                val userId = challengeParameters.getValue(KEY_USER_ID_FOR_SRP)
+                val deviceKey = challengeParameters.getOrDefault(KEY_DEVICE_KEY, "")
 
                 srpHelper.setUserPoolParams(userId, configuration.userPool?.poolId!!)
 
@@ -204,8 +209,8 @@ internal object SRPCognitoActions : SRPActions {
                     challengeName = ChallengeNameType.PasswordVerifier
                     clientId = configuration.userPool.appClient
                     challengeResponses = challengeParams
-                    clientMetadata = event.metadata
-                    session = event.session
+                    clientMetadata = metadata
+                    this.session = session
                     pinpointEndpointId?.let { analyticsMetadata { analyticsEndpointId = it } }
                     encodedContextData?.let { userContextData { encodedData = it } }
                 }
@@ -224,13 +229,28 @@ internal object SRPCognitoActions : SRPActions {
                     )
                 }
             } catch (e: Exception) {
-                val errorEvent = SRPEvent(SRPEvent.EventType.ThrowPasswordVerifierError(e))
-                logger.verbose("$id Sending event ${errorEvent.type}")
-                dispatcher.send(errorEvent)
-                val errorEvent2 = SignInEvent(SignInEvent.EventType.ThrowError(e))
-                logger.verbose("$id Sending event ${errorEvent.type}")
-                dispatcher.send(errorEvent2)
-                AuthenticationEvent(AuthenticationEvent.EventType.CancelSignIn())
+                if (e is ResourceNotFoundException) {
+                    val challengeParams: MutableMap<String, String> = challengeParameters.toMutableMap()
+                    challengeParams.remove(KEY_DEVICE_KEY)
+                    credentialStoreClient.clearCredentials(
+                        CredentialType.Device(
+                            challengeParams.getValue(
+                                KEY_USERNAME
+                            )
+                        )
+                    )
+                    SRPEvent(SRPEvent.EventType.RetryRespondPasswordVerifier(challengeParams, metadata, session))
+                } else {
+                    val errorEvent = SRPEvent(SRPEvent.EventType.ThrowPasswordVerifierError(e))
+                    logger.verbose("$id Sending event ${errorEvent.type}")
+                    dispatcher.send(errorEvent)
+
+                    val errorEvent2 = SignInEvent(SignInEvent.EventType.ThrowError(e))
+                    logger.verbose("$id Sending event ${errorEvent.type}")
+                    dispatcher.send(errorEvent2)
+
+                    AuthenticationEvent(AuthenticationEvent.EventType.CancelSignIn())
+                }
             }
             logger.verbose("$id Sending event ${evt.type}")
             dispatcher.send(evt)

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

@@ -16,6 +16,7 @@
 package com.amplifyframework.auth.cognito.actions
 
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.ResourceNotFoundException
 import aws.sdk.kotlin.services.cognitoidentityprovider.respondToAuthChallenge
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.helpers.AuthHelper
@@ -24,14 +25,16 @@ import com.amplifyframework.auth.exceptions.UnknownException
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.actions.SignInChallengeActions
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
+import com.amplifyframework.statemachine.codegen.data.CredentialType
 import com.amplifyframework.statemachine.codegen.events.CustomSignInEvent
 import com.amplifyframework.statemachine.codegen.events.SignInChallengeEvent
 
 internal object SignInChallengeCognitoActions : SignInChallengeActions {
     private const val KEY_SECRET_HASH = "SECRET_HASH"
     private const val KEY_USERNAME = "USERNAME"
     override fun verifyChallengeAuthAction(
-        event: SignInChallengeEvent.EventType.VerifyChallengeAnswer,
+        answer: String,
+        metadata: Map<String, String>,
         challenge: AuthChallenge
     ): Action = Action<AuthEnvironment>("VerifySignInChallenge") { id, dispatcher ->
         logger.verbose("$id Starting execution")
@@ -44,7 +47,7 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
             }
 
             getChallengeResponseKey(challenge.challengeName)?.also { responseKey ->
-                challengeResponses[responseKey] = event.answer
+                challengeResponses[responseKey] = answer
             }
 
             val secretHash = AuthHelper.getSecretHash(
@@ -61,7 +64,7 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
                 challengeName = ChallengeNameType.fromValue(challenge.challengeName)
                 this.challengeResponses = challengeResponses
                 session = challenge.session
-                clientMetadata = event.metadata
+                clientMetadata = metadata
                 pinpointEndpointId?.let { analyticsMetadata { analyticsEndpointId = it } }
                 encodedContextData?.let { this.userContextData { encodedData = it } }
             }
@@ -79,7 +82,20 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
                 )
             )
         } catch (e: Exception) {
-            SignInChallengeEvent(SignInChallengeEvent.EventType.ThrowError(e, challenge, true))
+            if (e is ResourceNotFoundException) {
+                challenge.username?.let { username ->
+                    credentialStoreClient.clearCredentials(CredentialType.Device(username))
+                }
+                SignInChallengeEvent(
+                    SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer(
+                        answer,
+                        metadata,
+                        challenge
+                    )
+                )
+            } else {
+                SignInChallengeEvent(SignInChallengeEvent.EventType.ThrowError(e, challenge, true))
+            }
         }
         logger.verbose("$id Sending event ${evt.type}")
         dispatcher.send(evt)

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/actions/SRPActions.kt

@@ -21,5 +21,9 @@ import com.amplifyframework.statemachine.codegen.events.SRPEvent
 internal interface SRPActions {
     fun initiateSRPAuthAction(event: SRPEvent.EventType.InitiateSRP): Action
     fun initiateSRPWithCustomAuthAction(event: SRPEvent.EventType.InitiateSRPWithCustom): Action
-    fun verifyPasswordSRPAction(event: SRPEvent.EventType.RespondPasswordVerifier): Action
+    fun verifyPasswordSRPAction(
+        challengeParameters: Map<String, String>,
+        metadata: Map<String, String>,
+        session: String?
+    ): Action
 }

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/actions/SignInChallengeActions.kt

@@ -17,11 +17,11 @@ package com.amplifyframework.statemachine.codegen.actions
 
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
-import com.amplifyframework.statemachine.codegen.events.SignInChallengeEvent
 
 internal interface SignInChallengeActions {
     fun verifyChallengeAuthAction(
-        event: SignInChallengeEvent.EventType.VerifyChallengeAnswer,
+        answer: String,
+        metadata: Map<String, String>,
         challenge: AuthChallenge
     ): Action
 }

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/events/SRPEvent.kt

@@ -37,6 +37,12 @@ internal class SRPEvent(val eventType: EventType, override val time: Date? = nul
             val session: String?
         ) : EventType()
 
+        data class RetryRespondPasswordVerifier(
+            val challengeParameters: Map<String, String>,
+            val metadata: Map<String, String>,
+            val session: String?
+        ) : EventType()
+
         data class ThrowAuthError(val exception: Exception) : EventType()
         data class CancelSRPSignIn(val id: String = "") : EventType()
         data class ThrowPasswordVerifierError(val exception: Exception) : EventType()

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/events/SignInChallengeEvent.kt

@@ -23,6 +23,12 @@ internal class SignInChallengeEvent(val eventType: EventType, override val time:
     sealed class EventType {
         data class WaitForAnswer(val challenge: AuthChallenge, val hasNewResponse: Boolean = false) : EventType()
         data class VerifyChallengeAnswer(val answer: String, val metadata: Map<String, String>) : EventType()
+
+        data class RetryVerifyChallengeAnswer(
+            val answer: String,
+            val metadata: Map<String, String>,
+            val authChallenge: AuthChallenge
+        ) : EventType()
         data class FinalizeSignIn(val accessToken: String) : EventType()
         data class Verified(val id: String = "") : EventType()
         data class ThrowError(

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/states/SRPSignInState.kt

@@ -54,14 +54,22 @@ internal sealed class SRPSignInState : State {
                 }
                 is InitiatingSRPA -> when (srpEvent) {
                     is SRPEvent.EventType.RespondPasswordVerifier -> {
-                        val action = srpActions.verifyPasswordSRPAction(srpEvent)
+                        val action = srpActions.verifyPasswordSRPAction(
+                            srpEvent.challengeParameters, srpEvent.metadata, srpEvent.session
+                        )
                         StateResolution(RespondingPasswordVerifier(), listOf(action))
                     }
                     is SRPEvent.EventType.ThrowAuthError -> StateResolution(Error(srpEvent.exception))
                     is SRPEvent.EventType.CancelSRPSignIn -> StateResolution(Cancelling())
                     else -> defaultResolution
                 }
                 is RespondingPasswordVerifier -> when (srpEvent) {
+                    is SRPEvent.EventType.RetryRespondPasswordVerifier -> {
+                        val action = srpActions.verifyPasswordSRPAction(
+                            srpEvent.challengeParameters, srpEvent.metadata, srpEvent.session
+                        )
+                        StateResolution(RespondingPasswordVerifier(), listOf(action))
+                    }
                     is SRPEvent.EventType.ThrowPasswordVerifierError -> StateResolution(Error(srpEvent.exception))
                     is SRPEvent.EventType.CancelSRPSignIn -> StateResolution(Cancelling())
                     else -> defaultResolution

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/states/SignInChallengeState.kt

@@ -59,15 +59,28 @@ internal sealed class SignInChallengeState : State {
                 }
                 is WaitingForAnswer -> when (challengeEvent) {
                     is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {
-                        val action = challengeActions.verifyChallengeAuthAction(challengeEvent, oldState.challenge)
+                        val action = challengeActions.verifyChallengeAuthAction(
+                            challengeEvent.answer, challengeEvent.metadata, oldState.challenge
+                        )
                         StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))
                     }
                     else -> defaultResolution
                 }
                 is Verifying -> when (challengeEvent) {
                     is SignInChallengeEvent.EventType.Verified -> StateResolution(Verified())
                     is SignInChallengeEvent.EventType.ThrowError -> {
-                        StateResolution(Error(challengeEvent.exception, challengeEvent.challenge, true), listOf())
+                        StateResolution(
+                            Error(
+                                challengeEvent.exception, challengeEvent.challenge, true
+                            ),
+                            listOf()
+                        )
+                    }
+                    is SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer -> {
+                        val action = challengeActions.verifyChallengeAuthAction(
+                            challengeEvent.answer, challengeEvent.metadata, challengeEvent.authChallenge
+                        )
+                        StateResolution(Verifying(challengeEvent.authChallenge.challengeName), listOf(action))
                     }
                     is SignInChallengeEvent.EventType.WaitForAnswer -> {
                         StateResolution(WaitingForAnswer(challengeEvent.challenge, true), listOf())
@@ -78,7 +91,9 @@ internal sealed class SignInChallengeState : State {
                 is Error -> {
                     when (challengeEvent) {
                         is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {
-                            val action = challengeActions.verifyChallengeAuthAction(challengeEvent, oldState.challenge)
+                            val action = challengeActions.verifyChallengeAuthAction(
+                                challengeEvent.answer, challengeEvent.metadata, oldState.challenge
+                            )
                             StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))
                         }
                         is SignInChallengeEvent.EventType.WaitForAnswer -> {

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/StateTransitionTestBase.kt

@@ -265,7 +265,15 @@ open class StateTransitionTestBase {
         Mockito.`when`(mockSignInActions.startSRPAuthAction(MockitoHelper.anyObject()))
             .thenReturn(
                 Action { dispatcher, _ ->
-                    dispatcher.send(SRPEvent(SRPEvent.EventType.InitiateSRP("username", "password", mapOf())))
+                    dispatcher.send(
+                        SRPEvent(
+                            SRPEvent.EventType.InitiateSRP(
+                                "username",
+                                "password",
+                                mapOf()
+                            )
+                        )
+                    )
                 }
             )
 
@@ -319,6 +327,7 @@ open class StateTransitionTestBase {
 
         Mockito.`when`(
             mockSignInChallengeActions.verifyChallengeAuthAction(
+                MockitoHelper.anyObject(),
                 MockitoHelper.anyObject(),
                 MockitoHelper.anyObject()
             )
@@ -355,7 +364,13 @@ open class StateTransitionTestBase {
                 }
             )
 
-        Mockito.`when`(mockSRPActions.verifyPasswordSRPAction(MockitoHelper.anyObject()))
+        Mockito.`when`(
+            mockSRPActions.verifyPasswordSRPAction(
+                MockitoHelper.anyObject(),
+                MockitoHelper.anyObject(),
+                MockitoHelper.anyObject()
+            )
+        )
             .thenReturn(
                 Action { dispatcher, _ ->
                     dispatcher.send(