graphql-hummingbird/Sources/GraphQLHummingbird/WebSocket/GraphQLHandler+handleWebSocket.swift at abc91550176f3a26fb54e091bc331fa07ce8e3b1 · 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
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(inbound: inbound, outbound: outbound, logger: logger)

        switch subProtocol {
        case .graphqlTransportWs:
            // https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md
            let server = GraphQLTransportWS.Server<WebSocketInit, AsyncThrowingStream<GraphQLResult, Error>>(
                messenger: messenger,
                onExecute: { graphQLRequest in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest
                    )
                    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 in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest
                    )
                    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()
                }
            )
            server.onMessage { message in
                logger.trace("GraphQL server received: \(String(message))")
            }
            server.auth(config.websocket.onWebSocketInit)
        case .graphqlWs:
            // https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
            let server = GraphQLWS.Server<WebSocketInit, AsyncThrowingStream<GraphQLResult, Error>>(
                messenger: messenger,
                onExecute: { graphQLRequest in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest
                    )
                    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 in
                    let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context>(
                        hummingbirdRequest: context.request,
                        hummingbirdContext: context.requestContext,
                        graphQLRequest: graphQLRequest
                    )
                    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()
                }
            )
            server.onMessage { message in
                logger.trace("GraphQL server received: \(String(message))")
            }
            server.auth(config.websocket.onWebSocketInit)
        }
        try await messenger.start()
    }

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