From f494bef9fd9703d734d17f63aea4d685c4f47423 Mon Sep 17 00:00:00 2001
From: Harshdeep Singh <6162866+harsh62@users.noreply.github.com>
Date: Mon, 24 Jun 2024 21:37:56 -0400
Subject: [PATCH] fix(Auth): Resolve AuthZ state correctly when in error state

---
 .../FetchAuthSessionOperationHelper.swift     | 20 ++++---
 ...AuthFetchSignInSessionOperationTests.swift | 53 +++++++++++++++++++
 2 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift
index d39751956a..b4a7aa9d4a 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift
@@ -66,23 +66,21 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
         authStateMachine: AuthStateMachine,
         forceRefresh: Bool) async throws -> AuthSession {
 
-            var event: AuthorizationEvent
             if forceRefresh || !credentials.areValid() {
-                if case .identityPoolWithFederation(
-                    let federatedToken,
-                    let identityId,
-                    _
-                ) = credentials {
-                    event = AuthorizationEvent(
-                        eventType: .startFederationToIdentityPool(federatedToken, identityId)
-                    )
-                } else {
+                var event: AuthorizationEvent
+                switch credentials {
+                case .identityPoolWithFederation(let federatedToken, let identityId, _):
+                    event = AuthorizationEvent(eventType: .startFederationToIdentityPool(federatedToken, identityId))
+                case .noCredentials:
+                    event = AuthorizationEvent(eventType: .fetchUnAuthSession)
+                case .userPoolOnly, .identityPoolOnly, .userPoolAndIdentityPool:
                     event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
                 }
                 await authStateMachine.send(event)
                 return try await listenForSession(authStateMachine: authStateMachine)
+            } else {
+                return credentials.cognitoSession
             }
-            return credentials.cognitoSession
         }
 
     func listenForSession(authStateMachine: AuthStateMachine) async throws -> AuthSession {
diff --git a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFetchSignInSessionOperationTests.swift b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFetchSignInSessionOperationTests.swift
index 9bee6c05b1..a67aa2522b 100644
--- a/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFetchSignInSessionOperationTests.swift
+++ b/AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/AuthorizationTests/AWSAuthFetchSignInSessionOperationTests.swift
@@ -584,6 +584,59 @@ class AWSAuthFetchSignInSessionOperationTests: BaseAuthorizationTests {
         }
     }
 
+    /// Test fetch session with authorization in error state
+    ///
+    /// - Given: An auth plugin with signedOut state
+    /// - When:
+    ///    - I invoke fetchAuthSession and mock notSignedIn for getTokens
+    /// - Then:
+    ///    - I should get an a valid session with the following details:
+    ///         - isSignedIn = false
+    ///         - aws credentails = valid values
+    ///         - identity id = valid values
+    ///         - cognito tokens = signedOut
+    ///
+    func testFetchSessionWithAuthorizationInErrorState() async throws {
+
+        let initialState = AuthState.configured(
+            AuthenticationState.signedOut(.testData),
+            AuthorizationState.error(.sessionError(.service(AuthError.unknown("error")), .noCredentials)))
+
+        let getId: MockIdentity.MockGetIdResponse = { _ in
+            return .init(identityId: "mockIdentityId")
+        }
+
+        let getCredentials: MockIdentity.MockGetCredentialsResponse = { _ in
+            let credentials = CognitoIdentityClientTypes.Credentials(accessKeyId: "accessKey",
+                                                                     expiration: Date(),
+                                                                     secretKey: "secret",
+                                                                     sessionToken: "session")
+            return .init(credentials: credentials, identityId: "responseIdentityID")
+        }
+
+        let plugin = configurePluginWith(identityPool: {
+            MockIdentity(mockGetIdResponse: getId,
+                         mockGetCredentialsResponse: getCredentials) },
+                                         initialState: initialState)
+
+        let session = try await plugin.fetchAuthSession(options: AuthFetchSessionRequest.Options())
+        XCTAssertFalse(session.isSignedIn)
+
+        let creds = try? (session as? AuthAWSCredentialsProvider)?.getAWSCredentials().get()
+        XCTAssertNotNil(creds?.accessKeyId)
+        XCTAssertNotNil(creds?.secretAccessKey)
+
+        let identityId = try? (session as? AuthCognitoIdentityProvider)?.getIdentityId().get()
+        XCTAssertNotNil(identityId)
+
+        let tokensResult = (session as? AuthCognitoTokensProvider)?.getCognitoTokens()
+        guard case .failure(let error) = tokensResult,
+              case .signedOut = error else {
+            XCTFail("Should return signed out error")
+            return
+        }
+    }
+
     /// Test signedOut state credential refresh
     ///
     /// - Given: Given an auth plugin with signedOut state and expired AWS credentials