Includes WebSocketInitResult in context computation inputs

Package.swift

@@ -15,8 +15,8 @@ let package = Package(
     ],
     dependencies: [
         .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/vapor/vapor.git", from: "4.0.0"),
     ],
     targets: [

Sources/GraphQLVapor/GraphQLConfig.swift

@@ -1,7 +1,11 @@
 import GraphQL
+import Vapor
 
 /// Configuration options for GraphQLVapor
-public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Sendable {
+public struct GraphQLConfig<
+    WebSocketInit: Equatable & Codable & Sendable,
+    WebSocketInitResult: Sendable
+>: Sendable {
     let allowGet: Bool
     let allowMissingAcceptHeader: Bool
     let ide: IDE
@@ -24,7 +28,7 @@ public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Send
         subscriptionProtocols: Set<SubscriptionProtocol> = [],
         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) in }
         ),
         additionalValidationRules: [@Sendable (ValidationContext) -> Visitor] = []
     ) {
@@ -67,16 +71,16 @@ public struct GraphQLConfig<WebSocketInit: Equatable & Codable & Sendable>: Send
     }
 
     public struct WebSocket: Sendable {
-        let onWebsocketInit: @Sendable (WebSocketInit) async throws -> Void
+        let onWebSocketInit: @Sendable (WebSocketInit, Request) async throws -> WebSocketInitResult
 
         /// GraphQL over WebSocket configuration
-        /// - Parameter onWebsocketInit: A custom callback run during `connection_init` resolution that allows
+        /// - 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) async throws -> WebSocketInitResult = { (_: EmptyWebSocketInit, _: Request) in }
         ) {
-            self.onWebsocketInit = onWebsocketInit
+            self.onWebSocketInit = onWebSocketInit
         }
     }
 }

Sources/GraphQLVapor/GraphQLHandler.swift

@@ -3,10 +3,11 @@ import Vapor
 
 struct GraphQLHandler<
     Context: 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) async throws -> Context
+    let config: GraphQLConfig<WebSocketInit, WebSocketInitResult>
+    let computeContext: @Sendable (GraphQLContextComputationInputs<WebSocketInitResult>) async throws -> Context
 }

Sources/GraphQLVapor/HTTP/GraphQLHandler+HTTP.swift

@@ -15,9 +15,10 @@ extension GraphQLHandler {
         guard operationType != .mutation else {
             throw Abort(.methodNotAllowed, reason: "Mutations using GET are disallowed")
         }
-        let graphQLContextComputationInputs = GraphQLContextComputationInputs(
+        let graphQLContextComputationInputs = GraphQLContextComputationInputs<WebSocketInitResult>(
             vaporRequest: request,
-            graphQLRequest: graphQLRequest
+            graphQLRequest: graphQLRequest,
+            websocketInitResult: nil
         )
         let context = try await computeContext(graphQLContextComputationInputs)
         let result = await execute(
@@ -34,9 +35,10 @@ extension GraphQLHandler {
             throw Abort(.unsupportedMediaType, reason: "Missing `Content-Type` header")
         }
         let graphQLRequest = try request.content.decode(GraphQLRequest.self)
-        let graphQLContextComputationInputs = GraphQLContextComputationInputs(
+        let graphQLContextComputationInputs = GraphQLContextComputationInputs<WebSocketInitResult>(
             vaporRequest: request,
-            graphQLRequest: graphQLRequest
+            graphQLRequest: graphQLRequest,
+            websocketInitResult: nil
         )
         let context = try await computeContext(graphQLContextComputationInputs)
         let result = await execute(

Sources/GraphQLVapor/RoutesBuilder+graphql.swift

@@ -20,19 +20,20 @@ public extension RoutesBuilder {
     ///   - computeContext: A closure used to compute the GraphQL context from incoming requests. This must be provided.
     func graphql<
         Context: Sendable,
-        WebSocketInit: Equatable & Codable & Sendable
+        WebSocketInit: Equatable & Codable & Sendable,
+        WebSocketInitResult: Sendable
     >(
         _ path: [PathComponent] = ["graphql"],
         schema: GraphQLSchema,
         rootValue: any Sendable = (),
-        config: GraphQLConfig<WebSocketInit> = GraphQLConfig<EmptyWebsocketInit>(),
-        computeContext: @Sendable @escaping (GraphQLContextComputationInputs) async throws -> Context
+        config: GraphQLConfig<WebSocketInit, WebSocketInitResult> = GraphQLConfig<EmptyWebSocketInit, Void>(),
+        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<WebSocketInitResult>) async throws -> Context
     ) {
         ContentConfiguration.global.use(encoder: GraphQLJSONEncoder(), for: .jsonGraphQL)
         ContentConfiguration.global.use(decoder: JSONDecoder(), for: .jsonGraphQL)
 
         // https://github.com/graphql/graphql-over-http/blob/main/spec/GraphQLOverHTTP.md#request
-        let handler = GraphQLHandler<Context, WebSocketInit>(schema: schema, rootValue: rootValue, config: config, computeContext: computeContext)
+        let handler = GraphQLHandler<Context, WebSocketInit, WebSocketInitResult>(schema: schema, rootValue: rootValue, config: config, computeContext: computeContext)
         get(path) { request in
             // WebSocket handling
             if
@@ -68,7 +69,17 @@ public extension RoutesBuilder {
     }
 }
 
-public struct GraphQLContextComputationInputs: Sendable {
+/// Request metadata that can be used to construct a GraphQL context
+public struct GraphQLContextComputationInputs<
+    WebSocketInitResult: Sendable
+>: Sendable {
+    /// The Vapor request that initiated the GraphQL request. In WebSockets, this is the upgrade GET request.
     public let vaporRequest: Request
+
+    /// 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/GraphQLVapor/WebSocket/EmptyWebSocketInit.swift

@@ -1 +1 @@
-public struct EmptyWebsocketInit: Equatable, Codable, Sendable {}
+public struct EmptyWebSocketInit: Equatable, Codable, Sendable {}

Sources/GraphQLVapor/WebSocket/GraphQLHandler+handleWebSocket.swift

@@ -18,16 +18,34 @@ extension GraphQLHandler {
                 }
             },
             onUpgrade: { websocket in
+                let messageStream = AsyncThrowingStream<String, Error> { continuation in
+                    websocket.onText { _, text in
+                        continuation.yield(text)
+                    }
+                    websocket.onClose.whenComplete { result in
+                        switch result {
+                        case .success:
+                            continuation.finish()
+                        case let .failure(error):
+                            continuation.finish(throwing: error)
+                        }
+                    }
+                }
+
                 let messenger = WebSocketMessenger(websocket: websocket)
                 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
+                        onInit: { initPayload in
+                            try await config.websocket.onWebSocketInit(initPayload, request)
+                        },
+                        onExecute: { graphQLRequest, initResult in
                             let graphQLContextComputationInputs = GraphQLContextComputationInputs(
                                 vaporRequest: request,
-                                graphQLRequest: graphQLRequest
+                                graphQLRequest: graphQLRequest,
+                                websocketInitResult: initResult
                             )
                             let context = try await computeContext(graphQLContextComputationInputs)
                             return try await graphql(
@@ -39,10 +57,11 @@ extension GraphQLHandler {
                                 operationName: graphQLRequest.operationName
                             )
                         },
-                        onSubscribe: { graphQLRequest in
+                        onSubscribe: { graphQLRequest, initResult in
                             let graphQLContextComputationInputs = GraphQLContextComputationInputs(
                                 vaporRequest: request,
-                                graphQLRequest: graphQLRequest
+                                graphQLRequest: graphQLRequest,
+                                websocketInitResult: initResult
                             )
                             let context = try await computeContext(graphQLContextComputationInputs)
                             return try await graphqlSubscribe(
@@ -55,15 +74,22 @@ extension GraphQLHandler {
                             ).get()
                         }
                     )
-                    server.auth(config.websocket.onWebsocketInit)
+                    Task {
+                        // This task completes upon websocket closure
+                        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
+                        onInit: { initPayload in
+                            try await config.websocket.onWebSocketInit(initPayload, request)
+                        },
+                        onExecute: { graphQLRequest, initResult in
                             let graphQLContextComputationInputs = GraphQLContextComputationInputs(
                                 vaporRequest: request,
-                                graphQLRequest: graphQLRequest
+                                graphQLRequest: graphQLRequest,
+                                websocketInitResult: initResult
                             )
                             let context = try await computeContext(graphQLContextComputationInputs)
                             return try await graphql(
@@ -75,10 +101,11 @@ extension GraphQLHandler {
                                 operationName: graphQLRequest.operationName
                             )
                         },
-                        onSubscribe: { graphQLRequest in
+                        onSubscribe: { graphQLRequest, initResult in
                             let graphQLContextComputationInputs = GraphQLContextComputationInputs(
                                 vaporRequest: request,
-                                graphQLRequest: graphQLRequest
+                                graphQLRequest: graphQLRequest,
+                                websocketInitResult: initResult
                             )
                             let context = try await computeContext(graphQLContextComputationInputs)
                             return try await graphqlSubscribe(
@@ -91,7 +118,10 @@ extension GraphQLHandler {
                             ).get()
                         }
                     )
-                    server.auth(config.websocket.onWebsocketInit)
+                    Task {
+                        // This task completes upon websocket closure
+                        try await server.listen(to: messageStream)
+                    }
                 }
             }
         )

Sources/GraphQLVapor/WebSocket/WebSocketMessenger.swift

@@ -3,47 +3,19 @@ import GraphQLWS
 import WebSocketKit
 
 /// Messenger wrapper for WebSockets
-class WebSocketMessenger: GraphQLTransportWS.Messenger, GraphQLWS.Messenger, @unchecked Sendable {
-    private weak var websocket: WebSocket?
-    private var onReceive: (String) async throws -> Void = { _ in }
-
-    init(websocket: WebSocket) {
-        self.websocket = websocket
-        websocket.onText { _, message in
-            // We must include self here, without a weak reference to prevent it from falling
-            // out of scope while the websocket is still alive
-            do {
-                try await self.onReceive(message)
-            } catch {
-                try? await self.error("\(error)", code: 4400)
-            }
-        }
-        websocket.onClose.whenComplete { [weak self] _ in
-            guard let self = self else {
-                return
-            }
-            self.onReceive { _ in }
-            websocket.onText { _, _ in }
-        }
-    }
+struct WebSocketMessenger: GraphQLTransportWS.Messenger, GraphQLWS.Messenger {
+    let websocket: WebSocket
 
     func send<S: Collection>(_ message: S) async throws where S.Element == Character {
-        guard let websocket = websocket else { return }
         try await websocket.send(message)
     }
 
-    func onReceive(callback: @escaping (String) async throws -> Void) {
-        onReceive = callback
-    }
-
     func error(_ message: String, code: Int) async throws {
-        guard let websocket = websocket else { return }
         try await websocket.send("\(code): \(message)")
         try await websocket.close(code: .init(codeNumber: code))
     }
 
     func close() async throws {
-        guard let websocket = websocket else { return }
         try await websocket.close()
     }
 }