fix(auth): Device binding add retry incase of device not found (#2699)

Author: royjit

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/SRPAuth/VerifyPasswordSRP.swift

@@ -25,7 +25,6 @@ struct VerifyPasswordSRP: Action {
                  environment: Environment) async {
 
         logVerbose("\(#fileID) Starting execution", environment: environment)
-
         let inputUsername = stateData.username
         var username = inputUsername
         var deviceMetadata = DeviceMetadata.noData
@@ -37,7 +36,6 @@ struct VerifyPasswordSRP: Action {
 
             username = parameters["USERNAME"] ?? inputUsername
             let userIdForSRP = parameters["USER_ID_FOR_SRP"] ?? inputUsername
-
             let saltHex = try saltHex(parameters)
             let secretBlockString = try secretBlockString(parameters)
             let secretBlock = try secretBlock(secretBlockString)
@@ -73,16 +71,42 @@ struct VerifyPasswordSRP: Action {
             let event = SignInEvent(
                 eventType: .throwPasswordVerifierError(error)
             )
+            logVerbose("\(#fileID) Sending event \(event)",
+                       environment: environment)
+            await dispatcher.send(event)
+        } catch let error where deviceNotFound(error: error, deviceMetadata: deviceMetadata) {
+            logVerbose("\(#fileID) Received device not found \(error)", environment: environment)
+            // Remove the saved device details and retry password verify
+            await DeviceMetadataHelper.removeDeviceMetaData(of: username, environment: environment)
+            let event = SignInEvent(eventType: .retryRespondPasswordVerifier(stateData, authResponse))
+            logVerbose("\(#fileID) Sending event \(event)",
+                       environment: environment)
             await dispatcher.send(event)
         } catch {
             logVerbose("\(#fileID) SRPSignInError Generic \(error)", environment: environment)
             let authError = SignInError.service(error: error)
             let event = SignInEvent(
                 eventType: .throwAuthError(authError)
             )
+            logVerbose("\(#fileID) Sending event \(event)",
+                       environment: environment)
             await dispatcher.send(event)
         }
     }
+
+    func deviceNotFound(error: Error, deviceMetadata: DeviceMetadata) -> Bool {
+
+        // If deviceMetadata was not send, the error returned is not from device not found.
+        if case .noData = deviceMetadata {
+            return false
+        }
+
+        if let serviceError: RespondToAuthChallengeOutputError = error.internalAWSServiceError(),
+           case .resourceNotFoundException = serviceError {
+            return true
+        }
+        return false
+    }
 }
 
 extension VerifyPasswordSRP: DefaultLogger { }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/VerifySignInChallenge.swift

@@ -54,20 +54,43 @@ struct VerifySignInChallenge: Action {
             logVerbose("\(#fileID) Sending event \(responseEvent)",
                        environment: environment)
             await dispatcher.send(responseEvent)
+        } catch let error where deviceNotFound(error: error, deviceMetadata: deviceMetadata) {
+            logVerbose("\(#fileID) Received device not found \(error)", environment: environment)
+            // Remove the saved device details and retry verify challenge
+            await DeviceMetadataHelper.removeDeviceMetaData(of: username, environment: environment)
+            let event = SignInChallengeEvent(
+                eventType: .retryVerifyChallengeAnswer(confirmSignEventData)
+            )
+            logVerbose("\(#fileID) Sending event \(event)", environment: environment)
+            await dispatcher.send(event)
         } catch let error as SignInError {
             let errorEvent = SignInEvent(eventType: .throwAuthError(error))
             logVerbose("\(#fileID) Sending event \(errorEvent)",
                        environment: environment)
             await dispatcher.send(errorEvent)
         } catch {
-            let error = SignInError.invalidServiceResponse(message: error.localizedDescription)
+            let error = SignInError.service(error: error)
             let errorEvent = SignInEvent(eventType: .throwAuthError(error))
             logVerbose("\(#fileID) Sending event \(errorEvent)",
                        environment: environment)
             await dispatcher.send(errorEvent)
         }
     }
 
+    func deviceNotFound(error: Error, deviceMetadata: DeviceMetadata) -> Bool {
+
+        // If deviceMetadata was not send, the error returned is not from device not found.
+        if case .noData = deviceMetadata {
+            return false
+        }
+
+        if let serviceError: RespondToAuthChallengeOutputError = error.internalAWSServiceError(),
+           case .resourceNotFoundException = serviceError {
+            return true
+        }
+        return false
+    }
+
 }
 
 extension VerifySignInChallenge: CustomDebugDictionaryConvertible {

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Events/SignInChallengeEvent.swift

@@ -15,6 +15,8 @@ struct SignInChallengeEvent: StateMachineEvent {
 
         case verifyChallengeAnswer(ConfirmSignInEventData)
 
+        case retryVerifyChallengeAnswer(ConfirmSignInEventData)
+
         case verified
 
     }
@@ -28,6 +30,7 @@ struct SignInChallengeEvent: StateMachineEvent {
         case .verified: return "SignInChallengeEvent.verified"
         case .verifyChallengeAnswer: return "SignInChallengeEvent.verifyChallengeAnswer"
         case .waitForAnswer: return "SignInChallengeEvent.waitForAnswer"
+        case .retryVerifyChallengeAnswer: return "SignInChallengeEvent.retryVerifyChallengeAnswer"
         }
     }
 

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Events/SignInEvent.swift

@@ -29,6 +29,8 @@ struct SignInEvent: StateMachineEvent {
 
         case respondPasswordVerifier(SRPStateData, InitiateAuthOutputResponse)
 
+        case retryRespondPasswordVerifier(SRPStateData, InitiateAuthOutputResponse)
+
         case initiateDeviceSRP(Username, SignInResponseBehavior)
 
         case respondDeviceSRPChallenge(Username, SignInResponseBehavior)
@@ -72,6 +74,7 @@ struct SignInEvent: StateMachineEvent {
         case .throwAuthError: return "SignInEvent.throwAuthError"
         case .receivedChallenge: return "SignInEvent.receivedChallenge"
         case .verifySMSChallenge: return "SignInEvent.verifySMSChallenge"
+        case .retryRespondPasswordVerifier: return "SignInEvent.retryRespondPasswordVerifier"
         }
     }
 
@@ -104,7 +107,8 @@ extension SignInEvent.EventType: Equatable {
             (.cancelSRPSignIn, .cancelSRPSignIn),
             (.throwAuthError, .throwAuthError),
             (.receivedChallenge, .receivedChallenge),
-            (.verifySMSChallenge, .verifySMSChallenge):
+            (.verifySMSChallenge, .verifySMSChallenge),
+            (.retryRespondPasswordVerifier, .retryRespondPasswordVerifier):
             return true
         default: return false
         }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/SRP/SRPSignInState+Resolver.swift

@@ -102,6 +102,13 @@ extension SRPSignInState {
             byApplying signInEvent: SignInEvent)
         -> StateResolution<SRPSignInState> {
             switch signInEvent.eventType {
+            case .retryRespondPasswordVerifier(let srpStateData, let authResponse):
+                let action = VerifyPasswordSRP(stateData: srpStateData,
+                                               authResponse: authResponse)
+                return StateResolution(
+                    newState: SRPSignInState.respondingPasswordVerifier(srpStateData),
+                    actions: [action]
+                )
             case .finalizeSignIn(let signedInData):
                 return .init(newState: .signedIn(signedInData),
                              actions: [SignInComplete(signedInData: signedInData)])

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/SignIn/SignInChallengeState+Resolver.swift

@@ -42,6 +42,17 @@ extension SignInChallengeState {
 
             case .verifying(let challenge, let signInMethod, _):
 
+                if case .retryVerifyChallengeAnswer(let answerEventData) = event.isChallengeEvent {
+                    let action = VerifySignInChallenge(
+                        challenge: challenge,
+                        confirmSignEventData: answerEventData,
+                        signInMethod: signInMethod)
+                    return .init(
+                        newState: .verifying(challenge, signInMethod, answerEventData.answer),
+                        actions: [action]
+                    )
+                }
+
                 if case .finalizeSignIn(let signedInData) = event.isSignInEvent {
                     return .init(newState: .verified,
                                  actions: [SignInComplete(signedInData: signedInData)])

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Helpers/UserPoolSignInHelper.swift

@@ -80,17 +80,13 @@ struct UserPoolSignInHelper: DefaultLogger {
         environment: UserPoolEnvironment) async throws -> StateMachineEvent {
 
             let client = try environment.cognitoUserPoolFactory()
-
-            do {
-                let response = try await client.respondToAuthChallenge(input: request)
-                let event = try self.parseResponse(response, for: username, signInMethod: signInMethod)
-                return event
-            } catch {
-                let authError = SignInError.service(error: error)
-                return SignInEvent(eventType: .throwAuthError(authError))
-            }
+            let response = try await client.respondToAuthChallenge(input: request)
+            let event = try self.parseResponse(response, for: username, signInMethod: signInMethod)
+            return event
         }
 
+
+
     static func parseResponse(
         _ response: SignInResponseBehavior,
         for username: String,

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/ActionTests/InitiateAuthSRP/VerifyPasswordSRPTests.swift

@@ -419,52 +419,52 @@ class VerifyPasswordSRPTests: XCTestCase {
         await waitForExpectations(timeout: 0.1)
     }
 
-//    /// Test verify password retry on device not found
-//    ///
-//    /// - Given: VerifyPasswordSRP action with mocked cognito client and configuration
-//    /// - When:
-//    ///    - I invoke the action with valid input and mock empty device not found error from Cognito
-//    /// - Then:
-//    ///    - Should send an event with retryRespondPasswordVerifier
-//    ///
-//    func testPasswordVerifierWithDeviceNotFound() async {
-//
-//        let identityProviderFactory: CognitoFactory = {
-//            MockIdentityProvider(
-//                mockRespondToAuthChallengeResponse: { _ in
-//                    throw RespondToAuthChallengeOutputError.resourceNotFoundException(
-//                        ResourceNotFoundException()
-//                    )
-//                })
-//        }
-//
-//        let environment = Defaults.makeDefaultAuthEnvironment(
-//            userPoolFactory: identityProviderFactory)
-//
-//        let data = InitiateAuthOutputResponse.validTestData
-//        let action = VerifyPasswordSRP(stateData: SRPStateData.testData,
-//                                       authResponse: data)
-//
-//        let passwordVerifierError = expectation(description: "passwordVerifierError")
-//
-//        let dispatcher = MockDispatcher { event in
-//            defer { passwordVerifierError.fulfill() }
-//
-//            guard let event = event as? SignInEvent else {
-//                XCTFail("Expected event to be SignInEvent but got \(event)")
-//                return
-//            }
-//
-//            guard case .retryRespondPasswordVerifier = event.eventType
-//            else {
-//                XCTFail("Should receive retryRespondPasswordVerifier")
-//                return
-//            }
-//        }
-//
-//        await action.execute(withDispatcher: dispatcher, environment: environment)
-//        await waitForExpectations(timeout: 0.1)
-//    }
+    /// Test verify password retry on device not found
+    ///
+    /// - Given: VerifyPasswordSRP action with mocked cognito client and configuration
+    /// - When:
+    ///    - I invoke the action with valid input and mock empty device not found error from Cognito
+    /// - Then:
+    ///    - Should send an event with retryRespondPasswordVerifier
+    ///
+    func testPasswordVerifierWithDeviceNotFound() async {
+
+        let identityProviderFactory: CognitoFactory = {
+            MockIdentityProvider(
+                mockRespondToAuthChallengeResponse: { _ in
+                    throw RespondToAuthChallengeOutputError.resourceNotFoundException(
+                        ResourceNotFoundException()
+                    )
+                })
+        }
+
+        let environment = Defaults.makeDefaultAuthEnvironment(
+            userPoolFactory: identityProviderFactory)
+
+        let data = InitiateAuthOutputResponse.validTestData
+        let action = VerifyPasswordSRP(stateData: SRPStateData.testData,
+                                       authResponse: data)
+
+        let passwordVerifierError = expectation(description: "passwordVerifierError")
+
+        let dispatcher = MockDispatcher { event in
+            defer { passwordVerifierError.fulfill() }
+
+            guard let event = event as? SignInEvent else {
+                XCTFail("Expected event to be SignInEvent but got \(event)")
+                return
+            }
+
+            guard case .retryRespondPasswordVerifier = event.eventType
+            else {
+                XCTFail("Should receive retryRespondPasswordVerifier")
+                return
+            }
+        }
+
+        await action.execute(withDispatcher: dispatcher, environment: environment)
+        await waitForExpectations(timeout: 0.1)
+    }
 
     /// Test  successful response from the VerifyPasswordSRP for confirmDevice
     ///