fix(auth): Enable retrying when confirm signIn fails (#2617)

royjit · web-flow · commit b01d22da37c8 · 2022-12-19T12:17:23.000-08:00
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift
@@ -89,7 +89,9 @@ extension AWSCognitoAuthPlugin: AuthCategoryBehavior {
         let options = options ?? AuthConfirmSignInRequest.Options()
         let request = AuthConfirmSignInRequest(challengeResponse: challengeResponse,
                                                options: options)
-        let task = AWSAuthConfirmSignInTask(request, stateMachine: authStateMachine)
+        let task = AWSAuthConfirmSignInTask(request,
+                                            stateMachine: authStateMachine,
+                                            configuration: authConfiguration)
         return try await taskQueue.sync {
             return try await task.value
         } as! AuthSignInResult
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/SignInChallengeState+Debug.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/SignInChallengeState+Debug.swift
@@ -14,9 +14,9 @@ extension SignInChallengeState {
         switch self {
 
         case .waitingForAnswer(let respondAuthChallenge, _),
-                .verifying(let respondAuthChallenge, _):
+                .verifying(let respondAuthChallenge, _, _):
             additionalMetadataDictionary = respondAuthChallenge.debugDictionary
-        case .error(let respondAuthChallenge, let error):
+        case .error(let respondAuthChallenge, _, let error):
             additionalMetadataDictionary = respondAuthChallenge.debugDictionary
             additionalMetadataDictionary["error"] = error
         default: additionalMetadataDictionary = [:]
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/SignInChallengeState.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/SignInChallengeState.swift
@@ -13,11 +13,11 @@ enum SignInChallengeState: State {
 
     case waitingForAnswer(RespondToAuthChallenge, SignInMethod)
 
-    case verifying(RespondToAuthChallenge, String)
+    case verifying(RespondToAuthChallenge, SignInMethod, String)
 
     case verified
 
-    case error(RespondToAuthChallenge, SignInError)
+    case error(RespondToAuthChallenge, SignInMethod, SignInError)
 }
 
 extension SignInChallengeState {
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/SignIn/SignInChallengeState+Resolver.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/SignIn/SignInChallengeState+Resolver.swift
@@ -33,23 +33,40 @@ extension SignInChallengeState {
                         challenge: challenge,
                         confirmSignEventData: answerEventData,
                         signInMethod: signInMethod)
-                    return .init(newState: .verifying(challenge, answerEventData.answer), actions: [action])
+                    return .init(
+                        newState: .verifying(challenge, signInMethod, answerEventData.answer),
+                        actions: [action]
+                    )
                 }
                 return .from(oldState)
 
-            case .verifying(let challenge, _):
+            case .verifying(let challenge, let signInMethod, _):
 
                 if case .finalizeSignIn(let signedInData) = event.isSignInEvent {
                     return .init(newState: .verified,
                                  actions: [SignInComplete(signedInData: signedInData)])
                 }
 
                 if case .throwAuthError(let error) = event.isSignInEvent {
-                    return .init(newState: .error(challenge, error))
+                    return .init(newState: .error(challenge, signInMethod, error))
                 }
                 return .from(oldState)
 
-            case .verified, .error:
+            case .error(let challenge, let signInMethod, _):
+                // If a verifyChallengeAnswer is received on error state we allow
+                // to retry the challenge.
+                if case .verifyChallengeAnswer(let answerEventData) = event.isChallengeEvent {
+                    let action = VerifySignInChallenge(
+                        challenge: challenge,
+                        confirmSignEventData: answerEventData,
+                        signInMethod: signInMethod)
+                    return .init(
+                        newState: .verifying(challenge, signInMethod, answerEventData.answer),
+                        actions: [action]
+                    )
+                }
+                return .from(oldState)
+            case .verified:
 
                 return .from(oldState)
             }
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthConfirmSignInTask.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthConfirmSignInTask.swift
@@ -13,32 +13,54 @@ class AWSAuthConfirmSignInTask: AuthConfirmSignInTask {
     private let request: AuthConfirmSignInRequest
     private let authStateMachine: AuthStateMachine
     private let taskHelper: AWSAuthTaskHelper
+    private let authConfiguration: AuthConfiguration
 
     var eventName: HubPayloadEventName {
         HubPayload.EventName.Auth.confirmSignInAPI
     }
 
-    init(_ request: AuthConfirmSignInRequest, stateMachine: AuthStateMachine) {
+    init(_ request: AuthConfirmSignInRequest,
+         stateMachine: AuthStateMachine,
+         configuration: AuthConfiguration) {
         self.request = request
         self.authStateMachine = stateMachine
         self.taskHelper = AWSAuthTaskHelper(authStateMachine: authStateMachine)
+        self.authConfiguration = configuration
     }
 
     func execute() async throws -> AuthSignInResult {
+
+        await taskHelper.didStateMachineConfigured()
+
+        //Check if we have a user pool configuration
+        guard authConfiguration.getUserPoolConfiguration() != nil else {
+            let message = AuthPluginErrorConstants.configurationError
+            let authError = AuthError.configuration(
+                "Could not find user pool configuration",
+                message)
+            throw authError
+        }
+
         if let validationError = request.hasError() {
             throw validationError
         }
         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 {
+
+        guard case .configured(let authNState, _) = await authStateMachine.currentState,
+              case .signingIn(let signInState) = authNState,
+              case .resolvingChallenge(let challengeState, _, _) = signInState else {
+            throw invalidStateError
+        }
+
+        switch challengeState {
+        case .waitingForAnswer, .error:
             await sendConfirmSignInEvent()
+        default:
+            throw invalidStateError
         }
 
         let stateSequences = await authStateMachine.listen()
@@ -56,7 +78,7 @@ class AWSAuthConfirmSignInTask: AuthConfirmSignInTask {
 
                case .signingIn(let signInState):
                    if case .resolvingChallenge(let challengeState, _, _) = signInState,
-                      case .error(_, let signInError) = challengeState {
+                      case .error(_, _, let signInError) = challengeState {
                        let authError = signInError.authError
                        if case .service(_, _, let serviceError) = authError,
                           let cognitoError = serviceError as? AWSCognitoAuthError,
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthSignInTask.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthSignInTask.swift
@@ -29,13 +29,17 @@ class AWSAuthSignInTask: AuthSignInTask {
     }
 
     func execute() async throws -> AuthSignInResult {
+        await taskHelper.didStateMachineConfigured()
+        //Check if we have a user pool configuration
         guard let userPoolConfiguration = authConfiguration.getUserPoolConfiguration() else {
             let message = AuthPluginErrorConstants.configurationError
-            let authError = AuthenticationError.configuration(message: message)
+            let authError = AuthError.configuration(
+                "Could not find user pool configuration",
+                message)
             throw authError
         }
 
-        await taskHelper.didStateMachineConfigured()
+
         try await validateCurrentState()
 
         let authflowType = authFlowType(userPoolConfiguration: userPoolConfiguration)
diff --git a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/ClientBehaviorTests/AWSAuthConfirmSignInTaskTests.swift b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/ClientBehaviorTests/AWSAuthConfirmSignInTaskTests.swift
@@ -59,7 +59,7 @@ class AuthenticationProviderConfirmSigninTests: BasePluginTest {
     /// - Then:
     ///    - I should get an .validation error
     ///
-    func testSuccessfullyConfirmSignIn() async {
+    func testConfirmSignInWithEmptyResponse() async {
         
         self.mockIdentityProvider = MockIdentityProvider(
             mockRespondToAuthChallengeResponse: { _ in
@@ -210,6 +210,40 @@ class AuthenticationProviderConfirmSigninTests: BasePluginTest {
             }
         }
     }
+
+    func testConfirmSignInRetryWithCodeMismatchException() async {
+        self.mockIdentityProvider = MockIdentityProvider(
+            mockRespondToAuthChallengeResponse: { _ in
+                throw RespondToAuthChallengeOutputError.codeMismatchException(
+                    .init(message: "Exception"))
+            })
+
+        do {
+            _ = try await plugin.confirmSignIn(challengeResponse: "code")
+            XCTFail("Should return an error if the result from service is invalid")
+        } catch {
+            guard case AuthError.service(_, _, let underlyingError) = error else {
+                XCTFail("Should produce service error instead of \(error)")
+                return
+            }
+            guard case .codeMismatch = (underlyingError as? AWSCognitoAuthError) else {
+                XCTFail("Underlying error should be codeMismatch \(error)")
+                return
+            }
+
+            self.mockIdentityProvider = MockIdentityProvider(
+                mockRespondToAuthChallengeResponse: { _ in
+                    return .testData()
+                })
+            do {
+                let confirmSignInResult = try await plugin.confirmSignIn(challengeResponse: "code")
+                XCTAssertTrue(confirmSignInResult.isSignedIn, "Signin result should be complete")
+            } catch {
+                XCTFail("Received failure with error \(error)")
+            }
+
+        }
+    }
     
     /// Test a confirmSignIn call with CodeExpiredException response from service
     ///

Original file line number	Diff line number	Diff line change
`@@ -13,11 +13,11 @@ enum SignInChallengeState: State {`
`13`	`13`
`14`	`14`	`case waitingForAnswer(RespondToAuthChallenge, SignInMethod)`
`15`	`15`
`16`		`- case verifying(RespondToAuthChallenge, String)`
	`16`	`+ case verifying(RespondToAuthChallenge, SignInMethod, String)`
`17`	`17`
`18`	`18`	`case verified`
`19`	`19`
`20`		`- case error(RespondToAuthChallenge, SignInError)`
	`20`	`+ case error(RespondToAuthChallenge, SignInMethod, SignInError)`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`extension SignInChallengeState {`