From 9dd5dd0356507e8fd1c769dc0b492481c1d08208 Mon Sep 17 00:00:00 2001
From: Artur Guseinov <arthur.guseinov@gmail.com>
Date: Wed, 1 Nov 2023 00:52:13 +0800
Subject: [PATCH] Keychain error refactor

---
 .../Crypto/KeyManagementService.swift         |  2 +-
 .../Keychain/KeychainError.swift              | 23 +++++++++++++------
 .../ClientAuth/ClientIdStorage.swift          | 11 +++++----
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift b/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift
index 53a271b14..975a39251 100644
--- a/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift
+++ b/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift
@@ -93,7 +93,7 @@ public class KeyManagementService: KeyManagementServiceProtocol {
     public func getPrivateKey(for publicKey: AgreementPublicKey) throws -> AgreementPrivateKey? {
         do {
             return try keychain.read(key: publicKey.hexRepresentation) as AgreementPrivateKey
-        } catch let error where (error as? KeychainError)?.status == errSecItemNotFound {
+        } catch KeychainError.itemNotFound {
             return nil
         } catch {
             throw error
diff --git a/Sources/WalletConnectKMS/Keychain/KeychainError.swift b/Sources/WalletConnectKMS/Keychain/KeychainError.swift
index fc3626e94..9b947819f 100644
--- a/Sources/WalletConnectKMS/Keychain/KeychainError.swift
+++ b/Sources/WalletConnectKMS/Keychain/KeychainError.swift
@@ -1,23 +1,32 @@
 import Foundation
 
-// TODO: Integrate with WalletConnectError
-public struct KeychainError: Error, LocalizedError {
-
-    public let status: OSStatus
+public enum KeychainError: Error, LocalizedError {
+    case itemNotFound
+    case other(OSStatus)
 
     public init(_ status: OSStatus) {
-        self.status = status
+        switch status {
+        case errSecItemNotFound:
+            self = .itemNotFound
+        default:
+            self = .other(status)
+        }
     }
 
     public var errorDescription: String? {
-        return "OSStatus: \(status), message: \(status.message)"
+        switch self {
+        case .itemNotFound:
+            return "Keychain item not found"
+        case .other(let status):
+            return "OSStatus: \(status), message: \(status.message)"
+        }
     }
 }
 
 extension KeychainError: CustomStringConvertible {
 
     public var description: String {
-        status.message
+        return errorDescription ?? ""
     }
 }
 
diff --git a/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift b/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
index 05125a486..e07333198 100644
--- a/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
+++ b/Sources/WalletConnectRelay/ClientAuth/ClientIdStorage.swift
@@ -25,14 +25,13 @@ public struct ClientIdStorage: ClientIdStoring {
         do {
             let publicPart = try getPublicPart()
             return try getPrivatePart(for: publicPart)
-        } catch {
-            guard let error = error as? KeychainError, error.status == errSecItemNotFound else {
-                throw error
-            }
+        } catch Errors.privatePartNotFound, Errors.publicPartNotFound {
             let privateKey = SigningPrivateKey()
             try setPrivatePart(privateKey)
             setPublicPart(privateKey.publicKey)
             return privateKey
+        } catch {
+            throw error
         }
     }
 
@@ -79,8 +78,10 @@ private extension ClientIdStorage {
     func getPrivatePart(for publicPart: SigningPublicKey) throws -> SigningPrivateKey {
         do {
             return try keychain.read(key: publicPart.storageId)
-        } catch {
+        } catch KeychainError.itemNotFound {
             throw Errors.privatePartNotFound
+        } catch {
+            throw error
         }
     }