From a8e27946b1e10fefb5f3fcf35a2cd686559550e5 Mon Sep 17 00:00:00 2001
From: Maksim Zdobnikau <43750648+DelevoXDG@users.noreply.github.com>
Date: Wed, 28 Feb 2024 14:45:40 +0100
Subject: [PATCH] Support string, selector base types in StarknetTypedData
 (#158)

---
 Sources/Starknet/Data/StarknetTypedData.swift | 34 +++++++++++++------
 Tests/StarknetTests/Data/TypedDataTests.swift | 28 +++++++++++----
 .../typed_data_validate_example.json          | 33 ++++++++++++++++++
 3 files changed, 78 insertions(+), 17 deletions(-)
 create mode 100644 Tests/StarknetTests/Resources/TypedData/typed_data_validate_example.json

diff --git a/Sources/Starknet/Data/StarknetTypedData.swift b/Sources/Starknet/Data/StarknetTypedData.swift
index 218a247a5..2e76165fd 100644
--- a/Sources/Starknet/Data/StarknetTypedData.swift
+++ b/Sources/Starknet/Data/StarknetTypedData.swift
@@ -63,8 +63,11 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
     public let message: [String: Element]
 
     private init?(types: [String: [TypeDeclaration]], primaryType: String, domain: [String: Element], message: [String: Element]) {
-        if types.keys.contains("felt") || types.keys.contains("felt*") {
-            return nil
+        let reservedTypeNames = ["felt", "felt*", "string", "selector"]
+        for typeName in reservedTypeNames {
+            if types.keys.contains(typeName) {
+                return nil
+            }
         }
 
         self.types = types
@@ -151,23 +154,21 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
             return hash
         }
 
-        if typeName == "felt*" {
+        switch typeName {
+        case "felt*":
             let array = try unwrapArray(from: element)
-
             let hashes = try array.map {
                 try unwrapFelt(from: $0)
             }
-
             let hash = StarknetCurve.pedersenOn(hashes)
-
             return hash
-        }
-
-        if typeName == "felt" {
+        case "felt", "string":
             return try unwrapFelt(from: element)
+        case "selector":
+            return try unwrapSelector(from: element)
+        default:
+            throw StarknetTypedDataError.decodingError
         }
-
-        throw StarknetTypedDataError.decodingError
     }
 
     private func encode(data: [String: Element], forType typeName: String) throws -> [Felt] {
@@ -309,6 +310,17 @@ private extension StarknetTypedData {
             throw StarknetTypedDataError.decodingError
         }
     }
+
+    func unwrapSelector(from element: Element) throws -> Felt {
+        switch element {
+        case let .felt(felt):
+            return felt
+        case let .string(string):
+            return starknetSelector(from: string)
+        default:
+            throw StarknetTypedDataError.decodingError
+        }
+    }
 }
 
 private extension String {
diff --git a/Tests/StarknetTests/Data/TypedDataTests.swift b/Tests/StarknetTests/Data/TypedDataTests.swift
index e8aee49b1..e77d5a6f6 100644
--- a/Tests/StarknetTests/Data/TypedDataTests.swift
+++ b/Tests/StarknetTests/Data/TypedDataTests.swift
@@ -13,15 +13,19 @@ final class TypedDataTests: XCTestCase {
     static let tdFeltArr = loadTypedDataFromFile(name: "typed_data_felt_array_example")!
     static let tdString = loadTypedDataFromFile(name: "typed_data_long_string_example")!
     static let tdStructArr = loadTypedDataFromFile(name: "typed_data_struct_array_example")!
+    static let tdValidate = loadTypedDataFromFile(name: "typed_data_validate_example")!
 
     func testInvalidTypes() {
-        XCTAssertNil(
-            StarknetTypedData(types: ["felt": []], primaryType: "felt", domain: "{}", message: "{\"felt\": 1}")
-        )
+        func testInvalidType(_ type: String) {
+            XCTAssertNil(
+                StarknetTypedData(types: [type: []], primaryType: type, domain: "{}", message: "{\"\(type)\": 1}")
+            )
+        }
 
-        XCTAssertNil(
-            StarknetTypedData(types: ["felt*": []], primaryType: "felt*", domain: "{}", message: "{\"felt*\": 1}")
-        )
+        testInvalidType("felt")
+        testInvalidType("felt*")
+        testInvalidType("string")
+        testInvalidType("selector")
     }
 
     func testMissingDependency() {
@@ -47,6 +51,7 @@ final class TypedDataTests: XCTestCase {
             (Self.tdFeltArr, "Mail", "0x5b03497592c0d1fe2f3667b63099761714a895c7df96ec90a85d17bfc7a7a0"),
             (Self.tdStructArr, "Post", "0x1d71e69bf476486b43cdcfaf5a85c00bb2d954c042b281040e513080388356d"),
             (Self.tdStructArr, "Mail", "0x873b878e35e258fc99e3085d5aaad3a81a0c821f189c08b30def2cde55ff27"),
+            (Self.tdValidate, "Validate", "0x2e86ac4735e6012fbeaa68cbd0e5a089014d0da150fa915769a35d5eba30593"),
         ]
 
         try cases.forEach { data, typeName, expectedResult in
@@ -83,6 +88,12 @@ final class TypedDataTests: XCTestCase {
                 "message",
                 "0x5650ec45a42c4776a182159b9d33118a46860a6e6639bb8166ff71f3c41eaef"
             ),
+            (
+                Self.tdValidate,
+                "Validate",
+                "message",
+                "0x87ecd5622070667d2534fa83dd9b16f6cb497b42998d301e8df0ed5875d02d"
+            ),
         ]
 
         try cases.forEach { data, typeName, dataSource, expectedResult in
@@ -116,6 +127,11 @@ final class TypedDataTests: XCTestCase {
                 "0xcd2a3d9f938e13cd947ec05abc7fe734df8dd826",
                 "0x5914ed2764eca2e6a41eb037feefd3d2e33d9af6225a9e7fe31ac943ff712c"
             ),
+            (
+                Self.tdValidate,
+                "0xcd2a3d9f938e13cd947ec05abc7fe734df8dd826",
+                "0x28e38c1c65783abb40b871705095584b96bcbf1f80c8268a0659d074b3afd92"
+            ),
         ]
 
         try cases.forEach { data, address, expectedResult in
diff --git a/Tests/StarknetTests/Resources/TypedData/typed_data_validate_example.json b/Tests/StarknetTests/Resources/TypedData/typed_data_validate_example.json
new file mode 100644
index 000000000..81f73d5d2
--- /dev/null
+++ b/Tests/StarknetTests/Resources/TypedData/typed_data_validate_example.json
@@ -0,0 +1,33 @@
+{
+  "types": {
+    "StarkNetDomain": [
+      {"name": "name", "type": "string"},
+      {"name": "version", "type": "felt"},
+      {"name": "chainId", "type": "felt"}
+    ],
+    "Validate": [
+      {"name": "id", "type": "felt"},
+      {"name": "from", "type": "felt"},
+      {"name": "amount", "type": "felt"},
+      {"name": "nameGamer", "type": "string"},
+      {"name": "endDate", "type": "felt"},
+      {"name": "itemsAuthorized", "type": "felt*"},
+      {"name": "chkFunction", "type": "selector"}
+    ]
+  },
+  "primaryType": "Validate",
+  "domain": {
+    "name": "myDapp",
+    "version": "1",
+    "chainId": "SN_GOERLI"
+  },
+  "message": {
+    "id": "0x0000004f000f",
+    "from": "0x2c94f628d125cd0e86eaefea735ba24c262b9a441728f63e5776661829a4066",
+    "amount": "400",
+    "nameGamer": "Hector26",
+    "endDate": "0x27d32a3033df4277caa9e9396100b7ca8c66a4ef8ea5f6765b91a7c17f0109c",
+    "itemsAuthorized": ["0x01", "0x03", "0x0a", "0x0e"],
+    "chkFunction": "check_authorization"
+  }
+}