graphql-hummingbird/Sources/GraphQLHummingbird/WebSocket/GraphQLHandler+handleWebSocket.swift at 3a5c927e95de5b203307284e2d9d8c8010beb935 · GraphQLSwift/graphql-hummingbird · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import GraphQL
import struct GraphQLTransportWS.EmptyInitPayload
import class GraphQLTransportWS.Server
import class GraphQLWS.Server
import Hummingbird
import HummingbirdWebSocket
import Logging

extension GraphQLHandler where Context: WebSocketRequestContext {
    func handleWebSocket(
        inbound: WebSocketInboundStream,
        outbound: WebSocketOutboundWriter,
        context: WebSocketRouterContext<Context>,
        subProtocol: WebSocketSubProtocol,
        logger: Logger
    ) async throws {
        let messenger = WebSocketMessenger(outbound: outbound, logger: logger)

        let messageStream = inbound.messages(maxSize: context.requestContext.maxUploadSize).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, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                messenger: messenger,
                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,
                        websocketInitResult: initResult
                    )
                    let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                    return try await graphql(
                        schema: self.schema,
                        request: graphQLRequest.query,
                        rootValue: self.rootValue,
                        context: graphQLContext,
                        variableValues: graphQLRequest.variables,
                        operationName: graphQLRequest.operationName
                    )
                },
                onSubscribe: { graphQLRequest, initResult in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest,
                        websocketInitResult: initResult
                    )
                    let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                    return try await graphqlSubscribe(
                        schema: self.schema,
                        request: graphQLRequest.query,
                        rootValue: self.rootValue,
                        context: graphQLContext,
                        variableValues: graphQLRequest.variables,
                        operationName: graphQLRequest.operationName
                    ).get()
                }
            )
            try await server.listen(to: messageStream)
        case .graphqlWs:
            // https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
            let server = GraphQLWS.Server<WebSocketInit, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                messenger: messenger,
                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,
                        websocketInitResult: initResult
                    )
                    let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                    return try await graphql(
                        schema: self.schema,
                        request: graphQLRequest.query,
                        rootValue: self.rootValue,
                        context: graphQLContext,
                        variableValues: graphQLRequest.variables,
                        operationName: graphQLRequest.operationName
                    )
                },
                onSubscribe: { graphQLRequest, initResult in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest,
                        websocketInitResult: initResult
                    )
                    let graphQLContext = try await computeContext(graphQLContextComputationInputs)
                    return try await graphqlSubscribe(
                        schema: self.schema,
                        request: graphQLRequest.query,
                        rootValue: self.rootValue,
                        context: graphQLContext,
                        variableValues: graphQLRequest.variables,
                        operationName: graphQLRequest.operationName
                    ).get()
                }
            )
            try await server.listen(to: messageStream)
        }
    }

    func shouldUpgrade(request: Request) throws -> RouterShouldUpgrade {
        let subProtocol = try negotiateSubProtocol(request: request)
        return .upgrade([.secWebSocketProtocol: subProtocol.rawValue])
    }

    func negotiateSubProtocol(request: Request) throws -> WebSocketSubProtocol {
        var subProtocol: WebSocketSubProtocol?
        let requestedSubProtocols = request.headers[values: .secWebSocketProtocol]
        if requestedSubProtocols.isEmpty {
            // Default
            subProtocol = .graphqlTransportWs
        } else {
            // Choose highest client preference that we understand
            for requestedSubProtocol in requestedSubProtocols {
                if let selectedSubProtocol = WebSocketSubProtocol(rawValue: requestedSubProtocol) {
                    subProtocol = selectedSubProtocol
                    break
                }
            }
        }
        guard let subProtocol = subProtocol else {
            // If they provided options but none matched, fail
            throw HTTPError(.badRequest, message: "Unable to negotiate subprotocol. \(WebSocketSubProtocol.allCases) are supported.")
        }
        return subProtocol
    }
}