From 77121b1a5a1d145f4107b9dae5884744e6c126d1 Mon Sep 17 00:00:00 2001
From: Artur Guseinov <arthur.guseinov@gmail.com>
Date: Wed, 4 Oct 2023 20:34:27 +0800
Subject: [PATCH] JWT tests fixed

---
 Sources/WalletConnectJWT/JSONEncoder+JWT.swift      | 11 +++++++++++
 Sources/WalletConnectJWT/JWT.swift                  |  6 +++---
 Sources/WalletConnectJWT/JWTEncodable.swift         |  7 ++-----
 Tests/RelayerTests/AuthTests/EdDSASignerTests.swift |  4 ++--
 Tests/RelayerTests/AuthTests/JWTTests.swift         |  7 +++++--
 5 files changed, 23 insertions(+), 12 deletions(-)
 create mode 100644 Sources/WalletConnectJWT/JSONEncoder+JWT.swift

diff --git a/Sources/WalletConnectJWT/JSONEncoder+JWT.swift b/Sources/WalletConnectJWT/JSONEncoder+JWT.swift
new file mode 100644
index 0000000000..7f73ad98cf
--- /dev/null
+++ b/Sources/WalletConnectJWT/JSONEncoder+JWT.swift
@@ -0,0 +1,11 @@
+import Foundation
+
+extension JSONEncoder {
+
+    public static var jwt: JSONEncoder {
+        let jsonEncoder = JSONEncoder()
+        jsonEncoder.outputFormatting = .withoutEscapingSlashes
+        jsonEncoder.dateEncodingStrategy = .secondsSince1970
+        return jsonEncoder
+    }
+}
diff --git a/Sources/WalletConnectJWT/JWT.swift b/Sources/WalletConnectJWT/JWT.swift
index 5e32594409..d377288402 100644
--- a/Sources/WalletConnectJWT/JWT.swift
+++ b/Sources/WalletConnectJWT/JWT.swift
@@ -7,12 +7,12 @@ struct JWT<JWTClaims: JWTEncodable>: Codable, Equatable {
     let signature: String
     let string: String
 
-    init(claims: JWTClaims, signer: JWTSigning) throws {
+    init(claims: JWTClaims, signer: JWTSigning, jsonEncoder: JSONEncoder = .jwt) throws {
         self.header = JWTHeader(alg: signer.alg)
         self.claims = claims
 
-        let headerString = try header.encode()
-        let claimsString = try claims.encode()
+        let headerString = try header.encode(jsonEncoder: jsonEncoder)
+        let claimsString = try claims.encode(jsonEncoder: jsonEncoder)
         let signature = try signer.sign(header: headerString, claims: claimsString)
         
         self.signature = signature
diff --git a/Sources/WalletConnectJWT/JWTEncodable.swift b/Sources/WalletConnectJWT/JWTEncodable.swift
index 04505fdb5f..e612e47f8b 100644
--- a/Sources/WalletConnectJWT/JWTEncodable.swift
+++ b/Sources/WalletConnectJWT/JWTEncodable.swift
@@ -1,17 +1,14 @@
 import Foundation
 
 public protocol JWTEncodable: Codable, Equatable {
-    func encode() throws -> String
+    func encode(jsonEncoder: JSONEncoder) throws -> String
 
     static func decode(from string: String) throws -> Self
 }
 
 extension JWTEncodable {
 
-    public func encode() throws -> String {
-        let jsonEncoder = JSONEncoder()
-        jsonEncoder.outputFormatting = .withoutEscapingSlashes
-        jsonEncoder.dateEncodingStrategy = .secondsSince1970
+    public func encode(jsonEncoder: JSONEncoder) throws -> String {
         let data = try jsonEncoder.encode(self)
         return JWTEncoder.base64urlEncodedString(data: data)
     }
diff --git a/Tests/RelayerTests/AuthTests/EdDSASignerTests.swift b/Tests/RelayerTests/AuthTests/EdDSASignerTests.swift
index 87d5ae4c3b..cbd8d2878c 100644
--- a/Tests/RelayerTests/AuthTests/EdDSASignerTests.swift
+++ b/Tests/RelayerTests/AuthTests/EdDSASignerTests.swift
@@ -11,8 +11,8 @@ final class EdDSASignerTests: XCTestCase {
         let keyRaw = Data(hex: "58e0254c211b858ef7896b00e3f36beeb13d568d47c6031c4218b87718061295")
         let signingKey = try! SigningPrivateKey(rawRepresentation: keyRaw)
         sut = EdDSASigner(signingKey)
-        let header = try! JWTHeader(alg: "EdDSA").encode()
-        let claims = try! RelayAuthPayload.Claims.stub().encode()
+        let header = try! JWTHeader(alg: "EdDSA").encode(jsonEncoder: .jwt)
+        let claims = try! RelayAuthPayload.Claims.stub().encode(jsonEncoder: .jwt)
         let signature = try! sut.sign(header: header, claims: claims)
         XCTAssertNotNil(signature)
     }
diff --git a/Tests/RelayerTests/AuthTests/JWTTests.swift b/Tests/RelayerTests/AuthTests/JWTTests.swift
index b0b044de11..9e662ff721 100644
--- a/Tests/RelayerTests/AuthTests/JWTTests.swift
+++ b/Tests/RelayerTests/AuthTests/JWTTests.swift
@@ -4,12 +4,15 @@ import XCTest
 @testable import WalletConnectJWT
 
 final class JWTTests: XCTestCase {
-    let expectedJWT =  "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJjNDc5ZmU1ZGM0NjRlNzcxZTc4YjE5M2QyMzlhNjViNThkMjc4Y2FkMWMzNGJmYjBiNTcxNmU1YmI1MTQ5MjhlIiwiYXVkIjoid3NzOi8vcmVsYXkud2FsbGV0Y29ubmVjdC5jb20iLCJpYXQiOjE2NTY5MTAwOTcsImlzcyI6ImRpZDprZXk6ejZNa29kSFp3bmVWUlNodGFMZjhKS1lreHBER3AxdkdabnBHbWRCcFg4TTJleHhIIiwiZXhwIjoxNjU2OTk2NDk3fQ.0JkxOM-FV21U7Hk-xycargj_qNRaYV2H5HYtE4GzAeVQYiKWj7YySY5AdSqtCgGzX4Gt98XWXn2kSr9rE1qvCA"
+    let expectedJWT =  "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJhdWQiOiJ3c3M6Ly9yZWxheS53YWxsZXRjb25uZWN0LmNvbSIsImV4cCI6MTY1Njk5NjQ5NywiaWF0IjoxNjU2OTEwMDk3LCJpc3MiOiJkaWQ6a2V5Ono2TWtvZEhad25lVlJTaHRhTGY4SktZa3hwREdwMXZHWm5wR21kQnBYOE0yZXh4SCIsInN1YiI6ImM0NzlmZTVkYzQ2NGU3NzFlNzhiMTkzZDIzOWE2NWI1OGQyNzhjYWQxYzM0YmZiMGI1NzE2ZTViYjUxNDkyOGUifQ.0JkxOM-FV21U7Hk-xycargj_qNRaYV2H5HYtE4GzAeVQYiKWj7YySY5AdSqtCgGzX4Gt98XWXn2kSr9rE1qvCA"
 
     func testJWTEncoding() throws {
         let signer = EdDSASignerMock()
         signer.signature = "0JkxOM-FV21U7Hk-xycargj_qNRaYV2H5HYtE4GzAeVQYiKWj7YySY5AdSqtCgGzX4Gt98XWXn2kSr9rE1qvCA"
-        let jwt = try JWT(claims: RelayAuthPayload.Claims.stub(), signer: signer)
+        let jsonEncoder = JSONEncoder()
+        jsonEncoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes]
+        jsonEncoder.dateEncodingStrategy = .secondsSince1970
+        let jwt = try JWT(claims: RelayAuthPayload.Claims.stub(), signer: signer, jsonEncoder: jsonEncoder)
         XCTAssertEqual(expectedJWT, jwt.string)
     }