feat(Auth) Keychain Sharing (App Reload Required)

* Remove migrateKeychainItemsOfUserSession bool from SecureStoragePreferences

Author: yaroluchko

Amplify/Categories/Auth/Models/AccessGroup.swift

@@ -0,0 +1,30 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+public struct AccessGroup {
+    public let name: String?
+    public let migrateKeychainItems: Bool
+
+    public init(name: String, migrateKeychainItemsOfUserSession: Bool = false) {
+        self.init(name: name, migrateKeychainItems: migrateKeychainItemsOfUserSession)
+    }
+
+    public static func none(migrateKeychainItemsOfUserSession: Bool) -> AccessGroup {
+        return .init(name: nil, migrateKeychainItems: migrateKeychainItemsOfUserSession)
+    }
+
+    public static var none: AccessGroup {
+        return .none(migrateKeychainItemsOfUserSession: false)
+    }
+
+    private init(name: String?, migrateKeychainItems: Bool) {
+        self.name = name
+        self.migrateKeychainItems = migrateKeychainItems
+    }
+}

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

@@ -177,7 +177,8 @@ 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 optionally custom network
+    /// preferences and 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/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,29 @@ struct AWSCognitoAuthCredentialStore {
     private var isKeychainConfiguredKey: String {
         "\(userDefaultsNameSpace).isKeychainConfigured"
     }
+    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)
+        }
+        
+        if migrateKeychainItemsOfUserSession {
+            try? migrateKeychainItemsToAccessGroup()
+        }
+            
+        try? saveStoredAccessGroup()
 
         if !userDefaults.bool(forKey: isKeychainConfiguredKey) {
             try? clearAllCredentials()
@@ -182,6 +198,81 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
     private func clearAllCredentials() throws {
         try keychain._removeAll()
     }
+    
+    private func retrieveStoredAccessGroup() throws -> String? {
+        return userDefaults.string(forKey: accessGroupKey)
+    }
+    
+    private func saveStoredAccessGroup() throws {
+        if let accessGroup {
+            userDefaults.set(accessGroup, forKey: accessGroupKey)
+        } else {
+            userDefaults.removeObject(forKey: accessGroupKey)
+        }
+    }
+    
+    private func migrateKeychainItemsToAccessGroup() throws {
+        let oldAccessGroup = try? retrieveStoredAccessGroup()
+        let oldKeychain: KeychainStoreBehavior
+        
+        if oldAccessGroup == accessGroup {
+            log.verbose("[AWSCognitoAuthCredentialStore] Stored access group is the same as current access group, aborting migration")
+            return
+        }
+        
+        if let oldAccessGroup {
+            oldKeychain = KeychainStore(service: sharedService, accessGroup: oldAccessGroup)
+        } else {
+            oldKeychain = KeychainStore(service: service)
+        }
+        
+        let authCredentialStoreKey = generateSessionKey(for: authConfiguration)
+        let authCredentialData: Data
+        let awsCredential: AmplifyCredentials
+        do {
+            authCredentialData = try oldKeychain._getData(authCredentialStoreKey)
+            awsCredential = try decode(data: authCredentialData)
+        } catch {
+            log.verbose("[AWSCognitoAuthCredentialStore] Could not retrieve previous credentials in keychain under old access group, nothing to migrate")
+            return
+        }
+        
+        guard awsCredential.areValid() else {
+            log.verbose("[AWSCognitoAuthCredentialStore] Credentials found are not valid (expired) in old access group keychain, aborting migration")
+            return
+        }
+        
+        let oldItems: [(key: String, value: Data)]
+        do {
+            oldItems = try oldKeychain._getAll()
+        } catch {
+            log.error("[AWSCognitoAuthCredentialStore] Error getting all items from keychain under old access group, aborting migration")
+            return
+        }
+
+        if oldItems.isEmpty {
+            log.verbose("[AWSCognitoAuthCredentialStore] No items in keychain under old access group, clearing keychain items under new access group")
+            return
+        }
+
+        for item in oldItems {
+            do {
+                try keychain._set(item.value, key: item.key)
+            } catch {
+                log.error("[AWSCognitoAuthCredentialStore] Error migrating one of the items, aborting migration: \(error)")
+                try? clearAllCredentials()
+                return
+            }
+        }
+        
+        do {
+            try oldKeychain._removeAll()
+        } catch {
+            log.error("[AWSCognitoAuthCredentialStore] Error deleting all items from keychain under old access group after migration")
+        }
+        
+        log.verbose("[AWSCognitoAuthCredentialStore] Migration of keychain items from old access group to new access group successful")
+    }
 
 }
 
@@ -205,3 +296,11 @@ private extension AWSCognitoAuthCredentialStore {
     }
 
 }
+
+extension AWSCognitoAuthCredentialStore: DefaultLogger {
+    public static var log: Logger {
+        Amplify.Logging.logger(forNamespace: String(describing: self))
+    }
+
+    public nonisolated var log: Logger { Self.log }
+}

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

@@ -0,0 +1,19 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import Amplify
+
+public struct AWSCognitoSecureStoragePreferences {
+
+    /// The access group that the keychain will use for auth items
+    public let accessGroup: AccessGroup?
+
+    public init(accessGroup: AccessGroup? = nil) {
+        self.accessGroup = accessGroup
+    }
+}

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/ActionTests/CredentialStore/MockCredentialStoreBehavior.swift

@@ -15,12 +15,15 @@ class MockKeychainStoreBehavior: KeychainStoreBehavior {
     typealias VoidHandler = () -> Void
 
     let data: String
+    let allData: [(key: String, value: Data)]
     let removeAllHandler: VoidHandler?
+    let mockKey: String = "mockKey"
 
     init(data: String,
          removeAllHandler: VoidHandler? = nil) {
         self.data = data
         self.removeAllHandler = removeAllHandler
+        self.allData = [(key: mockKey, value: Data(data.utf8))]
     }
 
     func _getString(_ key: String) throws -> String {
@@ -41,4 +44,8 @@ class MockKeychainStoreBehavior: KeychainStoreBehavior {
     func _removeAll() throws {
         removeAllHandler?()
     }
+    
+    func _getAll() throws -> [(key: String, value: Data)] {
+        return allData
+    }
 }

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)")
+        }
+    }
 }