fix(Auth): Add correct validation for initial state when executing confirm sign in (#2587)

* fix(Auth): Add correct validation for initial state when executing confirm sign in

* Update AWSAuthSignInPluginTests.swift

* Update AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AWSAuthSignInPluginTests.swift

Co-authored-by: Ian Saultz <[email protected]>

* Update AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AWSAuthSignInPluginTests.swift

Co-authored-by: Ian Saultz <[email protected]>

Co-authored-by: Ian Saultz <[email protected]>

Author: harsh62

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthConfirmSignInTask.swift

@@ -28,14 +28,19 @@ class AWSAuthConfirmSignInTask: AuthConfirmSignInTask {
         if let validationError = request.hasError() {
             throw validationError
         }
-
-        let pluginOptions = (request.options.pluginOptions as? AWSAuthConfirmSignInOptions)
-
-        await taskHelper.didStateMachineConfigured()
-
         let invalidStateError = AuthError.invalidState(
             "User is not attempting signIn operation",
             AuthPluginErrorConstants.invalidStateError, nil)
+
+        await taskHelper.didStateMachineConfigured()
+
+        if case .configured(let authNState, _) = await authStateMachine.currentState,
+           case .signingIn(let signInState) = authNState,
+           case .resolvingChallenge(let challengeState, _, _) = signInState,
+           case .waitingForAnswer = challengeState {
+            await sendConfirmSignInEvent()
+        }
+
         let stateSequences = await authStateMachine.listen()
         for await state in stateSequences {
                guard case .configured(let authNState, let authZState) = state else {
@@ -68,19 +73,11 @@ class AWSAuthConfirmSignInTask: AuthConfirmSignInTask {
                    } else if case .resolvingChallenge(let challengeState, _, _) = signInState {
                        switch challengeState {
                        case .waitingForAnswer:
-                           // Convert the attributes to [String: String]
-                           let attributePrefix = AuthPluginConstants.cognitoIdentityUserUserAttributePrefix
-                           let attributes = pluginOptions?.userAttributes?.reduce(
-                               into: [String: String]()) {
-                                   $0[attributePrefix + $1.key.rawValue] = $1.value
-                               } ?? [:]
-                           let confirmSignInData = ConfirmSignInEventData(
-                               answer: self.request.challengeResponse,
-                               attributes: attributes,
-                               metadata: pluginOptions?.metadata)
-                           let event = SignInChallengeEvent(
-                               eventType: .verifyChallengeAnswer(confirmSignInData))
-                           await authStateMachine.send(event)
+                           guard let result = try UserPoolSignInHelper.checkNextStep(signInState) else {
+                               continue
+                           }
+                           return result
+
                        default:
                            continue
                        }
@@ -96,4 +93,22 @@ class AWSAuthConfirmSignInTask: AuthConfirmSignInTask {
         throw invalidStateError
     }
 
+    func sendConfirmSignInEvent() async {
+        let pluginOptions = (request.options.pluginOptions as? AWSAuthConfirmSignInOptions)
+
+        // Convert the attributes to [String: String]
+        let attributePrefix = AuthPluginConstants.cognitoIdentityUserUserAttributePrefix
+        let attributes = pluginOptions?.userAttributes?.reduce(
+            into: [String: String]()) {
+                $0[attributePrefix + $1.key.rawValue] = $1.value
+            } ?? [:]
+        let confirmSignInData = ConfirmSignInEventData(
+            answer: self.request.challengeResponse,
+            attributes: attributes,
+            metadata: pluginOptions?.metadata)
+        let event = SignInChallengeEvent(
+            eventType: .verifyChallengeAnswer(confirmSignInData))
+        await authStateMachine.send(event)
+    }
+
 }

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AWSAuthSignInPluginTests.swift

@@ -585,6 +585,54 @@ class AWSAuthSignInPluginTests: BasePluginTest {
         }
     }
 
+    /// Test a signIn with customAuthWIthoutSRP
+    ///
+    /// - Given: An auth plugin with mocked service. Returning a new challenge after confirm sign in is called
+    ///
+    /// - When:
+    ///    - I invoke signIn and then confirm sign in
+    /// - Then:
+    ///    - The next step smsMfA should be triggered
+    ///
+    func testSignInWithCustomAuthIncorrectCode() async {
+
+        self.mockIdentityProvider = MockIdentityProvider(mockInitiateAuthResponse: { _ in
+            InitiateAuthOutputResponse(
+                authenticationResult: .none,
+                challengeName: .customChallenge,
+                challengeParameters: InitiateAuthOutputResponse.validChalengeParams,
+                session: "someSession")
+        }, mockRespondToAuthChallengeResponse: { _ in
+            RespondToAuthChallengeOutputResponse(
+                authenticationResult: .none,
+                challengeName: .smsMfa,
+                challengeParameters: ["paramKey": "value"],
+                session: "session")
+        })
+
+        let pluginOptions = AWSAuthSignInOptions(validationData: ["somekey": "somevalue"],
+                                                 metadata: ["somekey": "somevalue"],
+                                                 authFlowType: .customWithoutSRP)
+        let options = AuthSignInRequest.Options(pluginOptions: pluginOptions)
+        do {
+            let result = try await plugin.signIn(
+                username: "username",
+                password: "password",
+                options: options)
+            guard case .confirmSignInWithCustomChallenge = result.nextStep,
+                  case let confirmSignInResult = try await plugin.confirmSignIn(
+                    challengeResponse: "245234"
+                  ),
+                  case .confirmSignInWithSMSMFACode = confirmSignInResult.nextStep
+            else {
+                return XCTFail("Incorrect challenge type")
+            }
+        } catch {
+            XCTFail("Should not fail with \(error)")
+        }
+    }
+
+
     // MARK: - Service error for initiateAuth
 
     /// Test a signIn with `InternalErrorException` from service