feat(auth): adding support for keychain sharing using app groups (#3947)

* feat(Auth) Keychain Sharing (App Reload Required)

* Remove migrateKeychainItemsOfUserSession bool from SecureStoragePreferences

* Reconfigure when fetching auth session if sharing keychain

* Update API dumps for new version

* Indentation, clean up, and batch migration to avoid inconsistent state

* Update API dumps for new version

* Addressing review comments: documentation, no more credentials valid check, only delete items if absolutely necessary

* Style fixes

---------

Co-authored-by: Yaro Luchko <[email protected]>
Co-authored-by: aws-amplify-ops <[email protected]>

Author: 3 people

Amplify/Categories/Auth/Models/AccessGroup.swift

@@ -0,0 +1,51 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+/// A structure representing an access group for managing keychain items.
+public struct AccessGroup {
+    /// The name of the access group.
+    public let name: String?
+    
+    /// A flag indicating whether to migrate keychain items.
+    public let migrateKeychainItems: Bool
+
+    /**
+     Initializes an `AccessGroup` with the specified name and migration option.
+     
+     - Parameter name: The name of the access group.
+     - Parameter migrateKeychainItemsOfUserSession: A flag indicating whether to migrate keychain items. Defaults to `false`.
+     */
+    public init(name: String, migrateKeychainItemsOfUserSession: Bool = false) {
+        self.init(name: name, migrateKeychainItems: migrateKeychainItemsOfUserSession)
+    }
+
+    /**
+     Creates an `AccessGroup` instance with no specified name.
+     
+     - Parameter migrateKeychainItemsOfUserSession: A flag indicating whether to migrate keychain items.
+     - Returns: An `AccessGroup` instance with the migration option set.
+     */
+    public static func none(migrateKeychainItemsOfUserSession: Bool) -> AccessGroup {
+        return .init(migrateKeychainItems: migrateKeychainItemsOfUserSession)
+    }
+
+    /**
+     A static property representing an `AccessGroup` with no name and no migration.
+     
+     - Returns: An `AccessGroup` instance with no name and the migration option set to `false`.
+     */
+    public static var none: AccessGroup {
+        return .none(migrateKeychainItemsOfUserSession: false)
+    }
+    
+    private init(name: String? = nil, migrateKeychainItems: Bool) {
+        self.name = name
+        self.migrateKeychainItems = migrateKeychainItems
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin+Configure.swift

@@ -186,7 +186,11 @@ extension AWSCognitoAuthPlugin {
     }
 
     private func makeCredentialStore() -> AmplifyAuthCredentialStoreBehavior {
-        AWSCognitoAuthCredentialStore(authConfiguration: authConfiguration)
+        return AWSCognitoAuthCredentialStore(
+            authConfiguration: authConfiguration,
+            accessGroup: secureStoragePreferences?.accessGroup?.name,
+            migrateKeychainItemsOfUserSession: secureStoragePreferences?.accessGroup?.migrateKeychainItems ?? false
+        )
     }
 
     private func makeLegacyKeychainStore(service: String) -> KeychainStoreBehavior {

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin.swift

@@ -35,6 +35,9 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
     /// The user network preferences for timeout and retry
     let networkPreferences: AWSCognitoNetworkPreferences?
 
+    /// The user secure storage preferences for access group
+    let secureStoragePreferences: AWSCognitoSecureStoragePreferences?
+
     @_spi(InternalAmplifyConfiguration)
     internal(set) public var jsonConfiguration: JSONValue?
 
@@ -43,15 +46,14 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
         return "awsCognitoAuthPlugin"
     }
 
-    /// Instantiates an instance of the AWSCognitoAuthPlugin.
-    public init() {
-        self.networkPreferences = nil
-    }
-
-    /// Instantiates an instance of the AWSCognitoAuthPlugin with custom network preferences
+    /// Instantiates an instance of the AWSCognitoAuthPlugin with optional custom network
+    /// preferences and optional custom secure storage preferences
     /// - Parameters:
     ///   - networkPreferences: network preferences
-    public init(networkPreferences: AWSCognitoNetworkPreferences) {
+    ///   - secureStoragePreferences: secure storage preferences
+    public init(networkPreferences: AWSCognitoNetworkPreferences? = nil,
+                secureStoragePreferences: AWSCognitoSecureStoragePreferences = AWSCognitoSecureStoragePreferences()) {
         self.networkPreferences = networkPreferences
+        self.secureStoragePreferences = secureStoragePreferences
     }
 }

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

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AWSCognitoAuthCredentialStore.swift

@@ -13,6 +13,7 @@ struct AWSCognitoAuthCredentialStore {
 
     // Credential store constants
     private let service = "com.amplify.awsCognitoAuthPlugin"
+    private let sharedService = "com.amplify.awsCognitoAuthPluginShared"
     private let sessionKey = "session"
     private let deviceMetadataKey = "deviceMetadata"
     private let deviceASFKey = "deviceASF"
@@ -25,14 +26,40 @@ struct AWSCognitoAuthCredentialStore {
     private var isKeychainConfiguredKey: String {
         "\(userDefaultsNameSpace).isKeychainConfigured"
     }
+    /// This UserDefaults Key is use to retrieve the stored access group to determine
+    /// which access group the migration should happen from
+    /// If none is found, the unshared service is used for migration and all items
+    /// under that service are queried
+    private var accessGroupKey: String {
+        "\(userDefaultsNameSpace).accessGroup"
+    }
 
     private let authConfiguration: AuthConfiguration
     private let keychain: KeychainStoreBehavior
     private let userDefaults = UserDefaults.standard
+    private let accessGroup: String?
 
-    init(authConfiguration: AuthConfiguration, accessGroup: String? = nil) {
+    init(
+        authConfiguration: AuthConfiguration,
+        accessGroup: String? = nil,
+        migrateKeychainItemsOfUserSession: Bool = false
+    ) {
         self.authConfiguration = authConfiguration
-        self.keychain = KeychainStore(service: service, accessGroup: accessGroup)
+        self.accessGroup = accessGroup
+        if let accessGroup {
+            self.keychain = KeychainStore(service: sharedService, accessGroup: accessGroup)
+        } else {
+            self.keychain = KeychainStore(service: service)
+        }
+        
+        let oldAccessGroup = retrieveStoredAccessGroup()
+        if migrateKeychainItemsOfUserSession {
+            try? migrateKeychainItemsToAccessGroup()
+        } else if oldAccessGroup == nil && oldAccessGroup != accessGroup {
+            try? KeychainStore(service: service)._removeAll()
+        }
+            
+        saveStoredAccessGroup()
 
         if !userDefaults.bool(forKey: isKeychainConfiguredKey) {
             try? clearAllCredentials()
@@ -181,6 +208,39 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
     private func clearAllCredentials() throws {
         try keychain._removeAll()
     }
+    
+    private func retrieveStoredAccessGroup() -> String? {
+        return userDefaults.string(forKey: accessGroupKey)
+    }
+    
+    private func saveStoredAccessGroup() {
+        if let accessGroup {
+            userDefaults.set(accessGroup, forKey: accessGroupKey)
+        } else {
+            userDefaults.removeObject(forKey: accessGroupKey)
+        }
+    }
+    
+    private func migrateKeychainItemsToAccessGroup() throws {
+        let oldAccessGroup = retrieveStoredAccessGroup()
+        
+        if oldAccessGroup == accessGroup {
+            log.info("[AWSCognitoAuthCredentialStore] Stored access group is the same as current access group, aborting migration")
+            return
+        }
+        
+        let oldService = oldAccessGroup != nil ? sharedService : service
+        let newService = accessGroup != nil ? sharedService : service
+        
+        do {
+            try KeychainStoreMigrator(oldService: oldService, newService: newService, oldAccessGroup: oldAccessGroup, newAccessGroup: accessGroup).migrate()
+        } catch {
+            log.error("[AWSCognitoAuthCredentialStore] Migration has failed")
+            return
+        }
+        
+        log.verbose("[AWSCognitoAuthCredentialStore] Migration of keychain items from old access group to new access group successful")
+    }
 
 }
 
@@ -204,3 +264,5 @@ private extension AWSCognitoAuthCredentialStore {
     }
 
 }
+
+extension AWSCognitoAuthCredentialStore: DefaultLogger { }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AWSCognitoSecureStoragePreferences.swift

@@ -0,0 +1,23 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import Amplify
+
+/// A struct to store preferences for how the plugin uses storage
+public struct AWSCognitoSecureStoragePreferences {
+
+    /// The access group that the keychain will use for auth items
+    public let accessGroup: AccessGroup?
+
+    /// Creates an intstance of AWSCognitoSecureStoragePreferences
+    /// - Parameters:
+    ///   - accessGroup: access group to be used
+    public init(accessGroup: AccessGroup? = nil) {
+        self.accessGroup = accessGroup
+    }
+}

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

@@ -13,19 +13,34 @@ 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,

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/ConfigurationTests/AWSCognitoAuthPluginAmplifyOutputsConfigTests.swift

@@ -123,4 +123,84 @@ class AWSCognitoAuthPluginAmplifyOutputsConfigTests: XCTestCase {
             XCTFail("Should not throw error. \(error)")
         }
     }
+    
+    /// Test Auth configuration with valid config for user pool and identity pool, with secure storage preferences
+    ///
+    /// - Given: Given valid config for user pool and identity pool with secure storage preferences
+    /// - When:
+    ///    - I configure auth with the given configuration and secure storage preferences
+    /// - Then:
+    ///    - I should not get any error while configuring auth
+    ///
+    func testConfigWithUserPoolAndIdentityPoolWithSecureStoragePreferences() throws {
+        let plugin = AWSCognitoAuthPlugin(
+            secureStoragePreferences: .init(
+                accessGroup: AccessGroup(name: "xx")
+            )
+        )
+        try Amplify.add(plugin: plugin)
+
+        let amplifyConfig = AmplifyOutputsData(auth: .init(
+                awsRegion: "us-east-1",
+                userPoolId: "xx",
+                userPoolClientId: "xx",
+                identityPoolId: "xx"))
+
+        do {
+            try Amplify.configure(amplifyConfig)
+
+            let escapeHatch = plugin.getEscapeHatch()
+            guard case .userPoolAndIdentityPool(let userPoolClient, let identityPoolClient) = escapeHatch else {
+                XCTFail("Expected .userPool, got \(escapeHatch)")
+                return
+            }
+            XCTAssertNotNil(userPoolClient)
+            XCTAssertNotNil(identityPoolClient)
+
+        } catch {
+            XCTFail("Should not throw error. \(error)")
+        }
+    }
+    
+    /// Test Auth configuration with valid config for user pool and identity pool, with network preferences and secure storage preferences
+    ///
+    /// - Given: Given valid config for user pool and identity pool, network preferences, and secure storage preferences
+    /// - When:
+    ///    - I configure auth with the given configuration, network preferences, and secure storage preferences
+    /// - Then:
+    ///    - I should not get any error while configuring auth
+    ///
+    func testConfigWithUserPoolAndIdentityPoolWithNetworkPreferencesAndSecureStoragePreferences() throws {
+        let plugin = AWSCognitoAuthPlugin(
+            networkPreferences: .init(
+                maxRetryCount: 2,
+                timeoutIntervalForRequest: 60,
+                timeoutIntervalForResource: 60),
+            secureStoragePreferences: .init(
+                accessGroup: AccessGroup(name: "xx")
+            )
+        )
+        try Amplify.add(plugin: plugin)
+
+        let amplifyConfig = AmplifyOutputsData(auth: .init(
+                awsRegion: "us-east-1",
+                userPoolId: "xx",
+                userPoolClientId: "xx",
+                identityPoolId: "xx"))
+
+        do {
+            try Amplify.configure(amplifyConfig)
+
+            let escapeHatch = plugin.getEscapeHatch()
+            guard case .userPoolAndIdentityPool(let userPoolClient, let identityPoolClient) = escapeHatch else {
+                XCTFail("Expected .userPool, got \(escapeHatch)")
+                return
+            }
+            XCTAssertNotNil(userPoolClient)
+            XCTAssertNotNil(identityPoolClient)
+
+        } catch {
+            XCTFail("Should not throw error. \(error)")
+        }
+    }
 }