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

Author: yaroluchko

Amplify/Categories/Auth/Models/AccessGroup.swift

@@ -16,14 +16,14 @@ public struct AccessGroup {
     }
 
     public static func none(migrateKeychainItemsOfUserSession: Bool) -> AccessGroup {
-        return .init(name: nil, migrateKeychainItems: migrateKeychainItemsOfUserSession)
+        return .init(migrateKeychainItems: migrateKeychainItemsOfUserSession)
     }
 
     public static var none: AccessGroup {
         return .none(migrateKeychainItemsOfUserSession: false)
     }
 
-    private init(name: String?, migrateKeychainItems: Bool) {
+    private init(name: String? = nil, migrateKeychainItems: Bool) {
         self.name = name
         self.migrateKeychainItems = migrateKeychainItems
     }

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

@@ -186,8 +186,11 @@ extension AWSCognitoAuthPlugin {
     }
 
     private func makeCredentialStore() -> AmplifyAuthCredentialStoreBehavior {
-        return AWSCognitoAuthCredentialStore(authConfiguration: authConfiguration, accessGroup: secureStoragePreferences?.accessGroup?.name,
-                                             migrateKeychainItemsOfUserSession: secureStoragePreferences?.accessGroup?.migrateKeychainItems ?? false)
+        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

@@ -46,8 +46,8 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
         return "awsCognitoAuthPlugin"
     }
 
-    /// Instantiates an instance of the AWSCognitoAuthPlugin with optionally custom network
-    /// preferences and custom secure storage preferences
+    /// Instantiates an instance of the AWSCognitoAuthPlugin with optional custom network
+    /// preferences and optional custom secure storage preferences
     /// - Parameters:
     ///   - networkPreferences: network preferences
     ///   - secureStoragePreferences: secure storage preferences

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

@@ -35,7 +35,11 @@ struct AWSCognitoAuthCredentialStore {
     private let userDefaults = UserDefaults.standard
     private let accessGroup: String?
 
-    init(authConfiguration: AuthConfiguration, accessGroup: String? = nil, migrateKeychainItemsOfUserSession: Bool = false) {
+    init(
+        authConfiguration: AuthConfiguration,
+        accessGroup: String? = nil,
+        migrateKeychainItemsOfUserSession: Bool = false
+    ) {
         self.authConfiguration = authConfiguration
         self.accessGroup = accessGroup
         if let accessGroup {
@@ -241,33 +245,14 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
             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
-            }
-        }
+        let oldService = oldAccessGroup != nil ? sharedService : service
+        let newService = accessGroup != nil ? sharedService : service
 
         do {
-            try oldKeychain._removeAll()
+            try KeychainStoreMigrator(oldService: oldService, newService: newService, oldAccessGroup: oldAccessGroup, newAccessGroup: accessGroup).migrate()
         } catch {
-            log.error("[AWSCognitoAuthCredentialStore] Error deleting all items from keychain under old access group after migration")
+            log.error("[AWSCognitoAuthCredentialStore] Migration has failed")
+            return
         }
 
         log.verbose("[AWSCognitoAuthCredentialStore] Migration of keychain items from old access group to new access group successful")

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

@@ -20,7 +20,12 @@ class AWSAuthFetchSessionTask: AuthFetchSessionTask, DefaultLogger {
         HubPayload.EventName.Auth.fetchSessionAPI
     }
 
-    init(_ request: AuthFetchSessionRequest, authStateMachine: AuthStateMachine, configuration: AuthConfiguration, forceReconfigure: Bool = false) {
+    init(
+        _ request: AuthFetchSessionRequest,
+        authStateMachine: AuthStateMachine,
+        configuration: AuthConfiguration,
+        forceReconfigure: Bool = false
+    ) {
         self.request = request
         self.authStateMachine = authStateMachine
         self.fetchAuthSessionHelper = FetchAuthSessionOperationHelper()

AmplifyPlugins/Core/AWSPluginsCore/Keychain/KeychainStore.swift

@@ -75,13 +75,16 @@ public struct KeychainStore: KeychainStoreBehavior {
     }
 
     public init(service: String) {
-        attributes = KeychainStoreAttributes(service: service)
-        log.verbose("[KeychainStore] Initialized keychain with service=\(service), attributes=\(attributes), accessGroup=")
+        self.init(service: service, accessGroup: nil)
     }
 
     public init(service: String, accessGroup: String? = nil) {
         attributes = KeychainStoreAttributes(service: service, accessGroup: accessGroup)
-        log.verbose("[KeychainStore] Initialized keychain with service=\(service), attributes=\(attributes), accessGroup=\(attributes.accessGroup ?? "")")
+        log.verbose(
+            "[KeychainStore] Initialized keychain with service=\(service), " +
+            "attributes=\(attributes), " +
+            "accessGroup=\(attributes.accessGroup ?? "No access group specified")"
+        )
     }
 
     @_spi(KeychainStore)

AmplifyPlugins/Core/AWSPluginsCore/Keychain/KeychainStoreMigrator.swift

@@ -0,0 +1,56 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import Amplify
+
+public struct KeychainStoreMigrator {
+    let oldAttributes: KeychainStoreAttributes
+    let newAttributes: KeychainStoreAttributes
+    
+    public init(oldService: String, newService: String, oldAccessGroup: String?, newAccessGroup: String?) {
+        self.oldAttributes = KeychainStoreAttributes(service: oldService, accessGroup: oldAccessGroup)
+        self.newAttributes = KeychainStoreAttributes(service: newService, accessGroup: newAccessGroup)
+    }
+    
+    public func migrate() throws {
+        log.verbose("[KeychainStoreMigrator] Starting to migrate items")
+
+        var updateQuery = oldAttributes.defaultGetQuery()
+
+        var updateAttributes = [String: Any]()
+        updateAttributes[KeychainStore.Constants.AttributeService] = newAttributes.service
+        updateAttributes[KeychainStore.Constants.AttributeAccessGroup] = newAttributes.accessGroup
+
+        // Remove any current items to avoid duplicate item error
+        try? KeychainStore(service: newAttributes.service, accessGroup: newAttributes.accessGroup)._removeAll()
+        
+        let updateStatus = SecItemUpdate(updateQuery as CFDictionary, updateAttributes as CFDictionary)
+        switch updateStatus {
+        case errSecSuccess:
+            break
+        case errSecItemNotFound:
+            log.verbose("[KeychainStoreMigrator] No items to migrate, keychain under new access group is cleared")
+        case errSecDuplicateItem:
+            log.verbose("[KeychainStoreMigrator] Duplicate items found, could not migrate")
+            return
+        default:
+            log.error("[KeychainStoreMigrator] Error of status=\(updateStatus) occurred when attempting to migrate items in keychain")
+            throw KeychainStoreError.securityError(updateStatus)
+        }
+
+        log.verbose("[KeychainStoreMigrator] Successfully migrated items to new service and access group")
+    }
+}
+
+extension KeychainStoreMigrator: DefaultLogger {
+    public static var log: Logger {
+        Amplify.Logging.logger(forNamespace: String(describing: self))
+    }
+
+    public nonisolated var log: Logger { Self.log }
+}