fix(Auth): Fixing deviceSRP auth flow during MigrateAuth (#2584)

* fix(Auth): Fixing deviceSRP auth flow during MigrateAuth

* fix(Auth): Adding integration tests

* Update SignInState+Resolver.swift

Author: harsh62

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/DeviceSRPAuth/InitiateAuthDeviceSRP.swift

@@ -14,7 +14,6 @@ struct InitiateAuthDeviceSRP: Action {
     let identifier = "InitiateAuthDeviceSRP"
 
     let username: String
-    let deviceMetadata: DeviceMetadata
     let authResponse: SignInResponseBehavior
 
     func execute(withDispatcher dispatcher: EventDispatcher,
@@ -29,6 +28,11 @@ struct InitiateAuthDeviceSRP: Action {
             let gHexValue = srpEnv.srpConfiguration.gHexValue
             let srpKeyPair = srpClient.generateClientKeyPair()
 
+            // Get device metadata
+            let deviceMetadata = await DeviceMetadataHelper.getDeviceMetadata(
+                for: environment,
+                with: username)
+
             let srpStateData = SRPStateData(
                 username: username,
                 password: "",

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/DeviceSRPAuth/StartDeviceSRPFlow.swift

@@ -11,12 +11,13 @@ struct StartDeviceSRPFlow: Action {
 
     var identifier: String = "StartDeviceSRPFlow"
 
-    let srpStateData: SRPStateData
+    let username: Username
     let authResponse: SignInResponseBehavior
 
     func execute(withDispatcher dispatcher: EventDispatcher, environment: Environment) async {
         logVerbose("\(#fileID) Start execution", environment: environment)
-        let event = SignInEvent(id: UUID().uuidString, eventType: .respondDeviceSRPChallenge(srpStateData, authResponse))
+        let event = SignInEvent(id: UUID().uuidString, eventType: .respondDeviceSRPChallenge(
+            username, authResponse))
         logVerbose("\(#fileID) Sending event \(event.type)", environment: environment)
         await dispatcher.send(event)
     }
@@ -26,7 +27,7 @@ extension StartDeviceSRPFlow: CustomDebugDictionaryConvertible {
     var debugDictionary: [String: Any] {
         [
             "identifier": identifier,
-            "signInEventData": srpStateData.debugDictionary,
+            "username": username.masked(),
             "signInResponse": authResponse
         ]
     }

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

@@ -29,9 +29,9 @@ struct SignInEvent: StateMachineEvent {
 
         case respondPasswordVerifier(SRPStateData, InitiateAuthOutputResponse)
 
-        case initiateDeviceSRP(SignInResponseBehavior)
+        case initiateDeviceSRP(Username, SignInResponseBehavior)
 
-        case respondDeviceSRPChallenge(SRPStateData, SignInResponseBehavior)
+        case respondDeviceSRPChallenge(Username, SignInResponseBehavior)
 
         case respondDevicePasswordVerifier(SRPStateData, SignInResponseBehavior)
 

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/DeviceSRPState+Debug.swift

@@ -15,8 +15,8 @@ extension DeviceSRPState {
         switch self {
         case .notStarted:
             additionalMetadataDictionary = [:]
-        case .initiatingDeviceSRP(let signInEventData):
-            additionalMetadataDictionary = signInEventData.debugDictionary
+        case .initiatingDeviceSRP:
+            additionalMetadataDictionary = [:]
         case .cancelling:
             additionalMetadataDictionary = [:]
         case .respondingDevicePasswordVerifier(let srpStateData):

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DeviceSRPState.swift

@@ -10,7 +10,7 @@ import AWSCognitoIdentityProvider
 enum DeviceSRPState: State {
 
     case notStarted
-    case initiatingDeviceSRP(SRPStateData)
+    case initiatingDeviceSRP
     case respondingDevicePasswordVerifier(SRPStateData)
     case signedIn(SignedInData)
     case error(SignInError)
@@ -34,8 +34,8 @@ extension DeviceSRPState {
         switch (lhs, rhs) {
         case (.notStarted, .notStarted):
             return true
-        case (.initiatingDeviceSRP(let lhsData), .initiatingDeviceSRP(let rhsData)):
-            return lhsData == rhsData
+        case (.initiatingDeviceSRP, .initiatingDeviceSRP):
+            return true
         case (.cancelling, .cancelling):
             return true
         case (.respondingDevicePasswordVerifier(let lhsData), .respondingDevicePasswordVerifier(let rhsData)):

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

@@ -50,13 +50,12 @@ extension DeviceSRPState {
         private func resolveNotStarted(
             byApplying signInEvent: SignInEvent) -> StateResolution<DeviceSRPState> {
                 switch signInEvent.eventType {
-                case .respondDeviceSRPChallenge(let srpStateData, let authResponse):
+                case .respondDeviceSRPChallenge(let username, let authResponse):
                     let action = InitiateAuthDeviceSRP(
-                        username: srpStateData.username,
-                        deviceMetadata: srpStateData.deviceMetadata,
+                        username: username,
                         authResponse: authResponse)
                     return StateResolution(
-                        newState: DeviceSRPState.initiatingDeviceSRP(srpStateData),
+                        newState: DeviceSRPState.initiatingDeviceSRP,
                         actions: [action]
                     )
                 default:

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

@@ -89,10 +89,10 @@ extension SignInState {
                 }
 
                 if let signInEvent = event as? SignInEvent,
-                   case .initiateDeviceSRP(let challengeResponse) = signInEvent.eventType,
-                   case .respondingPasswordVerifier(let srpStateData) = srpSignInState {
+                   case .initiateDeviceSRP(let username, let challengeResponse) = signInEvent.eventType,
+                   case .respondingPasswordVerifier = srpSignInState {
                     let action = StartDeviceSRPFlow(
-                        srpStateData: srpStateData,
+                        username: username,
                         authResponse: challengeResponse)
                     return .init(newState: .resolvingDeviceSrpa(.notStarted),
                                  actions: [action])
@@ -125,6 +125,15 @@ extension SignInState {
                                  actions: [action])
                 }
 
+                if let signInEvent = event as? SignInEvent,
+                   case .initiateDeviceSRP(let username, let challengeResponse) = signInEvent.eventType {
+                    let action = StartDeviceSRPFlow(
+                        username: username,
+                        authResponse: challengeResponse)
+                    return .init(newState: .resolvingDeviceSrpa(.notStarted),
+                                 actions: [action])
+                }
+
                 let resolution = CustomSignInState.Resolver().resolve(
                     oldState: customSignInState, byApplying: event)
                 let signingInWithCustom = SignInState.signingInWithCustom(
@@ -151,6 +160,16 @@ extension SignInState {
                                  actions: [action])
                 }
 
+                if let signInEvent = event as? SignInEvent,
+                   case .initiateDeviceSRP(let username, let challengeResponse) = signInEvent.eventType,
+                   case .signingIn = migrateSignInState {
+                    let action = StartDeviceSRPFlow(
+                        username: username,
+                        authResponse: challengeResponse)
+                    return .init(newState: .resolvingDeviceSrpa(.notStarted),
+                                 actions: [action])
+                }
+
                 let resolution = MigrateSignInState.Resolver().resolve(
                     oldState: migrateSignInState, byApplying: event)
                 let signingInWithMigration = SignInState.signingInViaMigrateAuth(
@@ -166,6 +185,15 @@ extension SignInState {
                                  actions: [action])
                 }
 
+                if let signInEvent = event as? SignInEvent,
+                   case .initiateDeviceSRP(let username, let challengeResponse) = signInEvent.eventType {
+                    let action = StartDeviceSRPFlow(
+                        username: username,
+                        authResponse: challengeResponse)
+                    return .init(newState: .resolvingDeviceSrpa(.notStarted),
+                                 actions: [action])
+                }
+
                 // This could when we have nested challenges
                 // Example newPasswordRequired -> sms_mfa
                 if let signInEvent = event as? SignInEvent,
@@ -201,6 +229,15 @@ extension SignInState {
                     ), actions: [action])
                 }
 
+                if let signInEvent = event as? SignInEvent,
+                   case .initiateDeviceSRP(let username, let challengeResponse) = signInEvent.eventType {
+                    let action = StartDeviceSRPFlow(
+                        username: username,
+                        authResponse: challengeResponse)
+                    return .init(newState: .resolvingDeviceSrpa(.notStarted),
+                                 actions: [action])
+                }
+
                 if let signInEvent = event as? SignInEvent,
                    case .confirmDevice(let signedInData) = signInEvent.eventType {
                     let action = ConfirmDevice(signedInData: signedInData)

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

@@ -119,7 +119,7 @@ struct UserPoolSignInHelper {
                 case .smsMfa, .customChallenge, .newPasswordRequired:
                     return SignInEvent(eventType: .receivedChallenge(respondToAuthChallenge))
                 case .deviceSrpAuth:
-                    return SignInEvent(eventType: .initiateDeviceSRP(response))
+                    return SignInEvent(eventType: .initiateDeviceSRP(username, response))
                 default:
                     let message = "UnSupported challenge response \(challengeName)"
                     let error = SignInError.unknown(message: message)

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

@@ -80,8 +80,7 @@ class AWSAuthSignInTask: AuthSignInTask {
                     throw error
                 }
             case .error(let error):
-                let error = AuthError.unknown("Sign in reached an error state", error)
-                throw error
+                throw error.authError
 
             case .signingIn(let signInState):
                 guard let result = try UserPoolSignInHelper.checkNextStep(signInState) else {

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

@@ -542,6 +542,49 @@ class AWSAuthSignInPluginTests: BasePluginTest {
         }
     }
 
+    /// Test a signIn with deviceSRP response
+    ///
+    /// - Given: Given an auth plugin with mocked service. Mock unknown response for signIn result
+    ///
+    /// - When:
+    ///    - I invoke signIn
+    /// - Then:
+    ///    - I should get a .service error, because the response doesn't return valid data
+    ///
+    func testSignInWithNextDeviceSRP() async {
+
+        self.mockIdentityProvider = MockIdentityProvider(mockInitiateAuthResponse: { _ in
+            InitiateAuthOutputResponse(
+                authenticationResult: .none,
+                challengeName: .deviceSrpAuth,
+                challengeParameters: InitiateAuthOutputResponse.validChalengeParams,
+                session: "someSession")
+        }, mockRespondToAuthChallengeResponse: { _ in
+            RespondToAuthChallengeOutputResponse(
+                authenticationResult: .none,
+                challengeName: .devicePasswordVerifier,
+                challengeParameters: ["paramKey": "value"],
+                session: "session")
+        })
+
+        let pluginOptions = AWSAuthSignInOptions(validationData: ["somekey": "somevalue"],
+                                                 metadata: ["somekey": "somevalue"],
+                                                 authFlowType: .userPassword)
+        let options = AuthSignInRequest.Options(pluginOptions: pluginOptions)
+        do {
+            let result = try await plugin.signIn(
+                username: "username",
+                password: "password",
+                options: options)
+            XCTFail("Should not produce result - \(result)")
+        } catch {
+            guard case AuthError.service = error else {
+                XCTFail("Should produce as service error")
+                return
+            }
+        }
+    }
+
     // MARK: - Service error for initiateAuth
 
     /// Test a signIn with `InternalErrorException` from service