Add ability to transform values on migrate

README.md

@@ -209,7 +209,7 @@ However, because the Keychain is effectively disk storage, there is no guarantee
 
 ### Migrating Existing Keychain Values into Valet
 
-Already using the Keychain and no longer want to maintain your own Keychain code? We feel you. That’s why we wrote `migrateObjects(matching query: [String : AnyHashable], removeOnCompletion: Bool)`. This method allows you to migrate all your existing Keychain entries to a Valet instance in one line. Just pass in a Dictionary with the `kSecClass`, `kSecAttrService`, and any other `kSecAttr*` attributes you use – we’ll migrate the data for you.
+Already using the Keychain and no longer want to maintain your own Keychain code? We feel you. That’s why we wrote `migrateObjects(matching query: [String : AnyHashable], removeOnCompletion: Bool)`. This method allows you to migrate all your existing Keychain entries to a Valet instance in one line. Just pass in a Dictionary with the `kSecClass`, `kSecAttrService`, and any other `kSecAttr*` attributes you use – we’ll migrate the data for you. If you need more control over how your data is migrated, use `migrateObjects(matching query: [String : AnyHashable], compactMap: (MigratableKeyValuePair<AnyHashable>) throws -> MigratableKeyValuePair<String>?)` to filter or remap key:value pairs as part of your migration.
 
 ### Integrating Valet into a macOS application
 

Sources/Valet/Internal/Keychain.swift

@@ -174,55 +174,55 @@ internal final class Keychain {
     }
 
     // MARK: Migration
-    
-    internal static func migrateObjects(matching query: [String : AnyHashable], into destinationAttributes: [String : AnyHashable], removeOnCompletion: Bool) throws {
+
+    internal static func migrateObjects(matching query: [String : AnyHashable], into destinationAttributes: [String : AnyHashable], compactMap: (MigratableKeyValuePair<AnyHashable>) throws -> MigratableKeyValuePair<String>?) throws {
         guard !query.isEmpty else {
             // Migration requires secItemQuery to contain values.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecMatchLimit as String] as? String as CFString? != kSecMatchLimitOne else {
             // Migration requires kSecMatchLimit to be set to kSecMatchLimitAll.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecReturnData as String] as? Bool != true else {
             // kSecReturnData is not supported in a migration query.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecReturnAttributes as String] as? Bool != false else {
             // Migration requires kSecReturnAttributes to be set to kCFBooleanTrue.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecReturnRef as String] as? Bool != true else {
             // kSecReturnRef is not supported in a migration query.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecReturnPersistentRef as String] as? Bool != false else {
             // Migration requires kSecReturnPersistentRef to be set to kCFBooleanTrue.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecClass as String] as? String as CFString? == kSecClassGenericPassword else {
             // Migration requires kSecClass to be set to kSecClassGenericPassword to avoid data loss.
             throw MigrationError.invalidQuery
         }
-        
+
         guard query[kSecAttrAccessControl as String] == nil else {
             // kSecAttrAccessControl is not supported in a migration query. Keychain items can not be migrated en masse from the Secure Enclave.
             throw MigrationError.invalidQuery
         }
-        
+
         var secItemQuery = query
         secItemQuery[kSecMatchLimit as String] = kSecMatchLimitAll
         secItemQuery[kSecReturnAttributes as String] = true
         secItemQuery[kSecReturnData as String] = false
         secItemQuery[kSecReturnRef as String] = false
         secItemQuery[kSecReturnPersistentRef as String] = true
-        
+
         let collection: Any = try SecItem.copy(matching: secItemQuery)
         let retrievedItemsToMigrate: [[String: AnyHashable]]
         if let singleMatch = collection as? [String : AnyHashable] {
@@ -234,15 +234,15 @@ internal final class Keychain {
         } else {
             throw MigrationError.dataToMigrateInvalid
         }
-        
+
         // Now that we have the persistent refs with attributes, get the data associated with each keychain entry.
         var retrievedItemsToMigrateWithData = [[String : AnyHashable]]()
         for retrievedItem in retrievedItemsToMigrate {
             guard let retrievedPersistentRef = retrievedItem[kSecValuePersistentRef as String] else {
                 throw KeychainError.couldNotAccessKeychain
 
             }
-            
+
             let retrieveDataQuery: [String : AnyHashable] = [
                 kSecValuePersistentRef as String : retrievedPersistentRef,
                 kSecReturnData as String : true
@@ -253,7 +253,7 @@ internal final class Keychain {
                 guard !data.isEmpty else {
                     throw MigrationError.dataToMigrateInvalid
                 }
-                
+
                 var retrievedItemToMigrateWithData = retrievedItem
                 retrievedItemToMigrateWithData[kSecValueData as String] = data
                 retrievedItemsToMigrateWithData.append(retrievedItemToMigrateWithData)
@@ -265,73 +265,94 @@ internal final class Keychain {
                 throw error
             }
         }
-        
+
         // Sanity check that we are capable of migrating the data.
-        var keysToMigrate = Set<String>()
+        var keyValuePairsToMigrate = [String: Data]()
         for keychainEntry in retrievedItemsToMigrateWithData {
-            guard let key = keychainEntry[kSecAttrAccount as String] as? String, key != Keychain.canaryKey else {
+            guard let key = keychainEntry[kSecAttrAccount as String] else {
+                throw MigrationError.keyToMigrateInvalid
+            }
+
+            guard key as? String != Keychain.canaryKey else {
                 // We don't care about this key. Move along.
                 continue
             }
-
-            guard !key.isEmpty else {
+
+            guard let data = keychainEntry[kSecValueData as String] as? Data else {
+                // This state should be impossible, per Apple's documentation for `kSecValueData`.
+                throw MigrationError.dataToMigrateInvalid
+            }
+
+            guard let migratablePair = try compactMap(MigratableKeyValuePair<AnyHashable>(key: key, value: data)) else {
+                // We don't care about this key. Move along.
+                continue
+            }
+
+            guard !migratablePair.key.isEmpty else {
                 throw MigrationError.keyToMigrateInvalid
             }
-            
-            guard !keysToMigrate.contains(key) else {
+
+            guard keyValuePairsToMigrate[migratablePair.key] == nil else {
                 throw MigrationError.duplicateKeyToMigrate
             }
-            
-            guard let data = keychainEntry[kSecValueData as String] as? Data, !data.isEmpty else {
+
+            guard !migratablePair.value.isEmpty else {
                 throw MigrationError.dataToMigrateInvalid
             }
 
-            if Keychain.performCopy(forKey: key, options: destinationAttributes) == errSecItemNotFound {
-                keysToMigrate.insert(key)
+            if Keychain.performCopy(forKey: migratablePair.key, options: destinationAttributes) == errSecItemNotFound {
+                keyValuePairsToMigrate[migratablePair.key] = migratablePair.value
             } else {
                 throw MigrationError.keyToMigrateAlreadyExistsInValet
             }
         }
-
-        // All looks good. Time to actually migrate.
-        var alreadyMigratedKeys = [String]()
-        func revertMigration() {
-            // Something has gone wrong. Remove all migrated items.
-            for alreadyMigratedKey in alreadyMigratedKeys {
-                try? Keychain.removeObject(forKey: alreadyMigratedKey, options: destinationAttributes)
-            }
-        }
-
-        for keychainEntry in retrievedItemsToMigrateWithData {
-            guard let key = keychainEntry[kSecAttrAccount as String] as? String else {
-                revertMigration()
-                throw MigrationError.keyToMigrateInvalid
-            }
 
-            guard let value = keychainEntry[kSecValueData as String] as? Data else {
-                revertMigration()
-                throw MigrationError.dataToMigrateInvalid
-            }
+        // Capture the keys in the destination prior to migration beginning.
+        let keysInKeychainPreMigration = Set(try Keychain.allKeys(options: destinationAttributes))
 
+        // All looks good. Time to actually migrate.
+        for keyValuePair in keyValuePairsToMigrate {
             do {
-                try Keychain.setObject(value, forKey: key, options: destinationAttributes)
-                alreadyMigratedKeys.append(key)
+                try Keychain.setObject(keyValuePair.value, forKey: keyValuePair.key, options: destinationAttributes)
             } catch {
-                revertMigration()
+                revertMigration(into: destinationAttributes, keysInKeychainPreMigration: keysInKeychainPreMigration)
                 throw error
             }
         }
-
+    }
+
+    internal static func migrateObjects(matching query: [String : AnyHashable], into destinationAttributes: [String : AnyHashable], removeOnCompletion: Bool) throws {
+        // Capture the keys in the destination prior to migration beginning.
+        let keysInKeychainPreMigration = Set(try Keychain.allKeys(options: destinationAttributes))
+
+        // Attempt migration.
+        try migrateObjects(matching: query, into: destinationAttributes) { keychainKeyValuePair in
+            guard let key = keychainKeyValuePair.key as? String else {
+                throw MigrationError.keyToMigrateInvalid
+            }
+            return MigratableKeyValuePair(key: key, value: keychainKeyValuePair.value)
+        }
+
         // Remove data if requested.
         if removeOnCompletion {
             do {
                 try Keychain.removeAllObjects(matching: query)
             } catch {
-                revertMigration()
+                revertMigration(into: destinationAttributes, keysInKeychainPreMigration: keysInKeychainPreMigration)
+
                 throw MigrationError.removalFailed
             }
 
             // We're done!
         }
     }
+
+    internal static func revertMigration(into destinationAttributes: [String : AnyHashable], keysInKeychainPreMigration: Set<String>) {
+        if let allKeysPostPotentiallyPartialMigration = try? Keychain.allKeys(options: destinationAttributes) {
+            let migratedKeys = allKeysPostPotentiallyPartialMigration.subtracting(keysInKeychainPreMigration)
+            migratedKeys.forEach { migratedKey in
+                try? Keychain.removeObject(forKey: migratedKey, options: destinationAttributes)
+            }
+        }
+    }
 }

Sources/Valet/MigratableKeyValuePair.swift

@@ -0,0 +1,134 @@
+//
+//  MigratableKeyValuePair.swift
+//  Valet
+//
+//  Created by Dan Federman on 5/20/20.
+//  Copyright © 2020 Square, Inc.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import Foundation
+
+/// A struct that represented a key:value pair that can be migrated.
+public struct MigratableKeyValuePair<Key: Hashable>: Hashable {
+
+    // MARK: Initialization
+
+    /// Creates a migratable key:value pair with the provided inputs.
+    /// - Parameters:
+    ///   - key: The key in the key:value pair.
+    ///   - value: The value in the key:value pair.
+    public init(key: Key, value: Data) {
+        self.key = key
+        self.value = value
+    }
+
+    /// Creates a migratable key:value pair with the provided inputs.
+    /// - Parameters:
+    ///   - key: The key in the key:value pair.
+    ///   - value: The desired value in the key:value pair, represented as a String.
+    public init(key: Key, value: String) {
+        self.key = key
+        self.value = Data(value.utf8)
+    }
+
+    // MARK: Public
+
+    /// The key in the key:value pair.
+    public let key: Key
+    /// The value in the key:value pair.
+    public let value: Data
+}
+
+// MARK: - Objective-C Compatibility
+
+@objc(VALMigratableKeyValuePairInput)
+public final class ObjectiveCCompatibilityMigratableKeyValuePairInput: NSObject {
+
+    // MARK: Initialization
+
+    internal init(key: Any, value: Data) {
+        self.key = key
+        self.value = value
+    }
+
+    // MARK: Public
+
+    /// The key in the key:value pair.
+    @objc
+    public let key: Any
+    /// The value in the key:value pair.
+    @objc
+    public let value: Data
+}
+
+@objc(VALMigratableKeyValuePairOutput)
+public class ObjectiveCCompatibilityMigratableKeyValuePairOutput: NSObject {
+
+    // MARK: Initialization
+
+    /// Creates a migratable key:value pair with the provided inputs.
+    /// - Parameters:
+    ///   - key: The key in the key:value pair.
+    ///   - value: The value in the key:value pair.
+    @objc
+    public init(key: String, value: Data) {
+        self.key = key
+        self.value = value
+        preventMigration = false
+    }
+
+    /// Creates a migratable key:value pair with the provided inputs.
+    /// - Parameters:
+    ///   - key: The key in the key:value pair.
+    ///   - stringValue: The desired value in the key:value pair, represented as a String.
+    @objc
+    public init(key: String, stringValue: String) {
+        self.key = key
+        self.value = Data(stringValue.utf8)
+        preventMigration = false
+    }
+
+    // MARK: Public Static Methods
+
+    /// A sentinal `ObjectiveCCompatibilityMigratableKeyValuePairOutput` that conveys that the migration should be prevented.
+    @available(swift, obsoleted: 1.0)
+    @objc
+    public static func preventMigration() -> ObjectiveCCompatibilityMigratableKeyValuePairOutput {
+        ObjectiveCCompatibilityPreventMigrationOutput()
+    }
+
+    // MARK: Public
+
+    /// The key in the key:value pair.
+    @objc
+    public let key: String
+    /// The value in the key:value pair.
+    @objc
+    public let value: Data
+
+    // MARK: Internal
+
+    internal fileprivate(set) var preventMigration: Bool
+
+}
+
+private final class ObjectiveCCompatibilityPreventMigrationOutput: ObjectiveCCompatibilityMigratableKeyValuePairOutput {
+
+    init() {
+        super.init(key: "", stringValue: "")
+        preventMigration = true
+    }
+
+}