fix(auth): Fix confirm signin when incorrect MFA code is entered

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

@@ -588,11 +588,17 @@ internal class RealAWSCognitoAuthPlugin(
         authStateMachine.getCurrentState { authState ->
             val authNState = authState.authNState
             val signInState = (authNState as? AuthenticationState.SigningIn)?.signInState
-            when ((signInState as? SignInState.ResolvingChallenge)?.challengeState) {
-                is SignInChallengeState.WaitingForAnswer -> {
-                    _confirmSignIn(challengeResponse, options, onSuccess, onError)
+            if (signInState is SignInState.ResolvingChallenge) {
+                when (signInState.challengeState) {
+                    is SignInChallengeState.WaitingForAnswer -> {
+                        _confirmSignIn(challengeResponse, options, onSuccess, onError)
+                    }
+                    else -> {
+                        onError.accept(InvalidStateException())
+                    }
                 }
-                else -> onError.accept(InvalidStateException())
+            } else {
+                onError.accept(InvalidStateException())
             }
         }
     }
@@ -610,7 +616,6 @@ internal class RealAWSCognitoAuthPlugin(
                 val authNState = authState.authNState
                 val authZState = authState.authZState
                 val signInState = (authNState as? AuthenticationState.SigningIn)?.signInState
-                val challengeState = (signInState as? SignInState.ResolvingChallenge)?.challengeState
                 when {
                     authNState is AuthenticationState.SignedIn &&
                         authZState is AuthorizationState.SessionEstablished -> {
@@ -625,21 +630,34 @@ internal class RealAWSCognitoAuthPlugin(
                     signInState is SignInState.Error -> {
                         authStateMachine.cancel(token)
                         onError.accept(
-                            CognitoAuthExceptionConverter.lookup(signInState.exception, "Confirm Sign in failed.")
+                            CognitoAuthExceptionConverter.lookup(
+                                signInState.exception, "Confirm Sign in failed."
+                            )
+                        )
+                    }
+                    signInState is SignInState.ResolvingChallenge &&
+                        signInState.challengeState is SignInChallengeState.Error -> {
+                        authStateMachine.cancel(token)
+                        onError.accept(
+                            CognitoAuthExceptionConverter.lookup(
+                                (
+                                    signInState.challengeState as SignInChallengeState.Error
+                                    ).exception,
+                                "Confirm Sign in failed."
+                            )
                         )
                     }
                 }
-            },
-            {
-                val awsCognitoConfirmSignInOptions = options as? AWSCognitoAuthConfirmSignInOptions
-                val event = SignInChallengeEvent(
-                    SignInChallengeEvent.EventType.VerifyChallengeAnswer(
-                        challengeResponse,
-                        awsCognitoConfirmSignInOptions?.metadata ?: mapOf()
-                    )
+            }, {
+            val awsCognitoConfirmSignInOptions = options as? AWSCognitoAuthConfirmSignInOptions
+            val event = SignInChallengeEvent(
+                SignInChallengeEvent.EventType.VerifyChallengeAnswer(
+                    challengeResponse,
+                    awsCognitoConfirmSignInOptions?.metadata ?: mapOf()
                 )
-                authStateMachine.send(event)
-            }
+            )
+            authStateMachine.send(event)
+        }
         )
     }
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -26,7 +26,6 @@ import com.amplifyframework.statemachine.codegen.actions.SignInChallengeActions
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
 import com.amplifyframework.statemachine.codegen.events.CustomSignInEvent
 import com.amplifyframework.statemachine.codegen.events.SignInChallengeEvent
-import com.amplifyframework.statemachine.codegen.events.SignInEvent
 
 internal object SignInChallengeCognitoActions : SignInChallengeActions {
     private const val KEY_SECRET_HASH = "SECRET_HASH"
@@ -81,20 +80,25 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
                 )
             )
         } catch (e: Exception) {
-            SignInEvent(SignInEvent.EventType.ThrowError(e))
+            SignInChallengeEvent(SignInChallengeEvent.EventType.ThrowError(e, challenge))
         }
         logger.verbose("$id Sending event ${evt.type}")
         dispatcher.send(evt)
     }
 
+    override fun resetToWaitingForAnswer(
+        event: SignInChallengeEvent.EventType.ThrowError,
+        challenge: AuthChallenge
+    ): Action = Action<AuthEnvironment>("ResetToWaitingForAnswer") { id, dispatcher ->
+        logger.verbose("$id Starting execution")
+        dispatcher.send(SignInChallengeEvent(SignInChallengeEvent.EventType.WaitForAnswer(challenge)))
+    }
+
     private fun getChallengeResponseKey(challengeName: String): String? {
-        val VALUE_ANSWER = "ANSWER"
-        val VALUE_SMS_MFA = "SMS_MFA_CODE"
-        val VALUE_NEW_PASSWORD = "NEW_PASSWORD"
         return when (ChallengeNameType.fromValue(challengeName)) {
-            is ChallengeNameType.SmsMfa -> VALUE_SMS_MFA
-            is ChallengeNameType.NewPasswordRequired -> VALUE_NEW_PASSWORD
-            is ChallengeNameType.CustomChallenge -> VALUE_ANSWER
+            is ChallengeNameType.SmsMfa -> "SMS_MFA_CODE"
+            is ChallengeNameType.NewPasswordRequired -> "NEW_PASSWORD"
+            is ChallengeNameType.CustomChallenge -> "ANSWER"
             else -> null
         }
     }

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

@@ -24,4 +24,8 @@ internal interface SignInChallengeActions {
         event: SignInChallengeEvent.EventType.VerifyChallengeAnswer,
         challenge: AuthChallenge
     ): Action
+    fun resetToWaitingForAnswer(
+        event: SignInChallengeEvent.EventType.ThrowError,
+        challenge: AuthChallenge
+    ): Action
 }

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

@@ -25,6 +25,7 @@ internal class SignInChallengeEvent(val eventType: EventType, override val time:
         data class VerifyChallengeAnswer(val answer: String, val metadata: Map<String, String>) : EventType()
         data class FinalizeSignIn(val accessToken: String) : EventType()
         data class Verified(val id: String = "") : EventType()
+        data class ThrowError(val exception: Exception, val challenge: AuthChallenge) : EventType()
     }
 
     override val type: String = eventType.javaClass.simpleName

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

@@ -28,6 +28,7 @@ internal sealed class SignInChallengeState : State {
     data class WaitingForAnswer(val challenge: AuthChallenge) : SignInChallengeState()
     data class Verifying(val id: String = "") : SignInChallengeState()
     data class Verified(val id: String = "") : SignInChallengeState()
+    data class Error(val exception: Exception, val challenge: AuthChallenge) : SignInChallengeState()
 
     class Resolver(private val challengeActions: SignInChallengeActions) : StateMachineResolver<SignInChallengeState> {
         override val defaultState: SignInChallengeState = NotStarted()
@@ -58,8 +59,25 @@ internal sealed class SignInChallengeState : State {
                 }
                 is Verifying -> when (challengeEvent) {
                     is SignInChallengeEvent.EventType.Verified -> StateResolution(Verified())
+                    is SignInChallengeEvent.EventType.ThrowError -> {
+                        val action = challengeActions.resetToWaitingForAnswer(challengeEvent, challengeEvent.challenge)
+                        StateResolution(Error(challengeEvent.exception, challengeEvent.challenge), listOf(action))
+                    }
+
                     else -> defaultResolution
                 }
+                is Error -> {
+                    when (challengeEvent) {
+                        is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {
+                            val action = challengeActions.verifyChallengeAuthAction(challengeEvent, oldState.challenge)
+                            StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))
+                        }
+                        is SignInChallengeEvent.EventType.WaitForAnswer -> {
+                            StateResolution(WaitingForAnswer(challengeEvent.challenge))
+                        }
+                        else -> defaultResolution
+                    }
+                }
                 else -> defaultResolution
             }
         }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -149,6 +149,10 @@ internal sealed class SignInState : State {
                         val action = signInActions.confirmDevice(signInEvent)
                         StateResolution(ConfirmingDevice(), listOf(action))
                     }
+                    is SignInEvent.EventType.ReceivedChallenge -> {
+                        val action = signInActions.initResolveChallenge(signInEvent)
+                        StateResolution(ResolvingChallenge(oldState.challengeState), listOf(action))
+                    }
                     is SignInEvent.EventType.ThrowError -> StateResolution(Error(signInEvent.exception))
                     else -> defaultResolution
                 }

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/featuretest/generators/testcasegenerators/ConfirmSignInTestCaseGenerator.kt

@@ -15,6 +15,8 @@
 
 package com.amplifyframework.auth.cognito.featuretest.generators.testcasegenerators
 
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeMismatchException
+import com.amplifyframework.auth.cognito.CognitoAuthExceptionConverter
 import com.amplifyframework.auth.cognito.featuretest.API
 import com.amplifyframework.auth.cognito.featuretest.AuthAPI
 import com.amplifyframework.auth.cognito.featuretest.CognitoType
@@ -102,5 +104,37 @@ object ConfirmSignInTestCaseGenerator : SerializableProvider {
         )
     )
 
-    override val serializables: List<Any> = listOf(baseCase)
+    private val errorCase: FeatureTestCase
+        get() {
+            val exception = CodeMismatchException.invoke {
+                message = "Confirmation code entered is not correct."
+            }
+            return baseCase.copy(
+                description = "Test that invalid code on confirm SignIn with SMS challenge errors out",
+                preConditions = PreConditions(
+                    "authconfiguration.json",
+                    "SigningIn_SigningIn.json",
+                    mockedResponses = listOf(
+                        MockResponse(
+                            CognitoType.CognitoIdentityProvider,
+                            "respondToAuthChallenge",
+                            ResponseType.Failure,
+                            exception.toJsonElement()
+                        )
+                    )
+                ),
+                validations = listOf(
+                    ExpectationShapes.Amplify(
+                        AuthAPI.confirmSignIn,
+                        ResponseType.Failure,
+                        CognitoAuthExceptionConverter.lookup(
+                            exception,
+                            "Confirm Sign in failed."
+                        ).toJsonElement()
+                    )
+                )
+            )
+        }
+
+    override val serializables: List<Any> = listOf(baseCase, errorCase)
 }

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/featuretest/serializers/AuthStatesSerializer.kt

@@ -177,6 +177,7 @@ internal data class AuthStatesProxy(
                             type = "SignInChallengeState.WaitingForAnswer",
                             authChallenge = authState.challenge
                         )
+                        is SignInChallengeState.Error -> TODO()
                     }
                 }
                 else -> {

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/featuretest/serializers/CognitoExceptionSerializers.kt

@@ -18,8 +18,10 @@
 package com.amplifyframework.auth.cognito.featuretest.serializers
 
 import aws.sdk.kotlin.services.cognitoidentity.model.CognitoIdentityException
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeMismatchException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.CognitoIdentityProviderException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.NotAuthorizedException
+import com.amplifyframework.auth.exceptions.UnknownException
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.descriptors.SerialDescriptor
@@ -32,14 +34,19 @@ private data class CognitoExceptionSurrogate(
     val errorMessage: String?
 ) {
     fun <T> toRealException(): T {
-        return when (errorType) {
+        val exception = when (errorType) {
             NotAuthorizedException::class.java.simpleName -> NotAuthorizedException.invoke {
                 message = errorMessage
             } as T
+            UnknownException::class.java.simpleName -> UnknownException(message = errorMessage ?: "") as T
+            CodeMismatchException::class.java.simpleName -> CodeMismatchException.invoke {
+                message = errorMessage
+            } as T
             else -> {
                 error("Exception for $errorType not defined")
             }
         }
+        return exception
     }
 
     companion object {

aws-auth-cognito/src/test/java/featureTest/utilities/APIExecutor.kt

@@ -22,7 +22,6 @@ import com.amplifyframework.core.Consumer
 import com.google.gson.Gson
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
-import kotlin.Exception
 import kotlin.reflect.KClass
 import kotlin.reflect.KFunction
 import kotlin.reflect.KParameter

aws-auth-cognito/src/test/java/featureTest/utilities/CognitoMockFactory.kt

@@ -20,7 +20,6 @@ import aws.sdk.kotlin.services.cognitoidentity.model.Credentials
 import aws.sdk.kotlin.services.cognitoidentity.model.GetCredentialsForIdentityResponse
 import aws.sdk.kotlin.services.cognitoidentity.model.GetIdResponse
 import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
-import aws.sdk.kotlin.services.cognitoidentityprovider.forgetDevice
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AttributeType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AuthenticationResultType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
@@ -211,13 +210,14 @@ class CognitoMockFactory(
         responseObject: JsonObject
     ) {
         if (mockResponse.responseType == ResponseType.Failure) {
-            throw Json.decodeFromString(
+            val response = Json.decodeFromString(
                 when (mockResponse.type) {
                     CognitoType.CognitoIdentity -> CognitoIdentityExceptionSerializer
                     CognitoType.CognitoIdentityProvider -> CognitoIdentityProviderExceptionSerializer
                 },
                 responseObject.toString()
             )
+            throw response
         }
     }
 

aws-auth-cognito/src/test/resources/feature-test/testsuites/confirmSignIn/Test_that_invalid_code_on_confirm_SignIn_with_SMS_challenge_errors_out.json

@@ -0,0 +1,42 @@
+{
+    "description": "Test that invalid code on confirm SignIn with SMS challenge errors out",
+    "preConditions": {
+        "amplify-configuration": "authconfiguration.json",
+        "state": "SigningIn_SigningIn.json",
+        "mockedResponses": [
+            {
+                "type": "cognitoIdentityProvider",
+                "apiName": "respondToAuthChallenge",
+                "responseType": "failure",
+                "response": {
+                    "errorType": "CodeMismatchException",
+                    "errorMessage": "Confirmation code entered is not correct."
+                }
+            }
+        ]
+    },
+    "api": {
+        "name": "confirmSignIn",
+        "params": {
+            "challengeResponse": "000000"
+        },
+        "options": {
+        }
+    },
+    "validations": [
+        {
+            "type": "amplify",
+            "apiName": "confirmSignIn",
+            "responseType": "failure",
+            "response": {
+                "errorType": "CodeMismatchException",
+                "errorMessage": "Confirmation code entered is not correct.",
+                "recoverySuggestion": "Enter correct confirmation code.",
+                "cause": {
+                    "errorType": "CodeMismatchException",
+                    "errorMessage": "Confirmation code entered is not correct."
+                }
+            }
+        }
+    ]
+}

aws-auth-cognito/src/test/resources/feature-test/testsuites/fetchDevices/List_of_devices_returned_when_fetch_devices_API_succeeds.json

@@ -19,21 +19,21 @@
                             ],
                             "deviceCreateDate": {
                                 "value": {
-                                    "seconds": 1.671733506E9,
-                                    "nanos": 9.95178E8
+                                    "seconds": 1.676485556E9,
+                                    "nanos": 7.65001E8
                                 }
                             },
                             "deviceKey": "deviceKey",
                             "deviceLastAuthenticatedDate": {
                                 "value": {
-                                    "seconds": 1.671733506E9,
-                                    "nanos": 9.95186E8
+                                    "seconds": 1.676485556E9,
+                                    "nanos": 7.65007E8
                                 }
                             },
                             "deviceLastModifiedDate": {
                                 "value": {
-                                    "seconds": 1.671733506E9,
-                                    "nanos": 9.95188E8
+                                    "seconds": 1.676485556E9,
+                                    "nanos": 7.65008E8
                                 }
                             }
                         }