Cleanup Service + SequenceStore refactor

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))

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 {

Sources/WalletConnectSign/Services/CelanupService.swift

@@ -0,0 +1,21 @@
+import Foundation
+import WalletConnectKMS
+
+final class CleanupService {
+
+    private let pairingStore: WCPairingStorage
+    private let sessionStore: WCSessionStorage
+    private let kms: KeyManagementServiceProtocol
+
+    init(pairingStore: WCPairingStorage, sessionStore: WCSessionStorage, kms: KeyManagementServiceProtocol) {
+        self.pairingStore = pairingStore
+        self.sessionStore = sessionStore
+        self.kms = kms
+    }
+
+    func cleanup() throws {
+        pairingStore.deleteAll()
+        sessionStore.deleteAll()
+        try kms.deleteAll()
+    }
+}

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,7 +53,6 @@ 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)
@@ -66,6 +66,7 @@ public final class SignClient {
         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)
         setUpConnectionObserving(relayClient: relayClient)
         setUpEnginesCallbacks()
     }
@@ -97,6 +98,7 @@ public final class SignClient {
         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)
         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
 }

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()
+    }
 }

Sources/WalletConnectSign/Storage/SequenceStore.swift

@@ -41,8 +41,7 @@ final class SequenceStore<T> where T: ExpirableSequence {
     }
 
     func getAll() -> [T] {
-        return storage.dictionaryRepresentation().compactMap {
-            guard $0.key.hasPrefix(identifier) else {return nil}
+        return dictionaryForIdentifier().compactMap {
             if let data = $0.value as? Data, let sequence = try? JSONDecoder().decode(T.self, from: data) {
                 return verifyExpiry(on: sequence)
             }
@@ -54,6 +53,11 @@ final class SequenceStore<T> where T: ExpirableSequence {
         storage.removeObject(forKey: getKey(for: topic))
     }
 
+    func deleteAll() {
+        dictionaryForIdentifier()
+            .forEach { storage.removeObject(forKey: $0.key) }
+    }
+
     private func verifyExpiry(on sequence: T) -> T? {
         let now = dateInitializer()
         if now >= sequence.expiryDate {
@@ -67,4 +71,9 @@ final class SequenceStore<T> where T: ExpirableSequence {
     private func getKey(for topic: String) -> String {
         return "\(identifier).\(topic)"
     }
+
+    private func dictionaryForIdentifier() -> [String : Any] {
+        return storage.dictionaryRepresentation()
+            .filter { $0.key.hasPrefix("\(identifier).") }
+    }
 }

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()
+    }
 }

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 {

Tests/TestingUtils/Mocks/KeychainStorageMock.swift

@@ -26,4 +26,8 @@ final class KeychainStorageMock: KeychainStorageProtocol {
         didCallDelete = true
         storage[key] = nil
     }
+
+    func deleteAll() throws {
+        storage = [:]
+    }
 }

Tests/WalletConnectSignTests/Mocks/WCPairingStorageMock.swift

@@ -1,7 +1,7 @@
 @testable import WalletConnectSign
 
 final class WCPairingStorageMock: WCPairingStorage {
-    
+
     var onPairingExpiration: ((WCPairing) -> Void)?
 
     private(set) var pairings: [String: WCPairing] = [:]
@@ -25,4 +25,8 @@ final class WCPairingStorageMock: WCPairingStorage {
     func delete(topic: String) {
         pairings[topic] = nil
     }
+
+    func deleteAll() {
+        pairings = [:]
+    }
 }

Tests/WalletConnectSignTests/Mocks/WCSessionStorageMock.swift

@@ -25,5 +25,9 @@ final class WCSessionStorageMock: WCSessionStorage {
     func delete(topic: String) {
         sessions[topic] = nil
     }
+
+    func deleteAll() {
+        sessions = [:]
+    }
 }
 

Tests/WalletConnectSignTests/SequenceStoreTests.swift

@@ -21,7 +21,7 @@ final class SequenceStoreTests: XCTestCase {
     override func setUp() {
         timeTraveler = TimeTraveler()
         storageFake = RuntimeKeyValueStorage()
-        sut = SequenceStore<ExpirableSequenceStub>(storage: storageFake, identifier: "", dateInitializer: timeTraveler.generateDate)
+        sut = makeStore("test")
         sut.onSequenceExpiration = { _ in
             XCTFail("Unexpected expiration call")
         }
@@ -33,6 +33,14 @@ final class SequenceStoreTests: XCTestCase {
         sut = nil
     }
 
+    private func makeStore(_ identifier: String) -> SequenceStore<ExpirableSequenceStub> {
+        return SequenceStore<ExpirableSequenceStub>(
+            storage: storageFake,
+            identifier: identifier,
+            dateInitializer: timeTraveler.generateDate
+        )
+    }
+
     private func stubSequence(expiry: TimeInterval? = nil) -> ExpirableSequenceStub {
         ExpirableSequenceStub(
             topic: String.generateTopic(),
@@ -73,6 +81,22 @@ final class SequenceStoreTests: XCTestCase {
         XCTAssertNil(retrieved)
     }
 
+    func testDeleteAll() {
+        let sequence = stubSequence()
+        sut.setSequence(sequence)
+
+        let sut2 = makeStore("test2")
+        sut2.setSequence(sequence)
+
+        XCTAssertFalse(sut.getAll().isEmpty)
+        XCTAssertFalse(sut2.getAll().isEmpty)
+
+        sut.deleteAll()
+
+        XCTAssertTrue(sut.getAll().isEmpty)
+        XCTAssertFalse(sut2.getAll().isEmpty)
+    }
+
     // MARK: - Expiration Tests
 
     func testHasSequenceExpiration() {