Cleanup Service + SequenceStore refactor (#241)

* Cleanup service * Clean sessionToPairingTopic * SequenceStore refactor * Rename KeyValueStore -> CodableStore
WalletConnect · May 31, 2022 · 6bec5ef · 6bec5ef
1 parent a32b649
commit 6bec5ef
Show file tree

Hide file tree

Showing 23 changed files with 151 additions and 61 deletions.
diff --git a/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift b/Sources/WalletConnectKMS/Crypto/KeyManagementService.swift
@@ -13,6 +13,7 @@ public protocol KeyManagementServiceProtocol {
     func deletePrivateKey(for publicKey: String)
     func deleteAgreementSecret(for topic: String)
     func deleteSymmetricKey(for topic: String)
+    func deleteAll() throws
     func performKeyAgreement(selfPublicKey: AgreementPublicKey, peerPublicKey hexRepresentation: String) throws -> AgreementKeys
 }
 
@@ -122,7 +123,9 @@ public class KeyManagementService: KeyManagementServiceProtocol {
         return try KeyManagementService.generateAgreementKey(from: privateKey, peerPublicKey: hexRepresentation)
     }
 
-
+    public func deleteAll() throws {
+        try keychain.deleteAll()
+    }
 
     static func generateAgreementKey(from privateKey: AgreementPrivateKey, peerPublicKey hexRepresentation: String) throws -> AgreementKeys {
         let peerPublicKey = try AgreementPublicKey(rawRepresentation: Data(hex: hexRepresentation))

diff --git a/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift b/Sources/WalletConnectKMS/Keychain/KeychainStorage.swift
@@ -4,6 +4,7 @@ protocol KeychainStorageProtocol {
     func add<T: GenericPasswordConvertible>(_ item: T, forKey key: String) throws
     func read<T: GenericPasswordConvertible>(key: String) throws -> T
     func delete(key: String) throws
+    func deleteAll() throws
 }
 
 final class KeychainStorage: KeychainStorageProtocol {

diff --git a/Sources/WalletConnectRelay/RelayClient.swift b/Sources/WalletConnectRelay/RelayClient.swift
@@ -44,7 +44,7 @@ public final class RelayClient {
         self.logger = logger
         self.dispatcher = dispatcher
 
-        self.jsonRpcSubscriptionsHistory = JsonRpcHistory<RelayJSONRPC.SubscriptionParams>(logger: logger, keyValueStore: KeyValueStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: Self.historyIdentifier))
+        self.jsonRpcSubscriptionsHistory = JsonRpcHistory<RelayJSONRPC.SubscriptionParams>(logger: logger, keyValueStore: CodableStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: Self.historyIdentifier))
         setUpBindings()
     }
 

diff --git a/Sources/WalletConnectSign/Engine/Common/PairingEngine.swift b/Sources/WalletConnectSign/Engine/Common/PairingEngine.swift
@@ -9,11 +9,11 @@ final class PairingEngine {
     var onProposeResponse: ((String, SessionProposal)->())?
     var onSessionRejected: ((Session.Proposal, SessionType.Reason)->())?
 
-    private let proposalPayloadsStore: KeyValueStore<WCRequestSubscriptionPayload>
+    private let proposalPayloadsStore: CodableStore<WCRequestSubscriptionPayload>
     private let networkingInteractor: NetworkInteracting
     private let kms: KeyManagementServiceProtocol
     private let pairingStore: WCPairingStorage
-    private let sessionToPairingTopic: KeyValueStore<String>
+    private let sessionToPairingTopic: CodableStore<String>
     private var metadata: AppMetadata
     private var publishers = [AnyCancellable]()
     private let logger: ConsoleLogging
@@ -22,11 +22,11 @@ final class PairingEngine {
     init(networkingInteractor: NetworkInteracting,
          kms: KeyManagementServiceProtocol,
          pairingStore: WCPairingStorage,
-         sessionToPairingTopic: KeyValueStore<String>,
+         sessionToPairingTopic: CodableStore<String>,
          metadata: AppMetadata,
          logger: ConsoleLogging,
          topicGenerator: @escaping () -> String = String.generateTopic,
-         proposalPayloadsStore: KeyValueStore<WCRequestSubscriptionPayload> = KeyValueStore<WCRequestSubscriptionPayload>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.proposals.rawValue)) {
+         proposalPayloadsStore: CodableStore<WCRequestSubscriptionPayload> = CodableStore<WCRequestSubscriptionPayload>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.proposals.rawValue)) {
         self.networkingInteractor = networkingInteractor
         self.kms = kms
         self.metadata = metadata

diff --git a/Sources/WalletConnectSign/Engine/Common/SessionEngine.swift b/Sources/WalletConnectSign/Engine/Common/SessionEngine.swift
@@ -16,7 +16,7 @@ final class SessionEngine {
 
     private let sessionStore: WCSessionStorage
     private let pairingStore: WCPairingStorage
-    private let sessionToPairingTopic: KeyValueStore<String>
+    private let sessionToPairingTopic: CodableStore<String>
     private let networkingInteractor: NetworkInteracting
     private let kms: KeyManagementServiceProtocol
     private var metadata: AppMetadata
@@ -28,7 +28,7 @@ final class SessionEngine {
          kms: KeyManagementServiceProtocol,
          pairingStore: WCPairingStorage,
          sessionStore: WCSessionStorage,
-         sessionToPairingTopic: KeyValueStore<String>,
+         sessionToPairingTopic: CodableStore<String>,
          metadata: AppMetadata,
          logger: ConsoleLogging,
          topicGenerator: @escaping () -> String = String.generateTopic) {

diff --git a/Sources/WalletConnectSign/JsonRpcHistory/JsonRpcHistory.swift b/Sources/WalletConnectSign/JsonRpcHistory/JsonRpcHistory.swift
@@ -11,10 +11,10 @@ protocol JsonRpcHistoryRecording {
 }
 //TODO -remove and use jsonrpc history only from utils
 class JsonRpcHistory: JsonRpcHistoryRecording {
-    let storage: KeyValueStore<JsonRpcRecord>
+    let storage: CodableStore<JsonRpcRecord>
     let logger: ConsoleLogging
 
-    init(logger: ConsoleLogging, keyValueStore: KeyValueStore<JsonRpcRecord>) {
+    init(logger: ConsoleLogging, keyValueStore: CodableStore<JsonRpcRecord>) {
         self.logger = logger
         self.storage = keyValueStore
     }

diff --git a/Sources/WalletConnectSign/Services/CelanupService.swift b/Sources/WalletConnectSign/Services/CelanupService.swift
@@ -0,0 +1,25 @@
+import Foundation
+import WalletConnectKMS
+import WalletConnectUtils
+
+final class CleanupService {
+
+    private let pairingStore: WCPairingStorage
+    private let sessionStore: WCSessionStorage
+    private let kms: KeyManagementServiceProtocol
+    private let sessionToPairingTopic: CodableStore<String>
+
+    init(pairingStore: WCPairingStorage, sessionStore: WCSessionStorage, kms: KeyManagementServiceProtocol, sessionToPairingTopic: CodableStore<String>) {
+        self.pairingStore = pairingStore
+        self.sessionStore = sessionStore
+        self.sessionToPairingTopic = sessionToPairingTopic
+        self.kms = kms
+    }
+
+    func cleanup() throws {
+        pairingStore.deleteAll()
+        sessionStore.deleteAll()
+        sessionToPairingTopic.deleteAll()
+        try kms.deleteAll()
+    }
+}
diff --git a/Sources/WalletConnectSign/Sign/SignClient.swift b/Sources/WalletConnectSign/Sign/SignClient.swift
@@ -33,6 +33,7 @@ public final class SignClient {
     private let networkingInteractor: NetworkInteracting
     private let kms: KeyManagementService
     private let history: JsonRpcHistory
+    private let cleanupService: CleanupService
 
     // MARK: - Initializers
 
@@ -52,20 +53,20 @@ public final class SignClient {
     init(metadata: AppMetadata, projectId: String, relayHost: String, logger: ConsoleLogging, kms: KeyManagementService, keyValueStorage: KeyValueStorage) {
         self.metadata = metadata
         self.logger = logger
-//        try? keychain.deleteAll() // Use for cleanup while lifecycles are not handled yet, but FIXME whenever
         self.kms = kms
         let relayClient = RelayClient(relayHost: relayHost, projectId: projectId, keyValueStorage: keyValueStorage, logger: logger)
         let serializer = Serializer(kms: kms)
-        self.history = JsonRpcHistory(logger: logger, keyValueStore: KeyValueStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.jsonRpcHistory.rawValue))
+        self.history = JsonRpcHistory(logger: logger, keyValueStore: CodableStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.jsonRpcHistory.rawValue))
         self.networkingInteractor = NetworkInteractor(relayClient: relayClient, serializer: serializer, logger: logger, jsonRpcHistory: history)
-        let pairingStore = PairingStorage(storage: SequenceStore<WCPairing>(storage: keyValueStorage, identifier: StorageDomainIdentifiers.pairings.rawValue))
-        let sessionStore = SessionStorage(storage: SequenceStore<WCSession>(storage: keyValueStorage, identifier: StorageDomainIdentifiers.sessions.rawValue))
-        let sessionToPairingTopic = KeyValueStore<String>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.sessionToPairingTopic.rawValue)
+        let pairingStore = PairingStorage(storage: SequenceStore<WCPairing>(store: .init(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.pairings.rawValue)))
+        let sessionStore = SessionStorage(storage: SequenceStore<WCSession>(store: .init(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.sessions.rawValue)))
+        let sessionToPairingTopic = CodableStore<String>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.sessionToPairingTopic.rawValue)
         self.pairingEngine = PairingEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore, sessionToPairingTopic: sessionToPairingTopic, metadata: metadata, logger: logger)
         self.sessionEngine = SessionEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore, sessionStore: sessionStore, sessionToPairingTopic: sessionToPairingTopic, metadata: metadata, logger: logger)
         self.nonControllerSessionStateMachine = NonControllerSessionStateMachine(networkingInteractor: networkingInteractor, kms: kms, sessionStore: sessionStore, logger: logger)
         self.controllerSessionStateMachine = ControllerSessionStateMachine(networkingInteractor: networkingInteractor, kms: kms, sessionStore: sessionStore, logger: logger)
         self.pairEngine = PairEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore)
+        self.cleanupService = CleanupService(pairingStore: pairingStore, sessionStore: sessionStore, kms: kms, sessionToPairingTopic: sessionToPairingTopic)
         setUpConnectionObserving(relayClient: relayClient)
         setUpEnginesCallbacks()
     }
@@ -87,16 +88,17 @@ public final class SignClient {
         self.logger = logger
         self.kms = kms
         let serializer = Serializer(kms: kms)
-        self.history = JsonRpcHistory(logger: logger, keyValueStore: KeyValueStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.jsonRpcHistory.rawValue))
+        self.history = JsonRpcHistory(logger: logger, keyValueStore: CodableStore<JsonRpcRecord>(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.jsonRpcHistory.rawValue))
         self.networkingInteractor = NetworkInteractor(relayClient: relayClient, serializer: serializer, logger: logger, jsonRpcHistory: history)
-        let pairingStore = PairingStorage(storage: SequenceStore<WCPairing>(storage: keyValueStorage, identifier: StorageDomainIdentifiers.pairings.rawValue))
-        let sessionStore = SessionStorage(storage: SequenceStore<WCSession>(storage: keyValueStorage, identifier: StorageDomainIdentifiers.sessions.rawValue))
-        let sessionToPairingTopic = KeyValueStore<String>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.sessionToPairingTopic.rawValue)
+        let pairingStore = PairingStorage(storage: SequenceStore<WCPairing>(store: .init(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.pairings.rawValue)))
+        let sessionStore = SessionStorage(storage: SequenceStore<WCSession>(store: .init(defaults: keyValueStorage, identifier: StorageDomainIdentifiers.sessions.rawValue)))
+        let sessionToPairingTopic = CodableStore<String>(defaults: RuntimeKeyValueStorage(), identifier: StorageDomainIdentifiers.sessionToPairingTopic.rawValue)
         self.pairingEngine = PairingEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore, sessionToPairingTopic: sessionToPairingTopic, metadata: metadata, logger: logger)
         self.sessionEngine = SessionEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore, sessionStore: sessionStore, sessionToPairingTopic: sessionToPairingTopic, metadata: metadata, logger: logger)
         self.nonControllerSessionStateMachine = NonControllerSessionStateMachine(networkingInteractor: networkingInteractor, kms: kms, sessionStore: sessionStore, logger: logger)
         self.controllerSessionStateMachine = ControllerSessionStateMachine(networkingInteractor: networkingInteractor, kms: kms, sessionStore: sessionStore, logger: logger)
         self.pairEngine = PairEngine(networkingInteractor: networkingInteractor, kms: kms, pairingStore: pairingStore)
+        self.cleanupService = CleanupService(pairingStore: pairingStore, sessionStore: sessionStore, kms: kms, sessionToPairingTopic: sessionToPairingTopic)
         setUpConnectionObserving(relayClient: relayClient)
         setUpEnginesCallbacks()
     }
@@ -327,4 +329,13 @@ public final class SignClient {
             sessionEngine.setSubscription(topic: sessionTopic)
         }
     }
+
+#if DEBUG
+    /// Delete all stored data sach as: pairings, sessions, keys
+    ///
+    /// - Note: Doesn't unsubscribe from topics
+    public func cleanup() throws {
+        try cleanupService.cleanup()
+    }
+#endif
 }
diff --git a/Sources/WalletConnectSign/Storage/PairingStorage.swift b/Sources/WalletConnectSign/Storage/PairingStorage.swift
@@ -5,6 +5,7 @@ protocol WCPairingStorage: AnyObject {
     func getPairing(forTopic topic: String) -> WCPairing?
     func getAll() -> [WCPairing]
     func delete(topic: String)
+    func deleteAll()
 }
 
 final class PairingStorage: WCPairingStorage {
@@ -39,4 +40,8 @@ final class PairingStorage: WCPairingStorage {
     func delete(topic: String) {
         storage.delete(topic: topic)
     }
+
+    func deleteAll() {
+        storage.deleteAll()
+    }
 }
diff --git a/Sources/WalletConnectSign/Storage/SequenceStore.swift b/Sources/WalletConnectSign/Storage/SequenceStore.swift
@@ -14,57 +14,47 @@ final class SequenceStore<T> where T: ExpirableSequence {
 
     var onSequenceExpiration: ((_ sequence: T) -> Void)?
 
-    private let storage: KeyValueStorage
+    private let store: CodableStore<T>
     private let dateInitializer: () -> Date
-    private let identifier: String
 
-    init(storage: KeyValueStorage, identifier: String, dateInitializer: @escaping () -> Date = Date.init) {
-        self.storage = storage
+    init(store: CodableStore<T>, dateInitializer: @escaping () -> Date = Date.init) {
+        self.store = store
         self.dateInitializer = dateInitializer
-        self.identifier = identifier
     }
 
     func hasSequence(forTopic topic: String) -> Bool {
         (try? getSequence(forTopic: topic)) != nil
     }
 
-    //  This force-unwrap is safe because Expirable Sequances are JSON Encodable
     func setSequence(_ sequence: T) {
-        let encoded = try! JSONEncoder().encode(sequence)
-        storage.set(encoded, forKey: getKey(for: sequence.topic))
+        store.set(sequence, forKey: sequence.topic)
     }
 
     func getSequence(forTopic topic: String) throws -> T? {
-        guard let data = storage.object(forKey: getKey(for: topic)) as? Data else { return nil }
-        let sequence = try JSONDecoder().decode(T.self, from: data)
-        return verifyExpiry(on: sequence)
+        guard let value = try store.get(key: topic) else { return nil }
+        return verifyExpiry(on: value)
     }
 
     func getAll() -> [T] {
-        return storage.dictionaryRepresentation().compactMap {
-            guard $0.key.hasPrefix(identifier) else {return nil}
-            if let data = $0.value as? Data, let sequence = try? JSONDecoder().decode(T.self, from: data) {
-                return verifyExpiry(on: sequence)
-            }
-            return nil
-        }
+        let values = store.getAll()
+        return values.compactMap { verifyExpiry(on: $0) }
     }
 
     func delete(topic: String) {
-        storage.removeObject(forKey: getKey(for: topic))
+        store.delete(forKey: topic)
+    }
+
+    func deleteAll() {
+        store.deleteAll()
     }
 
     private func verifyExpiry(on sequence: T) -> T? {
         let now = dateInitializer()
         if now >= sequence.expiryDate {
-            storage.removeObject(forKey: getKey(for: sequence.topic))
+            store.delete(forKey: sequence.topic)
             onSequenceExpiration?(sequence)
             return nil
         }
         return sequence
     }
-
-    private func getKey(for topic: String) -> String {
-        return "\(identifier).\(topic)"
-    }
 }
diff --git a/Sources/WalletConnectSign/Storage/SessionStorage.swift b/Sources/WalletConnectSign/Storage/SessionStorage.swift
@@ -5,6 +5,7 @@ protocol WCSessionStorage: AnyObject {
     func getSession(forTopic topic: String) -> WCSession?
     func getAll() -> [WCSession]
     func delete(topic: String)
+    func deleteAll()
 }
 
 final class SessionStorage: WCSessionStorage {
@@ -39,4 +40,8 @@ final class SessionStorage: WCSessionStorage {
     func delete(topic: String) {
         storage.delete(topic: topic)
     }
+
+    func deleteAll() {
+        storage.deleteAll()
+    }
 }
diff --git a/...es/WalletConnectUtils/KeyValueStore.swift → ...ces/WalletConnectUtils/CodableStore.swift b/...es/WalletConnectUtils/KeyValueStore.swift → ...ces/WalletConnectUtils/CodableStore.swift
@@ -1,7 +1,7 @@
 
 import Foundation
 
-public final class KeyValueStore<T> where T: Codable {
+public final class CodableStore<T> where T: Codable {
     private let defaults: KeyValueStorage
     private let prefix: String
 
@@ -10,8 +10,9 @@ public final class KeyValueStore<T> where T: Codable {
         self.prefix = identifier
     }
 
-    public func set(_ item: T, forKey key: String) throws {
-        let encoded = try JSONEncoder().encode(item)
+    public func set(_ item: T, forKey key: String) {
+        // This force-unwrap is safe because T are JSON Encodable
+        let encoded = try! JSONEncoder().encode(item)
         defaults.set(encoded, forKey: getContextPrefixedKey(for: key))
     }
 
@@ -22,8 +23,7 @@ public final class KeyValueStore<T> where T: Codable {
     }
 
     public func getAll() -> [T] {
-        return defaults.dictionaryRepresentation().compactMap {
-            guard $0.key.hasPrefix(prefix) else {return nil}
+        return dictionaryForIdentifier().compactMap {
             if let data = $0.value as? Data,
                let item = try? JSONDecoder().decode(T.self, from: data) {
                 return item
@@ -36,7 +36,17 @@ public final class KeyValueStore<T> where T: Codable {
         defaults.removeObject(forKey: getContextPrefixedKey(for: key))
     }
 
+    public func deleteAll() {
+        dictionaryForIdentifier()
+            .forEach { defaults.removeObject(forKey: $0.key) }
+    }
+
     private func getContextPrefixedKey(for key: String) -> String {
         return "\(prefix).\(key)"
     }
+
+    private func dictionaryForIdentifier() -> [String : Any] {
+        return defaults.dictionaryRepresentation()
+            .filter { $0.key.hasPrefix("\(prefix).") }
+    }
 }
diff --git a/Sources/WalletConnectUtils/JsonRpcHistory.swift b/Sources/WalletConnectUtils/JsonRpcHistory.swift
@@ -6,10 +6,10 @@ public class JsonRpcHistory<T> where T: Codable&Equatable {
         case jsonRpcDuplicateDetected
         case noJsonRpcRequestMatchingResponse
     }
-    private let storage: KeyValueStore<JsonRpcRecord>
+    private let storage: CodableStore<JsonRpcRecord>
     private let logger: ConsoleLogging
 
-    public init(logger: ConsoleLogging, keyValueStore: KeyValueStore<JsonRpcRecord>) {
+    public init(logger: ConsoleLogging, keyValueStore: CodableStore<JsonRpcRecord>) {
         self.logger = logger
         self.storage = keyValueStore
     }

diff --git a/Tests/TestingUtils/Mocks/KeyManagementServiceMock.swift b/Tests/TestingUtils/Mocks/KeyManagementServiceMock.swift
@@ -85,6 +85,12 @@ final class KeyManagementServiceMock: KeyManagementServiceProtocol {
     func deleteAgreementSecret(for topic: String) {
         agreementKeys[topic] = nil
     }
+
+    func deleteAll() throws {
+        privateKeys = [:]
+        symmetricKeys = [:]
+        agreementKeys = [:]
+    }
 }
 
 extension KeyManagementServiceMock {

diff --git a/Tests/TestingUtils/Mocks/KeychainStorageMock.swift b/Tests/TestingUtils/Mocks/KeychainStorageMock.swift
@@ -26,4 +26,8 @@ final class KeychainStorageMock: KeychainStorageProtocol {
         didCallDelete = true
         storage[key] = nil
     }
+
+    func deleteAll() throws {
+        storage = [:]
+    }
 }
diff --git a/Tests/WalletConnectSignTests/JsonRpcHistoryTests.swift b/Tests/WalletConnectSignTests/JsonRpcHistoryTests.swift
@@ -10,7 +10,7 @@ final class JsonRpcHistoryTests: XCTestCase {
     var sut: WalletConnectSign.JsonRpcHistory!
 
     override func setUp() {
-        sut = JsonRpcHistory(logger: ConsoleLoggerMock(), keyValueStore: KeyValueStore<WalletConnectSign.JsonRpcRecord>(defaults: RuntimeKeyValueStorage(), identifier: ""))
+        sut = JsonRpcHistory(logger: ConsoleLoggerMock(), keyValueStore: CodableStore<WalletConnectSign.JsonRpcRecord>(defaults: RuntimeKeyValueStorage(), identifier: ""))
     }
 
     override func tearDown() {