fix(auth): don't remove items from keychain if it already has items (#4105)

Author: harsh62

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

@@ -56,7 +56,13 @@ struct AWSCognitoAuthCredentialStore {
         if migrateKeychainItemsOfUserSession {
             try? migrateKeychainItemsToAccessGroup()
         } else if oldAccessGroup == nil && oldAccessGroup != accessGroup {
-            try? KeychainStore(service: service)._removeAll()
+            // Only clear the old keychain if the shared keychain doesn't already have items.
+            // This prevents data loss when an app extension (e.g., widget) initializes before
+            // the main app has a chance to record the migration in UserDefaults, since
+            // UserDefaults is not shared between app and extensions.
+            if !sharedKeychainHasItems(accessGroup: accessGroup) {
+                try? KeychainStore(service: service)._removeAll()
+            }
         }
 
         saveStoredAccessGroup()
@@ -252,6 +258,15 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
             return
         }
 
+        // If the shared keychain already has items, migration has already occurred
+        // (likely by the main app). Skip migration to prevent data loss.
+        // This check is necessary because UserDefaults is not shared between app and extensions,
+        // so the extension may not know that migration already happened.
+        if sharedKeychainHasItems(accessGroup: accessGroup) {
+            log.info("[AWSCognitoAuthCredentialStore] Shared keychain already has items, migration already completed, aborting")
+            return
+        }
+
         let oldService = oldAccessGroup != nil ? sharedService : service
         let newService = accessGroup != nil ? sharedService : service
 
@@ -265,6 +280,17 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
         log.verbose("[AWSCognitoAuthCredentialStore] Migration of keychain items from old access group to new access group successful")
     }
 
+    /// Checks if the shared keychain (with the given access group) already contains items.
+    /// This is used to determine if migration has already occurred, which helps prevent
+    /// data loss when app extensions initialize with their own UserDefaults that don't
+    /// reflect the migration state recorded by the main app.
+    private func sharedKeychainHasItems(accessGroup: String?) -> Bool {
+        guard let accessGroup else { return false }
+
+        let sharedKeychain = KeychainStore(service: sharedService, accessGroup: accessGroup)
+        return (try? sharedKeychain._hasItems()) ?? false
+    }
+
 }
 
 /// Helpers for encode and decoding

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

@@ -43,4 +43,8 @@ class MockKeychainStoreBehavior: KeychainStoreBehavior {
     func _removeAll() throws {
         removeAllHandler?()
     }
+
+    func _hasItems() throws -> Bool {
+        return !data.isEmpty
+    }
 }

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/Support/DefaultConfig.swift

@@ -396,6 +396,10 @@ struct MockLegacyStore: KeychainStoreBehavior {
 
     }
 
+    func _hasItems() throws -> Bool {
+        return false
+    }
+
 }
 
 struct MockASF: AdvancedSecurityBehavior {

AmplifyPlugins/Core/AWSPluginsCore/Keychain/KeychainStore.swift

@@ -53,6 +53,12 @@ public protocol KeychainStoreBehavior {
     @_spi(KeychainStore)
     func _removeAll() throws
 
+    /// Checks if the Keychain contains any items for this service and access group.
+    /// This System Programming Interface (SPI) may have breaking changes in future updates.
+    /// - Returns: `true` if at least one item exists, `false` otherwise
+    @_spi(KeychainStore)
+    func _hasItems() throws -> Bool
+
 }
 
 public struct KeychainStore: KeychainStoreBehavior {
@@ -233,6 +239,29 @@ public struct KeychainStore: KeychainStoreBehavior {
         log.verbose("[KeychainStore] Successfully removed all items from keychain")
     }
 
+    /// Checks if the Keychain contains any items for this service and access group.
+    /// This System Programming Interface (SPI) may have breaking changes in future updates.
+    /// - Returns: `true` if at least one item exists, `false` otherwise
+    @_spi(KeychainStore)
+    public func _hasItems() throws -> Bool {
+        log.verbose("[KeychainStore] Checking if keychain has any items")
+        var query = attributes.defaultGetQuery()
+        query[Constants.MatchLimit] = Constants.MatchLimitOne
+
+        let status = SecItemCopyMatching(query as CFDictionary, nil)
+        switch status {
+        case errSecSuccess:
+            log.verbose("[KeychainStore] Keychain has items")
+            return true
+        case errSecItemNotFound:
+            log.verbose("[KeychainStore] Keychain has no items")
+            return false
+        default:
+            log.error("[KeychainStore] Error checking keychain items with status=\(status)")
+            throw KeychainStoreError.securityError(status)
+        }
+    }
+
 }
 
 extension KeychainStore {

AmplifyPlugins/Internal/Tests/InternalAWSPinpointUnitTests/Mocks/MockKeychainStore.swift

@@ -62,6 +62,10 @@ class MockKeychainStore: KeychainStoreBehavior {
         dataValues.removeAll()
     }
 
+    func _hasItems() throws -> Bool {
+        return !stringValues.isEmpty || !dataValues.isEmpty
+    }
+
     func resetCounters() {
         dataForKeyCount = 0
         stringForKeyCount = 0