[Sign] Feature: #252 Update session namespaces checks against require…

…d namespaces (#276) * Set required namespaces in session creation * Add basic checks to required namespaces * Changed updateNamespaces to throw errors * Add basic tests for valid & invalid update * Add failure tests, implemented checks on namespace extensions * Test small improvements Co-authored-by: André Vants <[email protected]>
WalletConnect · Jun 21, 2022 · 0aff90b · 0aff90b
1 parent 1773af7
commit 0aff90b
Show file tree

Hide file tree

Showing 6 changed files with 216 additions and 27 deletions.
diff --git a/Sources/WalletConnectSign/Engine/Controller/ControllerSessionStateMachine.swift b/Sources/WalletConnectSign/Engine/Controller/ControllerSessionStateMachine.swift
@@ -32,7 +32,7 @@ final class ControllerSessionStateMachine {
         try validateControlledAcknowledged(session)
         try Namespace.validate(namespaces)
         logger.debug("Controller will update methods")
-        session.updateNamespaces(namespaces)
+        try session.updateNamespaces(namespaces)
         sessionStore.setSession(session)
         try await networkingInteractor.request(.wcSessionUpdate(SessionType.UpdateParams(namespaces: namespaces)), onTopic: topic)
     }

diff --git a/Sources/WalletConnectSign/Engine/NonController/NonControllerSessionStateMachine.swift b/Sources/WalletConnectSign/Engine/NonController/NonControllerSessionStateMachine.swift
@@ -70,7 +70,11 @@ final class NonControllerSessionStateMachine {
         guard session.peerIsController else {
             throw Errors.respondError(payload: payload, reason: .unauthorizedUpdateNamespacesRequest)
         }
-        session.updateNamespaces(updateParams.namespaces)
+        do {
+            try session.updateNamespaces(updateParams.namespaces)
+        } catch {
+            throw Errors.respondError(payload: payload, reason: .invalidUpdateNamespaceRequest)
+        }
         sessionStore.setSession(session)
         networkingInteractor.respondSuccess(for: payload)
         onNamespacesUpdate?(session.topic, updateParams.namespaces)

diff --git a/Sources/WalletConnectSign/Namespace.swift b/Sources/WalletConnectSign/Namespace.swift
@@ -44,6 +44,12 @@ public struct SessionNamespace: Equatable, Codable {
             self.methods = methods
             self.events = events
         }
+
+        func isSuperset(of other: Extension) -> Bool {
+            self.accounts.isSuperset(of: other.accounts) &&
+            self.methods.isSuperset(of: other.methods) &&
+            self.events.isSuperset(of: other.events)
+        }
     }
 
     public init(accounts: Set<Account>, methods: Set<String>, events: Set<String>, extensions: [SessionNamespace.Extension]? = nil) {

diff --git a/Sources/WalletConnectSign/Types/Session/WCSession.swift b/Sources/WalletConnectSign/Types/Session/WCSession.swift
@@ -3,9 +3,12 @@ import WalletConnectKMS
 import WalletConnectUtils
 
 struct WCSession: ExpirableSequence {
+
     enum Error: Swift.Error {
         case controllerNotSet
+        case unsatisfiedUpdateNamespaceRequirement
     }
+
     let topic: String
     let relay: RelayProtocolOptions
     let selfParticipant: Participant
@@ -14,6 +17,7 @@ struct WCSession: ExpirableSequence {
     var acknowledged: Bool
     let controller: AgreementPeer
     private(set) var namespaces: [String: SessionNamespace]
+    private(set) var requiredNamespaces: [String: SessionNamespace]
 
     static var defaultTimeToLive: Int64 {
         Int64(7*Time.day)
@@ -34,6 +38,7 @@ struct WCSession: ExpirableSequence {
         self.selfParticipant = selfParticipant
         self.peerParticipant = peerParticipant
         self.namespaces = settleParams.namespaces
+        self.requiredNamespaces = settleParams.namespaces
         self.acknowledged = acknowledged
         self.expiryDate = Date(timeIntervalSince1970: TimeInterval(settleParams.expiry))
     }
@@ -46,6 +51,7 @@ struct WCSession: ExpirableSequence {
         self.selfParticipant = selfParticipant
         self.peerParticipant = peerParticipant
         self.namespaces = namespaces
+        self.requiredNamespaces = namespaces
         self.acknowledged = acknowledged
         self.expiryDate = Date(timeIntervalSince1970: TimeInterval(expiry))
     }
@@ -107,7 +113,27 @@ struct WCSession: ExpirableSequence {
         return false
     }
 
-    mutating func updateNamespaces(_ namespaces: [String: SessionNamespace]) {
+    mutating func updateNamespaces(_ namespaces: [String: SessionNamespace]) throws {
+        for item in requiredNamespaces {
+            guard
+                let compliantNamespace = namespaces[item.key],
+                compliantNamespace.accounts.isSuperset(of: item.value.accounts),
+                compliantNamespace.methods.isSuperset(of: item.value.methods),
+                compliantNamespace.events.isSuperset(of: item.value.events)
+            else {
+                throw Error.unsatisfiedUpdateNamespaceRequirement
+            }
+            if let extensions = item.value.extensions {
+                guard let compliantExtensions = compliantNamespace.extensions else {
+                    throw Error.unsatisfiedUpdateNamespaceRequirement
+                }
+                for existingExtension in extensions {
+                    guard compliantExtensions.contains(where: { $0.isSuperset(of: existingExtension) }) else {
+                        throw Error.unsatisfiedUpdateNamespaceRequirement
+                    }
+                }
+            }
+        }
         self.namespaces = namespaces
     }
 

diff --git a/Tests/WalletConnectSignTests/Stub/Session+Stub.swift b/Tests/WalletConnectSignTests/Stub/Session+Stub.swift
@@ -7,22 +7,24 @@ extension WCSession {
         isSelfController: Bool = false,
         expiryDate: Date = Date.distantFuture,
         selfPrivateKey: AgreementPrivateKey = AgreementPrivateKey(),
-        acknowledged: Bool = true) -> WCSession {
-            let peerKey = selfPrivateKey.publicKey.hexRepresentation
-            let selfKey = AgreementPrivateKey().publicKey.hexRepresentation
-            let controllerKey = isSelfController ? selfKey : peerKey
-            return WCSession(
-                topic: String.generateTopic(),
-                relay: RelayProtocolOptions.stub(),
-                controller: AgreementPeer(publicKey: controllerKey),
-                selfParticipant: Participant.stub(publicKey: selfKey),
-                peerParticipant: Participant.stub(publicKey: peerKey),
-                namespaces: [:],
-                events: [],
-                accounts: Account.stubSet(),
-                acknowledged: acknowledged,
-                expiry: Int64(expiryDate.timeIntervalSince1970))
-        }
+        namespaces: [String: SessionNamespace] = [:],
+        acknowledged: Bool = true
+    ) -> WCSession {
+        let peerKey = selfPrivateKey.publicKey.hexRepresentation
+        let selfKey = AgreementPrivateKey().publicKey.hexRepresentation
+        let controllerKey = isSelfController ? selfKey : peerKey
+        return WCSession(
+            topic: String.generateTopic(),
+            relay: RelayProtocolOptions.stub(),
+            controller: AgreementPeer(publicKey: controllerKey),
+            selfParticipant: Participant.stub(publicKey: selfKey),
+            peerParticipant: Participant.stub(publicKey: peerKey),
+            namespaces: namespaces,
+            events: [],
+            accounts: Account.stubSet(),
+            acknowledged: acknowledged,
+            expiry: Int64(expiryDate.timeIntervalSince1970))
+    }
 }
 
 extension Account {

diff --git a/Tests/WalletConnectSignTests/WCSessionTests.swift b/Tests/WalletConnectSignTests/WCSessionTests.swift
@@ -7,6 +7,8 @@ final class WCSessionTests: XCTestCase {
     let polyAccount = Account("eip155:137:0xab16a96d359ec26a11e2c2b3d8f8b8942d5bfcdb")!
     let cosmosAccount = Account("cosmos:cosmoshub-4:cosmos1t2uflqwqe0fsj0shcfkrvpukewcw40yjj6hdc0")!
 
+    // MARK: Namespace Permission Tests
+
     func testHasPermissionForMethod() {
         let namespace = [
             "eip155": SessionNamespace(
@@ -16,7 +18,7 @@ final class WCSessionTests: XCTestCase {
                 extensions: nil)
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertTrue(session.hasPermission(forMethod: "method", onChain: ethAccount.blockchain))
     }
 
@@ -33,7 +35,7 @@ final class WCSessionTests: XCTestCase {
                         events: [])])
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertTrue(session.hasPermission(forMethod: "method", onChain: ethAccount.blockchain))
     }
 
@@ -51,7 +53,7 @@ final class WCSessionTests: XCTestCase {
                 extensions: nil)
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertFalse(session.hasPermission(forMethod: "method", onChain: ethAccount.blockchain))
     }
 
@@ -68,7 +70,7 @@ final class WCSessionTests: XCTestCase {
                         events: [])])
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertFalse(session.hasPermission(forMethod: "method", onChain: ethAccount.blockchain))
     }
 
@@ -81,7 +83,7 @@ final class WCSessionTests: XCTestCase {
                 extensions: nil)
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertTrue(session.hasPermission(forEvent: "event", onChain: ethAccount.blockchain))
     }
 
@@ -98,7 +100,7 @@ final class WCSessionTests: XCTestCase {
                         events: ["event"])])
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertTrue(session.hasPermission(forEvent: "event", onChain: ethAccount.blockchain))
     }
 
@@ -116,7 +118,7 @@ final class WCSessionTests: XCTestCase {
                 extensions: nil)
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertFalse(session.hasPermission(forEvent: "event", onChain: ethAccount.blockchain))
     }
 
@@ -133,7 +135,156 @@ final class WCSessionTests: XCTestCase {
                         events: ["event"])])
         ]
         var session = WCSession.stub()
-        session.updateNamespaces(namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
         XCTAssertFalse(session.hasPermission(forEvent: "event", onChain: ethAccount.blockchain))
     }
+
+    // MARK: Namespace Update Tests
+
+    private func stubRequiredNamespaces() -> [String: SessionNamespace] {
+        return [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: nil)]
+    }
+
+    private func stubRequiredNamespacesWithExtension() -> [String: SessionNamespace] {
+        return [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount, polyAccount],
+                        methods: ["method-2", "newMethod-2"],
+                        events: ["event-2", "newEvent-2"])])]
+    }
+
+    func testUpdateEqualNamespaces() {
+        let namespace = stubRequiredNamespaces()
+        var session = WCSession.stub(namespaces: namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(namespace))
+    }
+
+    func testUpdateNamespacesOverRequirement() {
+        let namespace = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount],
+                methods: ["method"],
+                events: ["event"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount],
+                        methods: ["method-2"],
+                        events: ["event-2"])])]
+        let newNamespace = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "newMethod"],
+                events: ["event", "newEvent"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount, polyAccount],
+                        methods: ["method-2", "newMethod-2"],
+                        events: ["event-2", "newEvent-2"])])]
+        var session = WCSession.stub(namespaces: namespace)
+        XCTAssertNoThrow(try session.updateNamespaces(newNamespace))
+    }
+
+    func testUpdateLessThanRequiredChains() {
+        var session = WCSession.stub(namespaces: stubRequiredNamespaces())
+        XCTAssertThrowsError(try session.updateNamespaces([:]))
+    }
+
+    func testUpdateLessThanRequiredAccounts() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: nil)]
+        var session = WCSession.stub(namespaces: stubRequiredNamespaces())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredMethods() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method"],
+                events: ["event", "event-2"],
+                extensions: nil)]
+        var session = WCSession.stub(namespaces: stubRequiredNamespaces())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredEvents() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event"],
+                extensions: nil)]
+        var session = WCSession.stub(namespaces: stubRequiredNamespaces())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredExtension() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: nil)]
+        var session = WCSession.stub(namespaces: stubRequiredNamespacesWithExtension())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredExtensionAccounts() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount],
+                        methods: ["method-2", "newMethod-2"],
+                        events: ["event-2", "newEvent-2"])])]
+        var session = WCSession.stub(namespaces: stubRequiredNamespacesWithExtension())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredExtensionMethods() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount, polyAccount],
+                        methods: ["method-2"],
+                        events: ["event-2", "newEvent-2"])])]
+        var session = WCSession.stub(namespaces: stubRequiredNamespacesWithExtension())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
+
+    func testUpdateLessThanRequiredExtensionEvents() {
+        let invalid = [
+            "eip155": SessionNamespace(
+                accounts: [ethAccount, polyAccount],
+                methods: ["method", "method-2"],
+                events: ["event", "event-2"],
+                extensions: [
+                    SessionNamespace.Extension(
+                        accounts: [ethAccount, polyAccount],
+                        methods: ["method-2", "newMethod-2"],
+                        events: ["event-2"])])]
+        var session = WCSession.stub(namespaces: stubRequiredNamespacesWithExtension())
+        XCTAssertThrowsError(try session.updateNamespaces(invalid))
+    }
 }