Includes WebSocketInitResult in context computation inputs

Package.swift

@@ -16,8 +16,8 @@ let package = Package(
     dependencies: [
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.0.0"),
         .package(url: "https://github.com/GraphQLSwift/GraphQL.git", from: "4.0.0"),
-        .package(url: "https://github.com/GraphQLSwift/GraphQLTransportWS.git", from: "0.2.1"),
-        .package(url: "https://github.com/GraphQLSwift/GraphQLWS.git", from: "0.2.1"),
+        .package(url: "https://github.com/GraphQLSwift/GraphQLTransportWS.git", from: "1.0.0"),
+        .package(url: "https://github.com/GraphQLSwift/GraphQLWS.git", from: "1.0.0"),
         .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.0.0"),
         .package(url: "https://github.com/hummingbird-project/hummingbird-websocket.git", from: "2.0.0"),
     ],

Sources/GraphQLHummingbird/GraphQLConfig.swift

@@ -3,7 +3,11 @@ import GraphQL
 import Hummingbird
 
 /// Configuration options for GraphQLHummingbird
-public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Sendable {
+public struct GraphQLConfig<
+    Context: RequestContext,
+    WebSocketInit: Equatable & Codable & Sendable,
+    WebSocketInitResult: Sendable
+>: Sendable {
     let allowGet: Bool
     let allowMissingAcceptHeader: Bool
     let coders: Coders
@@ -28,7 +32,7 @@ public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Send
         subscriptionProtocols: Set<SubscriptionProtocol> = [.websocket],
         websocket: WebSocket = .init(
             // Including this strongly-typed argument is required to avoid compiler failures on Swift 6.2.3.
-            onWebSocketInit: { (_: EmptyWebSocketInit) in }
+            onWebSocketInit: { (_: EmptyWebSocketInit, _: Request, _: Context) in }
         ),
         additionalValidationRules: [@Sendable (ValidationContext) -> Visitor] = []
     ) {
@@ -94,14 +98,14 @@ public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Send
     }
 
     public struct WebSocket: Sendable {
-        let onWebSocketInit: @Sendable (WebSocketInit) async throws -> Void
+        let onWebSocketInit: @Sendable (WebSocketInit, Request, Context) async throws -> WebSocketInitResult
 
         /// GraphQL over WebSocket configuration
         /// - Parameter onWebSocketInit: A custom callback run during `connection_init` resolution that allows
         /// authorization using the `payload` field of the `connection_init` message.
         /// Throw from this closure to indicate that authorization has failed.
         public init(
-            onWebSocketInit: @Sendable @escaping (WebSocketInit) async throws -> Void = { (_: EmptyWebSocketInit) in }
+            onWebSocketInit: @Sendable @escaping (WebSocketInit, Request, Context) async throws -> WebSocketInitResult = { (_: EmptyWebSocketInit, _: Request, _: Context) in }
         ) {
             self.onWebSocketInit = onWebSocketInit
         }

Sources/GraphQLHummingbird/GraphQLHandler.swift

@@ -4,18 +4,30 @@ import Hummingbird
 struct GraphQLHandler<
     Context: RequestContext,
     GraphQLContext: Sendable,
-    WebSocketInit: Equatable & Codable & Sendable
+    WebSocketInit: Equatable & Codable & Sendable,
+    WebSocketInitResult: Sendable
 >: Sendable {
     let schema: GraphQLSchema
     let rootValue: any Sendable
-    let config: GraphQLConfig<WebSocketInit>
-    let computeContext: @Sendable (GraphQLContextComputationInputs<Context>) async throws -> GraphQLContext
+    let config: GraphQLConfig<Context, WebSocketInit, WebSocketInitResult>
+    let computeContext: @Sendable (GraphQLContextComputationInputs<Context, WebSocketInitResult>) async throws -> GraphQLContext
 }
 
+/// Request metadata that can be used to construct a GraphQL context
 public struct GraphQLContextComputationInputs<
-    Context: RequestContext
+    Context: RequestContext,
+    WebSocketInitResult: Sendable
 >: Sendable {
+    /// The Hummingbird request that initiated the GraphQL request. In WebSockets, this is the upgrade GET request.
     public let hummingbirdRequest: Request
+
+    /// The Hummingbird context from the request that initiated the GraphQL request. In WebSockets, this is the context from the upgrade GET request.
     public let hummingbirdContext: Context
+
+    /// The decoded GraphQL request, including the raw query, variables, and more
     public let graphQLRequest: GraphQLRequest
+
+    /// The result of the WebSocket's initialization closure. This can be used to customize GraphQL context creation based on the init
+    /// message metadata as opposed to only the upgrade request. In non-WebSocket contexts, this is nil.
+    public let websocketInitResult: WebSocketInitResult?
 }

Sources/GraphQLHummingbird/HTTP/GraphQLHandler+HTTP.swift

@@ -24,10 +24,11 @@ extension GraphQLHandler {
         guard operationType != .mutation else {
             throw HTTPError(.methodNotAllowed, message: "Mutations using GET are disallowed")
         }
-        let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+        let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
             hummingbirdRequest: request,
             hummingbirdContext: context,
-            graphQLRequest: graphQLRequest
+            graphQLRequest: graphQLRequest,
+            websocketInitResult: nil
         )
         let graphQLContext = try await computeContext(graphQLContextComputationInputs)
         let result = await execute(
@@ -63,10 +64,11 @@ extension GraphQLHandler {
             throw HTTPError(.unsupportedMediaType)
         }
 
-        let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+        let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
             hummingbirdRequest: request,
             hummingbirdContext: context,
-            graphQLRequest: graphQLRequest
+            graphQLRequest: graphQLRequest,
+            websocketInitResult: nil
         )
         let graphQLContext = try await computeContext(graphQLContextComputationInputs)
         let result = await execute(

Sources/GraphQLHummingbird/Router+graphql.swift

@@ -19,11 +19,11 @@ public extension RouterMethods {
         _ path: RouterPath = "graphql",
         schema: GraphQLSchema,
         rootValue: any Sendable = (),
-        config: GraphQLConfig<EmptyWebSocketInit> = .init(),
-        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context>) async throws -> GraphQLContext
+        config: GraphQLConfig<Context, EmptyWebSocketInit, Void> = .init(),
+        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, Void>) async throws -> GraphQLContext
     ) -> Self {
         // https://github.com/graphql/graphql-over-http/blob/main/spec/GraphQLOverHTTP.md#request
-        let handler = GraphQLHandler<Context, GraphQLContext, EmptyWebSocketInit>(
+        let handler = GraphQLHandler<Context, GraphQLContext, EmptyWebSocketInit, Void>(
             schema: schema,
             rootValue: rootValue,
             config: config,
@@ -78,15 +78,16 @@ public extension RouterMethods where Context: WebSocketRequestContext {
     @discardableResult
     func graphqlWebSocket<
         GraphQLContext: Sendable,
-        WebSocketInit: Equatable & Codable & Sendable
+        WebSocketInit: Equatable & Codable & Sendable,
+        WebSocketInitResult: Sendable
     >(
         _ path: RouterPath = "graphql",
         schema: GraphQLSchema,
         rootValue: any Sendable = (),
-        config: GraphQLConfig<WebSocketInit> = GraphQLConfig<EmptyWebSocketInit>(),
-        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context>) async throws -> GraphQLContext
+        config: GraphQLConfig<Context, WebSocketInit, WebSocketInitResult>,
+        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, WebSocketInitResult>) async throws -> GraphQLContext
     ) -> Self {
-        let handler = GraphQLHandler<Context, GraphQLContext, WebSocketInit>(
+        let handler = GraphQLHandler<Context, GraphQLContext, WebSocketInit, WebSocketInitResult>(
             schema: schema,
             rootValue: rootValue,
             config: config,
@@ -107,4 +108,34 @@ public extension RouterMethods where Context: WebSocketRequestContext {
         }
         return self
     }
+
+    /// Registers a graphql websocket route that responds using the provided schema.
+    ///
+    /// WebSocket requests support the
+    /// [`graphql-transport-ws`](https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md)
+    /// and [`graphql-ws`](https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md)
+    /// subprotocols.
+    ///
+    /// - Parameters:
+    ///   - path: The route that should respond to GraphQL requests. Both `GET` and `POST` routes are registered.
+    ///   - schema: The GraphQL schema that should be used to respond to requests.
+    ///   - rootValue: The `rootValue` GraphQL execution arg. This is the object passed to the root resolvers.
+    ///   - computeContext: A closure used to compute the GraphQL context from incoming requests. This must be provided.
+    @discardableResult
+    func graphqlWebSocket<GraphQLContext: Sendable>(
+        _ path: RouterPath = "graphql",
+        schema: GraphQLSchema,
+        rootValue: any Sendable = (),
+        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, Void>) async throws -> GraphQLContext
+    ) -> Self {
+        // This is just an overload that allows not passing `config`, since we cannot use a default argument that
+        // uses `Context` without getting `error: generic parameter 'Self' could not be inferred` compilation errors
+        graphqlWebSocket(
+            path,
+            schema: schema,
+            rootValue: rootValue,
+            config: GraphQLConfig<Context, EmptyWebSocketInit, Void>(),
+            computeContext: computeContext
+        )
+    }
 }

Sources/GraphQLHummingbird/WebSocket/GraphQLHandler+handleWebSocket.swift

@@ -14,18 +14,32 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
         subProtocol: WebSocketSubProtocol,
         logger: Logger
     ) async throws {
-        let messenger = WebSocketMessenger(inbound: inbound, outbound: outbound, logger: logger)
+        let messenger = WebSocketMessenger(outbound: outbound, logger: logger)
+
+        // TODO: Make maxSize configurable
+        let messageStream = inbound.messages(maxSize: 1024 * 1024).compactMap { message -> String? in
+            // TODO: Add binary support
+            guard case let .text(text) = message else {
+                return nil
+            }
+            logger.trace("GraphQL server received: \(message)")
+            return text
+        }
 
         switch subProtocol {
         case .graphqlTransportWs:
             // https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md
-            let server = GraphQLTransportWS.Server<WebSocketInit, AsyncThrowingStream<GraphQLResult, Error>>(
+            let server = GraphQLTransportWS.Server<WebSocketInit, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                 messenger: messenger,
-                onExecute: { graphQLRequest in
-                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+                onInit: { initPayload in
+                    try await config.websocket.onWebSocketInit(initPayload, context.request, context.requestContext)
+                },
+                onExecute: { graphQLRequest, initResult in
+                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                         hummingbirdRequest: context.request,
                         hummingbirdContext: context.requestContext,
-                        graphQLRequest: graphQLRequest
+                        graphQLRequest: graphQLRequest,
+                        websocketInitResult: initResult
                     )
                     let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                     return try await graphql(
@@ -37,11 +51,12 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
                         operationName: graphQLRequest.operationName
                     )
                 },
-                onSubscribe: { graphQLRequest in
-                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+                onSubscribe: { graphQLRequest, initResult in
+                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                         hummingbirdRequest: context.request,
                         hummingbirdContext: context.requestContext,
-                        graphQLRequest: graphQLRequest
+                        graphQLRequest: graphQLRequest,
+                        websocketInitResult: initResult
                     )
                     let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                     return try await graphqlSubscribe(
@@ -54,19 +69,20 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
                     ).get()
                 }
             )
-            server.onMessage { message in
-                logger.trace("GraphQL server received: \(String(message))")
-            }
-            server.auth(config.websocket.onWebSocketInit)
+            try await server.listen(to: messageStream)
         case .graphqlWs:
             // https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
-            let server = GraphQLWS.Server<WebSocketInit, AsyncThrowingStream<GraphQLResult, Error>>(
+            let server = GraphQLWS.Server<WebSocketInit, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                 messenger: messenger,
-                onExecute: { graphQLRequest in
-                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+                onInit: { initPayload in
+                    try await config.websocket.onWebSocketInit(initPayload, context.request, context.requestContext)
+                },
+                onExecute: { graphQLRequest, initResult in
+                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                         hummingbirdRequest: context.request,
                         hummingbirdContext: context.requestContext,
-                        graphQLRequest: graphQLRequest
+                        graphQLRequest: graphQLRequest,
+                        websocketInitResult: initResult
                     )
                     let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                     return try await graphql(
@@ -78,11 +94,12 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
                         operationName: graphQLRequest.operationName
                     )
                 },
-                onSubscribe: { graphQLRequest in
-                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
+                onSubscribe: { graphQLRequest, initResult in
+                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                         hummingbirdRequest: context.request,
                         hummingbirdContext: context.requestContext,
-                        graphQLRequest: graphQLRequest
+                        graphQLRequest: graphQLRequest,
+                        websocketInitResult: initResult
                     )
                     let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                     return try await graphqlSubscribe(
@@ -95,12 +112,8 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
                     ).get()
                 }
             )
-            server.onMessage { message in
-                logger.trace("GraphQL server received: \(String(message))")
-            }
-            server.auth(config.websocket.onWebSocketInit)
+            try await server.listen(to: messageStream)
         }
-        try await messenger.start()
     }
 
     func shouldUpgrade(request: Request) throws -> RouterShouldUpgrade {

Sources/GraphQLHummingbird/WebSocket/WebSocketMessenger.swift

@@ -6,44 +6,22 @@ import NIOCore
 
 /// Messenger wrapper for WebSockets
 class WebSocketMessenger: GraphQLTransportWS.Messenger, GraphQLWS.Messenger, @unchecked Sendable {
-    private let inbound: WebSocketInboundStream
     private let outbound: WebSocketOutboundWriter
     private let logger: Logger
 
-    private var onReceive: (String) async throws -> Void = { _ in }
-
     init(
-        inbound: WebSocketInboundStream,
         outbound: WebSocketOutboundWriter,
         logger: Logger
     ) {
-        self.inbound = inbound
         self.outbound = outbound
         self.logger = logger
     }
 
-    /// A blocking method to start listening to the inbound messages and bind them to the onReceive callback
-    func start() async throws {
-        // TODO: Make maxSize configurable
-        for try await message in inbound.messages(maxSize: 1024 * 1024) {
-            guard case let .text(text) = message else { continue }
-            do {
-                try await onReceive(text)
-            } catch {
-                try? await self.error("\(error)", code: 4400)
-            }
-        }
-    }
-
     func send<S: Collection>(_ message: S) async throws where S.Element == Character {
         logger.trace("GraphQL server sent: \(String(message))")
         try await outbound.write(.text(String(message)))
     }
 
-    func onReceive(callback: @escaping (String) async throws -> Void) {
-        onReceive = callback
-    }
-
     func error(_ message: String, code: Int) async throws {
         try await outbound.close(.init(codeNumber: code), reason: message)
     }