Return value after storing SecDataConvertible

dm-zharov · Apr 30, 2024 · 5e3b0b5 · 5e3b0b5
1 parent 1c58405
commit 5e3b0b5
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 33 deletions.
diff --git a/Sources/SwiftSecurity/Keychain/Keychain.swift b/Sources/SwiftSecurity/Keychain/Keychain.swift
@@ -184,6 +184,15 @@ extension Keychain: SecDataStore {
     public func store<T: SecDataConvertible>(_ data: T, query: SecItemQuery<GenericPassword>, accessPolicy: AccessPolicy = .default) throws {
         try store(.data(data.rawRepresentation), query: query, accessPolicy: accessPolicy)
     }
+
+    public func store<T: SecDataConvertible>(
+        _ data: T,
+        returning returnType: SecReturnType,
+        query: SecItemQuery<GenericPassword>,
+        accessPolicy: AccessPolicy = .default
+    ) throws -> SecValue<GenericPassword>? {
+        try store(.data(data.rawRepresentation), returning: returnType, query: query, accessPolicy: accessPolicy)
+    }
 
     public func retrieve<T: SecDataConvertible>(_ query: SecItemQuery<GenericPassword>, authenticationContext: LAContext? = nil) throws -> T? {
         if let value = try retrieve(.data, query: query, authenticationContext: authenticationContext), case let .data(data) = value {
@@ -334,12 +343,31 @@ extension Keychain: SecIdentityStore {
 // MARK: - Private
 
 private extension Keychain {
-    func store<SecItem>(_ value: SecValue<SecItem>, query: SecItemQuery<SecItem>, accessPolicy: AccessPolicy) throws {
+    @discardableResult
+    func store<SecItem>(
+        _ value: SecValue<SecItem>,
+        returning returnType: SecReturnType = [],
+        query: SecItemQuery<SecItem>,
+        accessPolicy: AccessPolicy = .default
+    ) throws -> SecValue<SecItem>? {
         var query = query
         query[.accessGroup] = accessGroup.rawValue
         query[.accessControl] = try accessPolicy.accessControl
         query[.accessible] = accessPolicy.accessibility
 
+        if returnType.contains(.data) {
+            query[kSecReturnData as String] = true
+        }
+        if returnType.contains(.info) {
+            query[kSecReturnAttributes as String] = true
+        }
+        if returnType.contains(.reference) {
+            query[kSecReturnRef as String] = true
+        }
+        if returnType.contains(.persistentReference) {
+            query[kSecReturnPersistentRef as String] = true
+        }
+
         switch value {
         case .data(let data):
             query[kSecValueData as String] = data
@@ -349,9 +377,35 @@ private extension Keychain {
             throw SwiftSecurityError.invalidParameter
         }
 
-        switch SecItemAdd(query.rawValue as CFDictionary, nil) {
+        var result: AnyObject?
+        switch SecItemAdd(query.rawValue as CFDictionary, &result) {
         case errSecSuccess:
-            return
+            switch returnType {
+            case .data:
+                if let data = result as? Data {
+                    return .data(data)
+                } else {
+                    return nil
+                }
+            case .reference:
+                if let result {
+                    return .reference(result)
+                } else {
+                    return nil
+                }
+            case .persistentReference:
+                if let data = result as? Data {
+                    return .persistentReference(data)
+                } else {
+                    return nil
+                }
+            default:
+                if let attributes = result as? [String: Any] {
+                    return .dictionary(SecItemInfo<SecItem>(rawValue: attributes))
+                } else {
+                    return nil
+                }
+            }
         case let status:
             throw SwiftSecurityError(rawValue: status)
         }

diff --git a/Sources/SwiftSecurity/Keychain/SecItemStore/SecItemStore.swift b/Sources/SwiftSecurity/Keychain/SecItemStore/SecItemStore.swift
@@ -16,22 +16,46 @@ public protocol SecItemStore {
     func removeAll() throws
 }
 
-// MARK: - SecData
+public extension SecItemStore {
+    func info<SecItem>(for query: SecItemQuery<SecItem>, authenticationContext: LAContext? = nil) throws -> SecItemInfo<SecItem>? {
+        if let value = try retrieve(.info, query: query, authenticationContext: authenticationContext), case let .dictionary(info) = value {
+            return info
+        } else {
+            return nil
+        }
+    }
+}
+
+// MARK: - Data
 
 public protocol SecDataStore: SecItemStore {
-    // MARK: - Generic
+    // MARK: - Generic Password
 
-    func store<T: SecDataConvertible>(_ data: T, query: SecItemQuery<GenericPassword>, accessPolicy: AccessPolicy) throws
+    func store<T: SecDataConvertible>(_ data: T, returning returnType: SecReturnType, query: SecItemQuery<GenericPassword>, accessPolicy: AccessPolicy) throws -> SecValue<GenericPassword>?
     func retrieve<T: SecDataConvertible>(_ query: SecItemQuery<GenericPassword>, authenticationContext: LAContext?) throws -> T?
     func remove(_ query: SecItemQuery<GenericPassword>) throws -> Bool
 
-    // MARK: - Internet
+    // MARK: - Internet Password
 
     func store<T: SecDataConvertible>(_ data: T, query: SecItemQuery<InternetPassword>, accessPolicy: AccessPolicy) throws
     func retrieve<T: SecDataConvertible>(_ query: SecItemQuery<InternetPassword>, authenticationContext: LAContext?) throws -> T?
     func remove(_ query: SecItemQuery<InternetPassword>) throws -> Bool
 }
 
+public extension SecDataStore {
+    func store<T: SecDataConvertible>(_ data: T, query: SecItemQuery<GenericPassword>, accessPolicy: AccessPolicy) throws {
+        try self.store(data, returning: [], query: query, accessPolicy: accessPolicy)
+    }
+
+    func retrieve(_ query: SecItemQuery<GenericPassword>) throws -> Data? {
+        try self.retrieve<Data>(query, authenticationContext: nil)
+    }
+
+    func retrieve(_ query: SecItemQuery<InternetPassword>) throws -> Data? {
+        try self.retrieve<Data>(query, authenticationContext: nil)
+    }
+}
+
 // MARK: - SecKey
 
 public protocol SecKeyStore: SecItemStore {
@@ -56,23 +80,3 @@ public protocol SecIdentityStore: SecItemStore {
     func retrieve(_ query: SecItemQuery<SecIdentity>, authenticationContext: LAContext?) throws -> SecIdentity?
     func remove(_ query: SecItemQuery<SecIdentity>) throws -> Bool
 }
-
-// MARK: - Convenient
-
-public extension SecDataStore {
-    func info<SecItem>(for query: SecItemQuery<SecItem>, authenticationContext: LAContext? = nil) throws -> SecItemInfo<SecItem>? {
-        if let value = try retrieve(.info, query: query, authenticationContext: authenticationContext), case let .dictionary(info) = value {
-            return info
-        } else {
-            return nil
-        }
-    }
-
-    func retrieve(_ query: SecItemQuery<GenericPassword>) throws -> Data? {
-        try self.retrieve<Data>(query, authenticationContext: nil)
-    }
-
-    func retrieve(_ query: SecItemQuery<InternetPassword>) throws -> Data? {
-        try self.retrieve<Data>(query, authenticationContext: nil)
-    }
-}
diff --git a/Tests/SwiftSecurityTests/AccessPolicyTests.swift b/Tests/SwiftSecurityTests/AccessPolicyTests.swift
@@ -13,35 +13,35 @@ import Security
 final class AccessPolicyTests: XCTestCase {
     func testAccessibility() throws {
         do {
-            let accessPolicy = SecAccessPolicy(.whenPasscodeSetThisDeviceOnly)
+            let accessPolicy = AccessPolicy(.whenPasscodeSetThisDeviceOnly)
             XCTAssertEqual(accessPolicy.accessibility, String(kSecAttrAccessibleWhenPasscodeSetThisDeviceOnly))
             XCTAssertNil(try accessPolicy.accessControl)
         }
         do {
-            let accessPolicy = SecAccessPolicy(.whenUnlocked)
+            let accessPolicy = AccessPolicy(.whenUnlocked)
             XCTAssertEqual(accessPolicy.accessibility, String(kSecAttrAccessibleWhenUnlocked))
             XCTAssertNil(try accessPolicy.accessControl)
         }
         do {
-            let accessPolicy = SecAccessPolicy(.whenUnlockedThisDeviceOnly)
+            let accessPolicy = AccessPolicy(.whenUnlockedThisDeviceOnly)
             XCTAssertEqual(accessPolicy.accessibility, String(kSecAttrAccessibleWhenUnlockedThisDeviceOnly))
             XCTAssertNil(try accessPolicy.accessControl)
         }
         do {
-            let accessPolicy = SecAccessPolicy(.afterFirstUnlock)
+            let accessPolicy = AccessPolicy(.afterFirstUnlock)
             XCTAssertEqual(accessPolicy.accessibility, String(kSecAttrAccessibleAfterFirstUnlock))
             XCTAssertNil(try accessPolicy.accessControl)
         }
         do {
-            let accessPolicy = SecAccessPolicy(.afterFirstUnlockThisDeviceOnly)
+            let accessPolicy = AccessPolicy(.afterFirstUnlockThisDeviceOnly)
             XCTAssertEqual(accessPolicy.accessibility, String(kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly))
             XCTAssertNil(try accessPolicy.accessControl)
         }
     }
 
     func testAccessControl() {
         do {
-            let accessPolicy = SecAccessPolicy(.afterFirstUnlock, options: .biometryAny)
+            let accessPolicy = AccessPolicy(.afterFirstUnlock, options: .biometryAny)
             XCTAssertNotNil(try accessPolicy.accessControl)
         }
     }