From 2d6dfa8b3a6900813771fa02b48de284e0f0bc73 Mon Sep 17 00:00:00 2001
From: Garrett Moseke <garrett@passivelogic.com>
Date: Fri, 4 Nov 2022 10:15:06 -0600
Subject: [PATCH] feat: update auth to allow eventloopfuture resolution

---
 Sources/GraphQLWS/Server.swift            | 25 +++++++++-----
 Tests/GraphQLWSTests/GraphQLWSTests.swift | 40 ++++++++++++++++++++---
 2 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/Sources/GraphQLWS/Server.swift b/Sources/GraphQLWS/Server.swift
index fb9ea17..188911d 100644
--- a/Sources/GraphQLWS/Server.swift
+++ b/Sources/GraphQLWS/Server.swift
@@ -15,8 +15,8 @@ public class Server<InitPayload: Equatable & Codable> {
     
     let onExecute: (GraphQLRequest) -> EventLoopFuture<GraphQLResult>
     let onSubscribe: (GraphQLRequest) -> EventLoopFuture<SubscriptionResult>
+    var auth: (InitPayload) throws -> EventLoopFuture<Void>
     
-    var auth: (InitPayload) throws -> Void = { _ in }
     var onExit: () -> Void = { }
     var onMessage: (String) -> Void = { _ in }
     var onOperationComplete: (String) -> Void = { _ in }
@@ -34,14 +34,17 @@ public class Server<InitPayload: Equatable & Codable> {
     ///   - messenger: The messenger to bind the server to.
     ///   - onExecute: Callback run during `start` resolution for non-streaming queries. Typically this is `API.execute`.
     ///   - onSubscribe: Callback run during `start` resolution for streaming queries. Typically this is `API.subscribe`.
+    ///   - eventLoop: EventLoop on which to perform server operations.
     public init(
         messenger: Messenger,
         onExecute: @escaping (GraphQLRequest) -> EventLoopFuture<GraphQLResult>,
-        onSubscribe: @escaping (GraphQLRequest) -> EventLoopFuture<SubscriptionResult>
+        onSubscribe: @escaping (GraphQLRequest) -> EventLoopFuture<SubscriptionResult>,
+        eventLoop: EventLoop
     ) {
         self.messenger = messenger
         self.onExecute = onExecute
         self.onSubscribe = onSubscribe
+        self.auth = { _ in eventLoop.makeSucceededVoidFuture() }
         
         messenger.onReceive { message in
             guard let messenger = self.messenger else { return }
@@ -100,10 +103,10 @@ public class Server<InitPayload: Equatable & Codable> {
         }
     }
     
-    /// Define the callback run during `connection_init` resolution that allows authorization using the `payload`.
-    /// Throw from this closure to indicate that authorization has failed.
+    /// Define a custom callback run during `connection_init` resolution that allows authorization using the `payload`.
+    /// Throw or fail the future from this closure to indicate that authorization has failed.
     /// - Parameter callback: The callback to assign
-    public func auth(_ callback: @escaping (InitPayload) throws -> Void) {
+    public func auth(_ callback: @escaping (InitPayload) throws -> EventLoopFuture<Void>) {
         self.auth = callback
     }
     
@@ -138,14 +141,20 @@ public class Server<InitPayload: Equatable & Codable> {
         }
         
         do {
-            try self.auth(connectionInitRequest.payload)
+            let authResult = try self.auth(connectionInitRequest.payload)
+            authResult.whenSuccess {
+                self.initialized = true
+                self.sendConnectionAck()
+            }
+            authResult.whenFailure { error in
+                self.error(.unauthorized())
+                return
+            }
         }
         catch {
             self.error(.unauthorized())
             return
         }
-        initialized = true
-        self.sendConnectionAck()
         // TODO: Should we send the `ka` message?
     }
     
diff --git a/Tests/GraphQLWSTests/GraphQLWSTests.swift b/Tests/GraphQLWSTests/GraphQLWSTests.swift
index d160aec..da2b94b 100644
--- a/Tests/GraphQLWSTests/GraphQLWSTests.swift
+++ b/Tests/GraphQLWSTests/GraphQLWSTests.swift
@@ -12,6 +12,7 @@ class GraphqlWsTests: XCTestCase {
     var clientMessenger: TestMessenger!
     var serverMessenger: TestMessenger!
     var server: Server<TokenInitPayload>!
+    var eventLoop: EventLoop!
     
     override func setUp() {
         // Point the client and server at each other
@@ -20,7 +21,7 @@ class GraphqlWsTests: XCTestCase {
         clientMessenger.other = serverMessenger
         serverMessenger.other = clientMessenger
         
-        let eventLoop = MultiThreadedEventLoopGroup(numberOfThreads: 1).next()
+        eventLoop = MultiThreadedEventLoopGroup(numberOfThreads: 1).next()
         let api = TestAPI()
         let context = TestContext()
         
@@ -30,16 +31,17 @@ class GraphqlWsTests: XCTestCase {
                 api.execute(
                     request: graphQLRequest.query,
                     context: context,
-                    on: eventLoop
+                    on: self.eventLoop
                 )
             },
             onSubscribe: { graphQLRequest in
                 api.subscribe(
                     request: graphQLRequest.query,
                     context: context,
-                    on: eventLoop
+                    on: self.eventLoop
                 )
-            }
+            },
+            eventLoop: self.eventLoop
         )
     }
     
@@ -73,7 +75,7 @@ class GraphqlWsTests: XCTestCase {
     }
     
     /// Tests that throwing in the authorization callback forces an unauthorized error
-    func testAuth() throws {
+    func testAuthWithThrow() throws {
         server.auth { payload in
             throw TestError.couldBeAnything
         }
@@ -100,6 +102,34 @@ class GraphqlWsTests: XCTestCase {
         )
     }
     
+    /// Tests that failing a future in the authorization callback forces an unauthorized error
+    func testAuthWithFailedFuture() throws {
+        server.auth { payload in
+            self.eventLoop.makeFailedFuture(TestError.couldBeAnything)
+        }
+        
+        var messages = [String]()
+        let completeExpectation = XCTestExpectation()
+        
+        let client = Client<TokenInitPayload>(messenger: clientMessenger)
+        client.onMessage { message, _ in
+            messages.append(message)
+            completeExpectation.fulfill()
+        }
+        
+        client.sendConnectionInit(
+            payload: TokenInitPayload(
+                authToken: ""
+            )
+        )
+        
+        wait(for: [completeExpectation], timeout: 2)
+        XCTAssertEqual(
+            messages,
+            ["\(ErrorCode.unauthorized): Unauthorized"]
+        )
+    }
+    
     /// Test single op message flow works as expected
     func testSingleOp() throws {
         let id = UUID().description