feat: update auth to allow eventloopfuture resolution

Author: Garrett Moseke

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?
     }
 

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