-
-
Notifications
You must be signed in to change notification settings - Fork 6
/
NIOWebSocketServerUpgrade+ext.swift
97 lines (93 loc) · 4.47 KB
/
NIOWebSocketServerUpgrade+ext.swift
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
//===----------------------------------------------------------------------===//
//
// This source file is part of the Hummingbird server framework project
//
// Copyright (c) 2023 the Hummingbird authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
import HTTPTypes
import NIOConcurrencyHelpers
import NIOCore
import NIOHTTP1
import NIOHTTPTypesHTTP1
import NIOWebSocket
/// Should HTTP channel upgrade to WebSocket
public enum ShouldUpgradeResult<Value: Sendable>: Sendable {
case dontUpgrade
case upgrade(HTTPFields, Value)
/// Map upgrade result to difference type
func map<Result>(_ map: (Value) throws -> Result) rethrows -> ShouldUpgradeResult<Result> {
switch self {
case .dontUpgrade:
return .dontUpgrade
case .upgrade(let headers, let value):
return try .upgrade(headers, map(value))
}
}
}
extension NIOTypedWebSocketServerUpgrader {
/// Create a new ``NIOTypedWebSocketServerUpgrader``.
///
/// - Parameters:
/// - maxFrameSize: The maximum frame size the decoder is willing to tolerate from the
/// remote peer. WebSockets in principle allows frame sizes up to `2**64` bytes, but
/// this is an objectively unreasonable maximum value (on AMD64 systems it is not
/// possible to even. Users may set this to any value up to `UInt32.max`.
/// - enableAutomaticErrorHandling: Whether the pipeline should automatically handle protocol
/// errors by sending error responses and closing the connection. Defaults to `true`,
/// may be set to `false` if the user wishes to handle their own errors.
/// - shouldUpgrade: A callback that determines whether the websocket request should be
/// upgraded. This callback is responsible for creating a `HTTPHeaders` object with
/// any headers that it needs on the response *except for* the `Upgrade`, `Connection`,
/// and `Sec-WebSocket-Accept` headers, which this upgrader will handle. It also
/// creates state that is passed onto the `upgradePipelineHandler`. Should return
/// an `EventLoopFuture` containing `.notUpgraded` if the upgrade should be refused.
/// - upgradePipelineHandler: A function that will be called once the upgrade response is
/// flushed, and that is expected to mutate the `Channel` appropriately to handle the
/// websocket protocol. This only needs to add the user handlers: the
/// `WebSocketFrameEncoder` and `WebSocketFrameDecoder` will have been added to the
/// pipeline automatically.
convenience init<Value>(
maxFrameSize: Int = 1 << 14,
enableAutomaticErrorHandling: Bool = true,
shouldUpgrade: @escaping @Sendable (Channel, HTTPRequest) -> EventLoopFuture<ShouldUpgradeResult<Value>>,
upgradePipelineHandler: @escaping @Sendable (Channel, Value) -> EventLoopFuture<UpgradeResult>
) {
let shouldUpgradeResult = NIOLockedValueBox<Value?>(nil)
self.init(
maxFrameSize: maxFrameSize,
enableAutomaticErrorHandling: enableAutomaticErrorHandling,
shouldUpgrade: { (channel, head: HTTPRequestHead) in
let request: HTTPRequest
do {
request = try HTTPRequest(head, secure: false, splitCookie: false)
} catch {
return channel.eventLoop.makeFailedFuture(error)
}
return shouldUpgrade(channel, request).map { result in
switch result {
case .dontUpgrade:
return nil
case .upgrade(let headers, let value):
shouldUpgradeResult.withLockedValue { $0 = value }
return .init(headers)
}
}
},
upgradePipelineHandler: { channel, _ in
let result = shouldUpgradeResult.withLockedValue { $0 }
if let result {
return upgradePipelineHandler(channel, result)
} else {
preconditionFailure("Should not be running updatePipelineHander if shouldUpgrade failed")
}
}
)
}
}