From 8dd74ed1be74ad0f8a7f15d36007f310aa140eb8 Mon Sep 17 00:00:00 2001
From: Yaro Luchko <yaluchko@amazon.com>
Date: Mon, 12 Aug 2024 17:02:53 -0700
Subject: [PATCH] Reconfigure when fetching auth session if sharing keychain

---
 .../AWSCognitoAuthPlugin+ClientBehavior.swift         |  6 +++++-
 .../Task/AWSAuthFetchSessionTask.swift                | 11 ++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift
index 75bed80e89..70530e76e0 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/ClientBehavior/AWSCognitoAuthPlugin+ClientBehavior.swift
@@ -114,7 +114,11 @@ extension AWSCognitoAuthPlugin: AuthCategoryBehavior {
     public func fetchAuthSession(options: AuthFetchSessionRequest.Options?) async throws -> AuthSession {
         let options = options ?? AuthFetchSessionRequest.Options()
         let request = AuthFetchSessionRequest(options: options)
-        let task = AWSAuthFetchSessionTask(request, authStateMachine: authStateMachine)
+        let forceReconfigure = secureStoragePreferences?.accessGroup?.name != nil
+        let task = AWSAuthFetchSessionTask(request,
+                                           authStateMachine: authStateMachine,
+                                           configuration: authConfiguration,
+                                           forceReconfigure: forceReconfigure)
         return try await taskQueue.sync {
             return try await task.value
         } as! AuthSession
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthFetchSessionTask.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthFetchSessionTask.swift
index 7a00941ad9..0350f08e5a 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthFetchSessionTask.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthFetchSessionTask.swift
@@ -13,20 +13,29 @@ class AWSAuthFetchSessionTask: AuthFetchSessionTask, DefaultLogger {
     private let authStateMachine: AuthStateMachine
     private let fetchAuthSessionHelper: FetchAuthSessionOperationHelper
     private let taskHelper: AWSAuthTaskHelper
+    private let configuration: AuthConfiguration
+    private let forceReconfigure: Bool
 
     var eventName: HubPayloadEventName {
         HubPayload.EventName.Auth.fetchSessionAPI
     }
 
-    init(_ request: AuthFetchSessionRequest, authStateMachine: AuthStateMachine) {
+    init(_ request: AuthFetchSessionRequest, authStateMachine: AuthStateMachine, configuration: AuthConfiguration, forceReconfigure: Bool = false) {
         self.request = request
         self.authStateMachine = authStateMachine
         self.fetchAuthSessionHelper = FetchAuthSessionOperationHelper()
         self.taskHelper = AWSAuthTaskHelper(authStateMachine: authStateMachine)
+        self.configuration = configuration
+        self.forceReconfigure = forceReconfigure
     }
 
     func execute() async throws -> AuthSession {
         log.verbose("Starting execution")
+        if forceReconfigure {
+            log.verbose("Reconfiguring auth state machine for keychain sharing")
+            let event = AuthEvent(eventType: .reconfigure(configuration))
+            await authStateMachine.send(event)
+        }
         await taskHelper.didStateMachineConfigured()
         let doesNeedForceRefresh = request.options.forceRefresh
         return try await fetchAuthSessionHelper.fetch(authStateMachine,