Merge pull request #45 from WalletConnect/#35-update-checks

#35 Update permission checks
WalletConnect · Jan 24, 2022 · 58f675b · 58f675b
2 parents 496f6a8 + e5ae263
commit 58f675b
Show file tree

Hide file tree

Showing 8 changed files with 197 additions and 41 deletions.
diff --git a/Sources/WalletConnect/Engine/SessionEngine.swift b/Sources/WalletConnect/Engine/SessionEngine.swift
@@ -198,21 +198,19 @@ final class SessionEngine {
         }
     }
 
-    func update(topic: String, accounts: Set<String>) {
-        guard var session = try? sequencesStore.getSequence(forTopic: topic) else {
-            logger.debug("Could not find session for topic \(topic)")
-            return
-        }
-        session.update(accounts)
-        relayer.request(.wcSessionUpdate(SessionType.UpdateParams(state: SessionState(accounts: accounts))), onTopic: topic) { [unowned self] result in
-            switch result {
-            case .success(_):
-                sequencesStore.setSequence(session)
-                onSessionUpdate?(topic, accounts)
-            case .failure(_):
-                break
+    func update(topic: String, accounts: Set<String>) throws {
+        var session = try sequencesStore.getSequence(forTopic: topic)
+        for account in accounts {
+            if !String.conformsToCAIP10(account) {
+                throw WalletConnectError.internal(.notApproved) // TODO: Use a suitable error cases
             }
         }
+        if !isController || session.settled?.status != .acknowledged {
+            throw WalletConnectError.unauthrorized(.unauthorizedUpdateRequest)
+        }
+        session.update(accounts)
+        sequencesStore.setSequence(session)
+        relayer.request(.wcSessionUpdate(SessionType.UpdateParams(accounts: accounts)), onTopic: topic)
     }
 
     func upgrade(topic: String, permissions: Session.Permissions) {
@@ -270,7 +268,7 @@ final class SessionEngine {
             case .sessionReject(let rejectParams):
                 handleSessionReject(rejectParams, topic: topic)
             case .sessionUpdate(let updateParams):
-                handleSessionUpdate(topic: topic, updateParams: updateParams, requestId: requestId)
+                handleSessionUpdate(payload: subscriptionPayload, updateParams: updateParams)
             case .sessionUpgrade(let upgradeParams):
                 handleSessionUpgrade(topic: topic, upgradeParams: upgradeParams, requestId: requestId)
             case .sessionDelete(let deleteParams):
@@ -320,24 +318,34 @@ final class SessionEngine {
         }
     }
 
-    private func handleSessionUpdate(topic: String, updateParams: SessionType.UpdateParams, requestId: Int64) {
-        guard var session = try? sequencesStore.getSequence(forTopic: topic) else {
-            logger.debug("Could not find session for topic \(topic)")
+    // TODO: Use standard protocol reason codes when responding
+    private func handleSessionUpdate(payload: WCRequestSubscriptionPayload, updateParams: SessionType.UpdateParams) {
+        let topic = payload.topic
+        let requestId = payload.wcRequest.id
+
+        for account in updateParams.state.accounts {
+            if !String.conformsToCAIP10(account) {
+                relayer.respondError(for: payload, reason: Reason(code: 0, message: "invalid accounts"))
+                return
+            }
+        }
+        guard var session = try? sequencesStore.getSequence(forTopic: topic), session.isSettled else {
+            relayer.respondError(for: payload, reason: Reason(code: 0, message: "session not found"))
             return
         }
-        guard session.peerIsController else {
-            let error = WalletConnectError.unauthrorized(.unauthorizedUpdateRequest)
-            logger.error(error)
-            respond(error: error, requestId: requestId, topic: topic)
+        guard !isController, session.peerIsController else {
+            relayer.respondError(for: payload, reason: Reason(code: 0, message: "unauthorized update"))
             return
         }
+
+        session.settled?.state = updateParams.state
+        sequencesStore.setSequence(session)
+
         let response = JSONRPCResponse<AnyCodable>(id: requestId, result: AnyCodable(true))
         relayer.respond(topic: topic, response: JsonRpcResponseTypes.response(response)) { [unowned self] error in
             if let error = error {
                 logger.error(error)
             } else {
-                session.settled?.state = updateParams.state
-                sequencesStore.setSequence(session)
                 onSessionUpdate?(topic, updateParams.state.accounts)
             }
         }
@@ -509,6 +517,8 @@ final class SessionEngine {
             handleProposeResponse(topic: response.topic, proposeParams: proposeParams, result: response.result)
         case .sessionApprove(let approveParams):
             handleApproveResponse(topic: response.topic, result: response.result)
+        case .sessionUpdate:
+            handleUpdateResponse(topic: response.topic, result: response.result)
         default:
             break
         }
@@ -556,4 +566,16 @@ final class SessionEngine {
             crypto.deletePrivateKey(for: pendingSession.publicKey)
         }
     }
+
+    private func handleUpdateResponse(topic: String, result: Result<JSONRPCResponse<AnyCodable>, Error>) {
+        guard let session = try? sequencesStore.getSequence(forTopic: topic), let accounts = session.settled?.state.accounts else {
+            return
+        }
+        switch result {
+        case .success:
+            onSessionUpdate?(topic, accounts)
+        case .failure:
+            logger.error("Peer failed to update state.")
+        }
+    }
 }
diff --git a/Sources/WalletConnect/Relay/WalletConnectRelay.swift b/Sources/WalletConnect/Relay/WalletConnectRelay.swift
@@ -18,6 +18,7 @@ protocol WalletConnectRelaying: AnyObject {
     func request(_ wcMethod: WCMethod, onTopic topic: String, completion: ((Result<JSONRPCResponse<AnyCodable>, JSONRPCErrorResponse>)->())?)
     func request(topic: String, payload: WCRequest, completion: ((Result<JSONRPCResponse<AnyCodable>, JSONRPCErrorResponse>)->())?) 
     func respond(topic: String, response: JsonRpcResponseTypes, completion: @escaping ((Error?)->()))
+    func respondError(for payload: WCRequestSubscriptionPayload, reason: Reason)
     func subscribe(topic: String)
     func unsubscribe(topic: String)
 }
@@ -116,6 +117,11 @@ class WalletConnectRelay: WalletConnectRelaying {
         }
     }
 
+    func respondError(for payload: WCRequestSubscriptionPayload, reason: Reason) {
+        let response = JSONRPCErrorResponse(id: payload.wcRequest.id, error: JSONRPCErrorResponse.Error(code: reason.code, message: reason.message))
+        respond(topic: payload.topic, response: JsonRpcResponseTypes.error(response)) { _ in } // TODO: Move error handling to relayer package
+    }
+
     func subscribe(topic: String)  {
         networkRelayer.subscribe(topic: topic) { [weak self] error in
             if let error = error {

diff --git a/Sources/WalletConnect/Storage/SessionStorage.swift b/Sources/WalletConnect/Storage/SessionStorage.swift
@@ -2,7 +2,7 @@ protocol SessionSequenceStorage: AnyObject {
     var onSequenceExpiration: ((_ topic: String, _ pubKey: String) -> Void)? { get set }
     func hasSequence(forTopic topic: String) -> Bool
     func setSequence(_ sequence: SessionSequence)
-    func getSequence(forTopic topic: String) throws -> SessionSequence?
+    func getSequence(forTopic topic: String) throws -> SessionSequence
     func getAll() -> [SessionSequence]
     func delete(topic: String)
 }
@@ -28,8 +28,11 @@ final class SessionStorage: SessionSequenceStorage {
         storage.setSequence(sequence)
     }
 
-    func getSequence(forTopic topic: String) throws -> SessionSequence? {
-        try storage.getSequence(forTopic: topic)
+    func getSequence(forTopic topic: String) throws -> SessionSequence {
+        guard let sequence = try storage.getSequence(forTopic: topic) else {
+            throw WalletConnectError.internal(.noSequenceForTopic)
+        }
+        return sequence
     }
 
     func getAll() -> [SessionSequence] {

diff --git a/Sources/WalletConnect/Types/Session/SessionType.swift b/Sources/WalletConnect/Types/Session/SessionType.swift
@@ -19,6 +19,10 @@ internal enum SessionType {
 
     struct UpdateParams: Codable, Equatable {
         let state: SessionState
+
+        init(accounts: Set<String>) {
+            self.state = SessionState(accounts: accounts)
+        }
     }
 
     struct UpgradeParams: Codable, Equatable {

diff --git a/Sources/WalletConnect/WalletConnectClient.swift b/Sources/WalletConnect/WalletConnectClient.swift
@@ -137,7 +137,11 @@ public final class WalletConnectClient {
     ///   - topic: Topic of the session that is intended to be updated.
     ///   - accounts: Set of accounts that will be allowed to be used by the session after the update.
     public func update(topic: String, accounts: Set<String>) {
-        sessionEngine.update(topic: topic, accounts: accounts)
+        do {
+            try sessionEngine.update(topic: topic, accounts: accounts)
+        } catch {
+            print("Error on session update call: \(error)")
+        }
     }
 
     /// For the responder to upgrade session permissions

diff --git a/Tests/WalletConnectTests/Mocks/MockedRelay.swift b/Tests/WalletConnectTests/Mocks/MockedRelay.swift
@@ -20,24 +20,31 @@ class MockedWCRelay: WalletConnectRelaying {
     var wcRequestPublisher: AnyPublisher<WCRequestSubscriptionPayload, Never> {
         wcRequestPublisherSubject.eraseToAnyPublisher()
     }
+    var didCallRequest = false
     var didCallSubscribe = false
     var didCallUnsubscribe = false
+    var didRespondError = false
     var error: Error? = nil
 
     private(set) var requests: [(topic: String, request: WCRequest)] = []
-
+    
     func request(_ wcMethod: WCMethod, onTopic topic: String, completion: ((Result<JSONRPCResponse<AnyCodable>, JSONRPCErrorResponse>) -> ())?) {
         request(topic: topic, payload: wcMethod.asRequest(), completion: completion)
     }
 
     func request(topic: String, payload: WCRequest, completion: ((Result<JSONRPCResponse<AnyCodable>, JSONRPCErrorResponse>) -> ())?) {
+        didCallRequest = true
         requests.append((topic, payload))
     }
 
     func respond(topic: String, response: JsonRpcResponseTypes, completion: @escaping ((Error?) -> ())) {
         completion(error)
     }
 
+    func respondError(for payload: WCRequestSubscriptionPayload, reason: Reason) {
+        didRespondError = true
+    }
+
     func subscribe(topic: String) {
         didCallSubscribe = true
     }

diff --git a/Tests/WalletConnectTests/Mocks/SessionSequenceStorageMock.swift b/Tests/WalletConnectTests/Mocks/SessionSequenceStorageMock.swift
@@ -1,4 +1,5 @@
 @testable import WalletConnect
+import Foundation
 
 final class SessionSequenceStorageMock: SessionSequenceStorage {
 
@@ -14,8 +15,11 @@ final class SessionSequenceStorageMock: SessionSequenceStorage {
         sessions[sequence.topic] = sequence
     }
 
-    func getSequence(forTopic topic: String) throws -> SessionSequence? {
-        sessions[topic]
+    func getSequence(forTopic topic: String) throws -> SessionSequence {
+        guard let session = sessions[topic] else {
+            throw NSError.mock()
+        }
+        return session
     }
 
     func getAll() -> [SessionSequence] {