From 8ad09bbeb94b61466c9bb3e3ac55ea687313f3fc Mon Sep 17 00:00:00 2001
From: Artur Guseinov <arthur.guseinov@gmail.com>
Date: Tue, 31 Oct 2023 21:42:10 +0800
Subject: [PATCH] Update instead of replacing

---
 .../Keychain/KeychainError.swift              | 10 ++---
 .../Keychain/KeychainStorage.swift            | 44 ++++++-------------
 .../ClientAuth/ClientIdStorage.swift          |  3 ++
 3 files changed, 22 insertions(+), 35 deletions(-)

diff --git a/Sources/WalletConnectKMS/Keychain/KeychainError.swift b/Sources/WalletConnectKMS/Keychain/KeychainError.swift
index 78c16dcf7..fc3626e94 100644
--- a/Sources/WalletConnectKMS/Keychain/KeychainError.swift
+++ b/Sources/WalletConnectKMS/Keychain/KeychainError.swift
@@ -1,22 +1,22 @@
 import Foundation
 
 // TODO: Integrate with WalletConnectError
-struct KeychainError: Error, LocalizedError {
+public struct KeychainError: Error, LocalizedError {
 
-    let status: OSStatus
+    public let status: OSStatus
 
-    init(_ status: OSStatus) {
+    public init(_ status: OSStatus) {
         self.status = status
     }
 
-    var errorDescription: String? {
+    public var errorDescription: String? {
         return "OSStatus: \(status), message: \(status.message)"
     }
 }
 
 extension KeychainError: CustomStringConvertible {
 
-    var description: String {
+    public var description: String {
         status.message
     }
 }
diff --git a/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift b/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift
index 4573172d3..f52f4c828 100644
--- a/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift
+++ b/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift
@@ -55,8 +55,7 @@ public final class KeychainStorage: KeychainStorageProtocol {
         case errSecSuccess:
             return item as? Data
         case errSecItemNotFound:
-            // TODO: Replace with nil once migration period ends
-            return try tryMigrateAttrAccessible(key: key)
+            return tryMigrateAttrAccessible(key: key) // TODO: Replace with nil once migration period ends
         default:
             throw KeychainError(status)
         }
@@ -109,38 +108,23 @@ public final class KeychainStorage: KeychainStorageProtocol {
         ]
     }
 
-    private func tryMigrateAttrAccessible(key: String) throws -> Data? {
-        var query = buildBaseServiceQuery(for: key)
-        query[kSecReturnData] = true
-        query[kSecAttrAccessible] = kSecAttrAccessibleWhenUnlockedThisDeviceOnly
-
-        var item: CFTypeRef?
-        let status = secItem.copyMatching(query as CFDictionary, &item)
+    private func tryMigrateAttrAccessible(key: String) -> Data? {
+        var updateQuery = buildBaseServiceQuery(for: key)
+        updateQuery[kSecAttrAccessible] = kSecAttrAccessibleWhenUnlockedThisDeviceOnly
 
-        switch status {
-        case errSecSuccess: // Migration needed
-            guard let data = item as? Data else { return nil }
-            
-            // Fetching old value
-            var deleteQuery = buildBaseServiceQuery(for: key)
-            deleteQuery[kSecAttrAccessible] = kSecAttrAccessibleWhenUnlockedThisDeviceOnly
+        let attributes = [kSecAttrAccessible: kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly]
+        let status = secItem.update(updateQuery as CFDictionary, attributes as CFDictionary)
 
-            // Deleting old value
-            let status = secItem.delete(deleteQuery as CFDictionary)
+        guard status == errSecSuccess else {
+            return nil
+        }
 
-            guard status == errSecSuccess || status == errSecItemNotFound else {
-                throw KeychainError(status)
-            }
+        var readQuery = buildBaseServiceQuery(for: key)
+        readQuery[kSecReturnData] = true
 
-            // Replacing with new value
-            try add(data: data, forKey: key)
+        var item: CFTypeRef?
+        _ = secItem.copyMatching(readQuery as CFDictionary, &item)
 
-            // Continue `readData` execution
-            return item as? Data
-        case errSecItemNotFound:
-            return nil
-        default:
-            throw KeychainError(status)
-        }
+        return item as? Data
     }
 }
diff --git a/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift b/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
index 272451b96..05125a486 100644
--- a/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
+++ b/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
@@ -26,6 +26,9 @@ public struct ClientIdStorage: ClientIdStoring {
             let publicPart = try getPublicPart()
             return try getPrivatePart(for: publicPart)
         } catch {
+            guard let error = error as? KeychainError, error.status == errSecItemNotFound else {
+                throw error
+            }
             let privateKey = SigningPrivateKey()
             try setPrivatePart(privateKey)
             setPublicPart(privateKey.publicKey)