chore(auth): Cleanup on FetchAuthSessionOperationhelper (#2714)

Author: royjit

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/RefreshAuthorizationSession/InitializeRefreshSession.swift

@@ -44,7 +44,7 @@ struct InitializeRefreshSession: Action {
                                                     region: config.region,
                                                     poolId: config.poolId)
             if isForceRefresh ||
-                tokens.doesExpire(in: FetchAuthSessionOperationHelper.expiryBufferInSeconds) {
+                tokens.doesExpire(in: AmplifyCredentials.expiryBufferInSeconds) {
                 event = .init(eventType: .refreshCognitoUserPoolWithIdentityId(signedInData, identityID))
             } else {
                 event = .init(eventType: .refreshAWSCredentialsWithUserPool(identityID,

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AmplifyCredentials+CognitoSession.swift

@@ -9,7 +9,7 @@ import Foundation
 import Amplify
 
 extension AmplifyCredentials {
-
+    static let expiryBufferInSeconds = TimeInterval.seconds(2 * 60)
     var cognitoSession: AWSAuthCognitoSession {
 
         switch self {
@@ -48,4 +48,36 @@ extension AmplifyCredentials {
             return AuthCognitoSignedOutSessionHelper.makeSessionWithNoGuestAccess()
         }
     }
+
+    func areValid() -> Bool {
+        return self != .noCredentials &&
+        !doesExpire(in: Self.expiryBufferInSeconds)
+    }
+
+    private func doesExpire(in expiryBuffer: TimeInterval) -> Bool {
+        var doesExpire = false
+        switch self {
+
+        case .userPoolOnly(signedInData: let data):
+            doesExpire = data.cognitoUserPoolTokens.doesExpire(in: expiryBuffer)
+
+        case .identityPoolOnly(identityID: _, credentials: let awsCredentials):
+            doesExpire = awsCredentials.doesExpire(in: expiryBuffer)
+
+        case .userPoolAndIdentityPool(signedInData: let data,
+                                      identityID: _,
+                                      credentials: let awsCredentials):
+            doesExpire = (
+                data.cognitoUserPoolTokens.doesExpire(in: expiryBuffer) ||
+                awsCredentials.doesExpire(in: expiryBuffer)
+            )
+
+        case .identityPoolWithFederation(_, _, let awsCredentials):
+            doesExpire = awsCredentials.doesExpire(in: expiryBuffer)
+
+        case .noCredentials:
+            doesExpire = true
+        }
+        return doesExpire
+    }
 }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/FetchAuthSessionOperationHelper.swift

@@ -10,7 +10,7 @@ import Amplify
 
 class FetchAuthSessionOperationHelper: DefaultLogger {
 
-    static let expiryBufferInSeconds = TimeInterval.seconds(2 * 60)
+
 
     typealias FetchAuthSessionCompletion = (Result<AuthSession, AuthError>) -> Void
 
@@ -35,8 +35,8 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
 
         case .sessionEstablished(let credentials):
             log.verbose("Session exists, checking validity")
-            return try await postAuthSessionEvent(
-                forCredential: credentials,
+            return try await refreshIfRequired(
+                existingCredentials: credentials,
                 authStateMachine: authStateMachine,
                 forceRefresh: forceRefresh)
 
@@ -46,8 +46,8 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
                 let session = AuthCognitoSignedInSessionHelper.makeExpiredSignedInSession()
                 return session
             } else if case .sessionError(_, let credentials) = error {
-                return try await postAuthSessionEvent(
-                    forCredential: credentials,
+                return try await refreshIfRequired(
+                    existingCredentials: credentials,
                     authStateMachine: authStateMachine,
                     forceRefresh: forceRefresh)
             } else {
@@ -62,60 +62,28 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
         }
     }
 
-    func postAuthSessionEvent(
-        forCredential credentials: AmplifyCredentials,
+    func refreshIfRequired(
+        existingCredentials credentials: AmplifyCredentials,
         authStateMachine: AuthStateMachine,
         forceRefresh: Bool) async throws -> AuthSession {
-            switch credentials {
-
-            case .userPoolOnly(signedInData: let data):
-                if data.cognitoUserPoolTokens.doesExpire(in: Self.expiryBufferInSeconds) ||
-                    forceRefresh {
-                    let event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
-                    await authStateMachine.send(event)
-                    return try await listenForSession(authStateMachine: authStateMachine)
-                } else {
-                    return credentials.cognitoSession
-                }
-
-            case .identityPoolOnly(identityID: _, credentials: let awsCredentials):
-                if awsCredentials.doesExpire(in: Self.expiryBufferInSeconds) ||
-                    forceRefresh {
-                    let event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
-                    await authStateMachine.send(event)
-                    return try await listenForSession(authStateMachine: authStateMachine)
-                } else {
-                   return credentials.cognitoSession
-                }
 
-            case .userPoolAndIdentityPool(signedInData: let data,
-                                          identityID: _,
-                                          credentials: let awsCredentials):
-                if  data.cognitoUserPoolTokens.doesExpire(in: Self.expiryBufferInSeconds) ||
-                        awsCredentials.doesExpire(in: Self.expiryBufferInSeconds) ||
-                        forceRefresh {
-                    let event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
-                    await authStateMachine.send(event)
-                    return try await listenForSession(authStateMachine: authStateMachine)
+            var event: AuthorizationEvent
+            if forceRefresh || !credentials.areValid() {
+                if case .identityPoolWithFederation(
+                    let federatedToken,
+                    let identityId,
+                    _
+                ) = credentials {
+                    event = AuthorizationEvent(
+                        eventType: .startFederationToIdentityPool(federatedToken, identityId)
+                    )
                 } else {
-                    return credentials.cognitoSession
+                    event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
                 }
-
-            case .identityPoolWithFederation(let federatedToken, let identityId, let awsCredentials):
-                if awsCredentials.doesExpire() || forceRefresh {
-                    let event = AuthorizationEvent.init(
-                        eventType: .startFederationToIdentityPool(federatedToken, identityId))
-                    await authStateMachine.send(event)
-                    return try await listenForSession(authStateMachine: authStateMachine)
-                } else {
-                    return credentials.cognitoSession
-                }
-
-            case .noCredentials:
-                let event = AuthorizationEvent(eventType: .refreshSession(forceRefresh))
                 await authStateMachine.send(event)
                 return try await listenForSession(authStateMachine: authStateMachine)
             }
+            return credentials.cognitoSession
         }
 
     func listenForSession(authStateMachine: AuthStateMachine) async throws -> AuthSession {
@@ -155,8 +123,8 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
         switch error {
         case .sessionError(let fetchError, let credentials):
             return try sessionResultWithFetchError(fetchError,
-                                        authenticationState: authenticationState,
-                                        existingCredentials: credentials)
+                                                   authenticationState: authenticationState,
+                                                   existingCredentials: credentials)
         case .sessionExpired:
             let session = AuthCognitoSignedInSessionHelper.makeExpiredSignedInSession()
             return session
@@ -172,8 +140,8 @@ class FetchAuthSessionOperationHelper: DefaultLogger {
     }
 
     func sessionResultWithFetchError(_ error: FetchSessionError,
-                                authenticationState: AuthenticationState,
-                                existingCredentials: AmplifyCredentials)
+                                     authenticationState: AuthenticationState,
+                                     existingCredentials: AmplifyCredentials)
     throws -> AuthSession {
 
         var isSignedIn = false