fix(Auth): Making improvements to federation flow (#2488)

harsh62 · web-flow · commit cbd49eaa712c · 2022-10-24T12:48:26.000-04:00
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/Federation/ClearFederationToIdentityPool.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/Federation/ClearFederationToIdentityPool.swift
@@ -18,8 +18,18 @@ struct ClearFederationToIdentityPool: Action {
 
         let event: StateMachineEvent
         do {
-            try await credentialStoreClient?.deleteData(type: .amplifyCredentials)
-            event = AuthenticationEvent.init(eventType: .clearedFederationToIdentityPool)
+            // Check if we have credentials are of type federation, otherwise throw an error
+            if case .amplifyCredentials(let credentials) = try await credentialStoreClient?.fetchData(
+                type: .amplifyCredentials),
+               case .identityPoolWithFederation = credentials {
+
+                try await credentialStoreClient?.deleteData(type: .amplifyCredentials)
+                event = AuthenticationEvent.init(eventType: .clearedFederationToIdentityPool)
+            } else {
+                let error = AuthenticationError.service(message: "Unable to find credentials to clear for federation.")
+                event = AuthenticationEvent.init(eventType: .error(error))
+                logError("\(#fileID) Sending event \(event.type) with error \(error)", environment: environment)
+            }
 
         } catch {
             let error = AuthenticationError.unknown(message: "Unable to clear credential store: \(error)")
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AmplifyCredentials+CognitoSession.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AmplifyCredentials+CognitoSession.swift
@@ -31,8 +31,13 @@ extension AmplifyCredentials {
                 identityIdResult: .success(identityId),
                 awsCredentialsResult: .success(awsCredentials),
                 cognitoTokensResult: .failure(
-                    .signedOut("Cognito tokens unavailable when federated to identity pool",
-                               "Clear federation and sign to get user pool tokens.", nil)))
+                    .invalidState(
+                        "Users Federated to Identity Pool do not have User Pool access.",
+                        "To access User Pool data, you must use a Sign In method.",
+                        nil
+                    )
+                )
+            )
         case .userPoolAndIdentityPool(let signedInData, let identityID, let credentials):
             return AWSAuthCognitoSession(
                 isSignedIn: true,
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/ClearFederationOperationHelper.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/ClearFederationOperationHelper.swift
@@ -14,11 +14,18 @@ struct ClearFederationOperationHelper {
 
         let currentState = await authStateMachine.currentState
 
-        if case .configured(let authNState, let authZState) = currentState,
-           case .federatedToIdentityPool = authNState,
-           case .sessionEstablished = authZState {
+        guard case .configured(let authNState, let authZState) = currentState else {
+            let authError = AuthError.invalidState(
+                "Clearing of federation failed.",
+                AuthPluginErrorConstants.invalidStateError, nil)
+            throw authError
+        }
+
+        switch (authNState, authZState) {
+        case (.federatedToIdentityPool, .sessionEstablished),
+            (.error, .error):
             try await startClearingFederation(with: authStateMachine)
-        } else {
+        default:
             let authError = AuthError.invalidState(
                 "Clearing of federation failed.",
                 AuthPluginErrorConstants.invalidStateError, nil)
@@ -27,9 +34,9 @@ struct ClearFederationOperationHelper {
     }
 
     private func startClearingFederation(with authStateMachine: AuthStateMachine) async throws {
-        let stateSequences = await authStateMachine.listen()
         let event = AuthenticationEvent.init(eventType: .clearFederationToIdentityPool)
         await authStateMachine.send(event)
+        let stateSequences = await authStateMachine.listen()
         for await state in stateSequences {
             guard  case .configured(let authNState, _) = state else {
                 continue
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift
@@ -31,7 +31,34 @@ class FetchAuthSessionOperationHelper {
             await authStateMachine.send(event)
             return try await listenForSession(authStateMachine: authStateMachine)
         case .sessionEstablished(let credentials):
+            return try await postAuthSessionEvent(
+                forCredential: credentials,
+                authStateMachine: authStateMachine,
+                forceRefresh: forceRefresh)
+        case .error(let error):
+            if case .sessionExpired = error {
+                let session = AuthCognitoSignedInSessionHelper.makeExpiredSignedInSession()
+                return session
+            } else if case .sessionError(_, let credentials) = error {
+                return try await postAuthSessionEvent(
+                    forCredential: credentials,
+                    authStateMachine: authStateMachine,
+                    forceRefresh: forceRefresh)
+            } else {
+                let event = AuthorizationEvent(eventType: .fetchUnAuthSession)
+                await authStateMachine.send(event)
+                return try await listenForSession(authStateMachine: authStateMachine)
+            }
+
+        default:
+            return try await listenForSession(authStateMachine: authStateMachine)
+        }
+    }
 
+    func postAuthSessionEvent(
+        forCredential credentials: AmplifyCredentials,
+        authStateMachine: AuthStateMachine,
+        forceRefresh: Bool) async throws -> AuthSession {
             switch credentials {
 
             case .userPoolOnly(signedInData: let data):
@@ -82,34 +109,9 @@ class FetchAuthSessionOperationHelper {
                 await authStateMachine.send(event)
                 return try await listenForSession(authStateMachine: authStateMachine)
             }
-
-        case .error(let error):
-            if case .sessionExpired = error {
-                let session = AuthCognitoSignedInSessionHelper.makeExpiredSignedInSession()
-                return session
-            } else if case .sessionError(_, let amplifyCredentials) = error {
-                var event: AuthorizationEvent
-                // If there was no credentials before, we try to get the unauth session
-                // else, we try to refresh the credentials.
-                if case .noCredentials = amplifyCredentials {
-                    event = AuthorizationEvent(eventType: .fetchUnAuthSession)
-                } else {
-                    event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
-                }
-                await authStateMachine.send(event)
-                return try await listenForSession(authStateMachine: authStateMachine)
-            } else {
-                let event = AuthorizationEvent(eventType: .fetchUnAuthSession)
-                await authStateMachine.send(event)
-                return try await listenForSession(authStateMachine: authStateMachine)
-            }
-
-        default:
-            return try await listenForSession(authStateMachine: authStateMachine)
         }
-    }
 
-    func listenForSession(authStateMachine: AuthStateMachine) async throws ->AuthSession {
+    func listenForSession(authStateMachine: AuthStateMachine) async throws -> AuthSession {
 
         let stateSequences = await authStateMachine.listen()
         for await state in stateSequences {
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Errors/FetchSessionError.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Errors/FetchSessionError.swift
@@ -77,7 +77,7 @@ extension FetchSessionError: AuthErrorConvertible {
                 AmplifyErrorMessages.reportBugToAWS())
         case .federationNotSupportedDuringRefresh:
             return .unknown(
-                "Federation triggered during refresh session that is not supported. \(AmplifyErrorMessages.reportBugToAWS())")
+                "Refreshing credentials from federationToIdentityPool is not supported \(AmplifyErrorMessages.reportBugToAWS())")
         case .service(let error):
             return .service(
                 "Service error occurred",
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/AuthorizationState.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/AuthorizationState.swift
@@ -19,8 +19,10 @@ enum AuthorizationState: State {
 
     case clearingFederation
 
-    case federatingToIdentityPool(FetchAuthSessionState,
-                                  FederatedToken)
+    case federatingToIdentityPool(
+        FetchAuthSessionState,
+        FederatedToken,
+        existingCredentials: AmplifyCredentials)
 
     case fetchingUnAuthSession(FetchAuthSessionState)
 
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/AuthorizationState+Debug.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/AuthorizationState+Debug.swift
@@ -24,7 +24,7 @@ extension AuthorizationState: CustomDebugDictionaryConvertible {
                                             "refreshState": state.debugDictionary]
         case .fetchingUnAuthSession(let state),
                 .fetchingAuthSessionWithUserPool(let state, _),
-                .federatingToIdentityPool(let state, _):
+                .federatingToIdentityPool(let state, _, _):
             additionalMetadataDictionary = state.debugDictionary
 
         case .error(let error):
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/Authentication/AuthenticationState+Resolver.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/Authentication/AuthenticationState+Resolver.swift
@@ -122,6 +122,13 @@ extension AuthenticationState {
                 } else if let authZEvent = event.isAuthorizationEvent,
                          case .startFederationToIdentityPool = authZEvent {
                     return .init(newState: .federatingToIdentityPool)
+                } else if let authEvent = event as? AuthenticationEvent,
+                   case .clearFederationToIdentityPool = authEvent.eventType {
+                    return .init(
+                        newState: .clearingFederation,
+                        actions: [
+                            ClearFederationToIdentityPool()
+                        ])
                 } else {
                     return .from(oldState)
                 }
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/Authorization/AuthorizationState+Resolver.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/Authorization/AuthorizationState+Resolver.swift
@@ -46,7 +46,10 @@ extension AuthorizationState {
                         federatedToken: federatedToken,
                         developerProvidedIdentityId: identityId)
                     return .init(
-                        newState: .federatingToIdentityPool(.notStarted, federatedToken),
+                        newState: .federatingToIdentityPool(
+                            .notStarted,
+                            federatedToken,
+                            existingCredentials: .noCredentials),
                         actions: [action])
                 }
 
@@ -69,7 +72,10 @@ extension AuthorizationState {
                         federatedToken: federatedToken,
                         developerProvidedIdentityId: identityId)
                     return .init(
-                        newState: .federatingToIdentityPool(.notStarted, federatedToken),
+                        newState: .federatingToIdentityPool(
+                            .notStarted,
+                            federatedToken,
+                            existingCredentials: credentials),
                         actions: [action])
                 }
 
@@ -89,7 +95,8 @@ extension AuthorizationState {
 
                 return .from(oldState)
 
-            case .federatingToIdentityPool(let fetchSessionState, let federatedToken):
+            case .federatingToIdentityPool(
+                let fetchSessionState, let federatedToken, let credentials):
 
                 if case .fetched(let identityID,
                                  let credentials) = event.isAuthorizationEvent {
@@ -103,7 +110,7 @@ extension AuthorizationState {
                 }
 
                 if case .receivedSessionError(let error) = event.isAuthorizationEvent {
-                    return .init(newState: .error(.sessionError(error, .noCredentials)))
+                    return .init(newState: .error(.sessionError(error, credentials)))
                 }
 
                 if case .throwError(let error) = event.isAuthorizationEvent {
@@ -112,8 +119,12 @@ extension AuthorizationState {
 
                 let resolver = FetchAuthSessionState.Resolver()
                 let resolution = resolver.resolve(oldState: fetchSessionState, byApplying: event)
-                return .init(newState: .federatingToIdentityPool(resolution.newState, federatedToken),
-                             actions: resolution.actions)
+                return .init(
+                    newState: .federatingToIdentityPool(
+                        resolution.newState,
+                        federatedToken,
+                        existingCredentials: credentials),
+                    actions: resolution.actions)
 
             case .signingIn:
                 if let authEvent = event.isAuthenticationEvent {
@@ -247,13 +258,24 @@ extension AuthorizationState {
                 return .from(oldState)
 
             case .error(let error):
+                var existingCredentials: AmplifyCredentials = .noCredentials
+                if case .sessionError(_, let credentials) = error {
+                    existingCredentials = credentials
+                }
+
                 if let authNEvent = event.isAuthenticationEvent {
-                    if case .signInRequested = authNEvent {
+
+                    switch authNEvent {
+                    case .signInRequested:
                         return .from(.signingIn)
-                    } else if case .signOutRequested = authNEvent {
+                    case .signOutRequested:
                         return .from(.signingOut(nil))
-                    } else if case .cancelSignIn = authNEvent {
+                    case .cancelSignIn:
                         return .from(.configured)
+                    case .clearFederationToIdentityPool:
+                        return .from(.clearingFederation)
+                    default: break
+
                     }
                 }
                 if case .fetchUnAuthSession = event.isAuthorizationEvent {
@@ -269,7 +291,10 @@ extension AuthorizationState {
                         federatedToken: federatedToken,
                         developerProvidedIdentityId: identityId)
                     return .init(
-                        newState: .federatingToIdentityPool(.notStarted, federatedToken),
+                        newState: .federatingToIdentityPool(
+                            .notStarted,
+                            federatedToken,
+                            existingCredentials: existingCredentials),
                         actions: [action])
                 }
 
diff --git a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/HubEventTests/AuthHubEventHandlerTests.swift b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/HubEventTests/AuthHubEventHandlerTests.swift
@@ -278,6 +278,7 @@ class AuthHubEventHandlerTests: XCTestCase {
             }
         }
 
+        _ = try await plugin.federateToIdentityPool(withProviderToken: "something", for: .facebook)
         try await plugin.clearFederationToIdentityPool()
 
         await waitForExpectations(timeout: networkTimeout)
diff --git a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/Support/DefaultConfig.swift b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/Support/DefaultConfig.swift
@@ -230,23 +230,50 @@ enum Defaults {
 
 struct MockCredentialStoreOperationClient: CredentialStoreStateBehavior {
 
+    let mockAmplifyStore = MockAmplifyStore()
+
     func fetchData(type: CredentialStoreDataType) async throws -> CredentialStoreData {
-        switch type {
-        case .amplifyCredentials:
-            return .amplifyCredentials(.testData)
-        case .deviceMetadata(username: let username):
-            return .deviceMetadata(.metadata(.init(
-                deviceKey: "key",
-                deviceGroupKey: "key",
-                deviceSecret: "secret")), username)
-        case .asfDeviceId(username: let username):
-            return .asfDeviceId("id", username)
-        }
 
+        do {
+            switch type {
+            case .amplifyCredentials:
+                let amplifyCredentials = try mockAmplifyStore.retrieveCredential()
+                return .amplifyCredentials(amplifyCredentials)
+            case .deviceMetadata(username: let username):
+                let deviceMetadata = try mockAmplifyStore.retrieveDevice(for: username)
+                return .deviceMetadata(deviceMetadata, username)
+            case .asfDeviceId(username: let username):
+                let device = try mockAmplifyStore.retrieveASFDevice(for: username)
+                return .asfDeviceId(device, username)
+            }
+        }
+        catch KeychainStoreError.itemNotFound {
+            switch type {
+            case .amplifyCredentials:
+                return .amplifyCredentials(.testData)
+            case .deviceMetadata(username: let username):
+                return .deviceMetadata(.metadata(.init(
+                    deviceKey: "key",
+                    deviceGroupKey: "key",
+                    deviceSecret: "secret")), username)
+            case .asfDeviceId(username: let username):
+                return .asfDeviceId("id", username)
+            }
+        }
+        catch {
+            fatalError()
+        }
     }
 
     func storeData(data: CredentialStoreData) async throws {
-
+        switch data {
+        case .amplifyCredentials(let amplifyCredentials):
+            try mockAmplifyStore.saveCredential(amplifyCredentials)
+        case .deviceMetadata(let deviceMetadata, let username):
+            try mockAmplifyStore.saveDevice(deviceMetadata, for: username)
+        case .asfDeviceId(let string, let username):
+            try mockAmplifyStore.saveASFDevice(string, for: username)
+        }
     }
 
     func deleteData(type: CredentialStoreDataType) async throws {
@@ -266,7 +293,7 @@ class MockAmplifyStore: AmplifyAuthCredentialStoreBehavior {
     func retrieveCredential() throws -> AmplifyCredentials {
         guard let data = Self.dict.getValue(forKey: credentialsKey),
               let cred = (try? JSONDecoder().decode(AmplifyCredentials.self, from: data)) else {
-            return AmplifyCredentials.noCredentials
+            throw KeychainStoreError.itemNotFound
         }
         return cred
     }
@@ -287,7 +314,7 @@ class MockAmplifyStore: AmplifyAuthCredentialStoreBehavior {
     func retrieveDevice(for username: String) throws -> DeviceMetadata {
         guard let data = Self.dict.getValue(forKey: username),
               let device = (try? JSONDecoder().decode(DeviceMetadata.self, from: data)) else {
-            return .noData
+            throw KeychainStoreError.itemNotFound
         }
         return device
     }
diff --git a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFederationToIdentityPoolTests.swift b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFederationToIdentityPoolTests.swift

Original file line number	Diff line number	Diff line change
`@@ -278,6 +278,7 @@ class AuthHubEventHandlerTests: XCTestCase {`
`278`	`278`	`}`
`279`	`279`	`}`
`280`	`280`
	`281`	`+ _ = try await plugin.federateToIdentityPool(withProviderToken: "something", for: .facebook)`
`281`	`282`	`try await plugin.clearFederationToIdentityPool()`
`282`	`283`
`283`	`284`	`await waitForExpectations(timeout: networkTimeout)`