add support for buffered bytes from websocket channel

johnnzhou · johnnzhou · commit c7ed19d21781 · 2024-09-05T00:15:52.000-07:00
diff --git a/Package.swift b/Package.swift
@@ -7,17 +7,17 @@ let package = Package(
         .macOS(.v10_15),
         .iOS(.v13),
         .watchOS(.v6),
-        .tvOS(.v13),
+        .tvOS(.v13)
     ],
     products: [
-        .library(name: "WebSocketKit", targets: ["WebSocketKit"]),
+        .library(name: "WebSocketKit", targets: ["WebSocketKit"])
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.53.0"),
+        .package(url: "https://github.com/johnnzhou/swift-nio.git", branch: "main"),
         .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.16.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.24.0"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.16.0"),
-        .package(url: "https://github.com/apple/swift-atomics.git", from: "1.1.0"),
+        .package(url: "https://github.com/apple/swift-atomics.git", from: "1.1.0")
     ],
     targets: [
         .target(name: "WebSocketKit", dependencies: [
@@ -32,7 +32,7 @@ let package = Package(
             .product(name: "Atomics", package: "swift-atomics")
         ]),
         .testTarget(name: "WebSocketKitTests", dependencies: [
-            .target(name: "WebSocketKit"),
-        ]),
+            .target(name: "WebSocketKit")
+        ])
     ]
 )
diff --git a/Sources/WebSocketKit/BufferWritableMonitorDelegate.swift b/Sources/WebSocketKit/BufferWritableMonitorDelegate.swift
diff --git a/Sources/WebSocketKit/BufferWritableMonitorHandler.swift b/Sources/WebSocketKit/BufferWritableMonitorHandler.swift
diff --git a/Sources/WebSocketKit/WebSocket.swift b/Sources/WebSocketKit/WebSocket.swift
@@ -22,20 +22,24 @@ public final class WebSocket: Sendable {
     public var closeCode: WebSocketErrorCode? {
         _closeCode.withLockedValue { $0 }
     }
-    
+
     private let _closeCode: NIOLockedValueBox<WebSocketErrorCode?>
 
     public var onClose: EventLoopFuture<Void> {
         self.channel.closeFuture
     }
 
+    public var allocator: ByteBufferAllocator {
+        self.channel.allocator
+    }
+
     @usableFromInline
     /* private but @usableFromInline */
     internal let channel: Channel
-    private let onTextCallback: NIOLoopBoundBox<@Sendable (WebSocket, String) -> ()>
-    private let onBinaryCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> ()>
-    private let onPongCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> ()>
-    private let onPingCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> ()>
+    private let onTextCallback: NIOLoopBoundBox<@Sendable (WebSocket, String) -> Void>
+    private let onBinaryCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> Void>
+    private let onPongCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> Void>
+    private let onPingCallback: NIOLoopBoundBox<@Sendable (WebSocket, ByteBuffer) -> Void>
     private let type: PeerType
     private let waitingForPong: NIOLockedValueBox<Bool>
     private let waitingForClose: NIOLockedValueBox<Bool>
@@ -60,29 +64,29 @@ public final class WebSocket: Sendable {
         self._bufferedBytes = .init(0)
     }
 
-    @preconcurrency public func onText(_ callback: @Sendable @escaping (WebSocket, String) -> ()) {
+    @preconcurrency public func onText(_ callback: @Sendable @escaping (WebSocket, String) -> Void) {
         self.onTextCallback.value = callback
     }
 
-    @preconcurrency public func onBinary(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> ()) {
+    @preconcurrency public func onBinary(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> Void) {
         self.onBinaryCallback.value = callback
     }
-    
-    public func onPong(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> ()) {
+
+    public func onPong(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> Void) {
         self.onPongCallback.value = callback
     }
-    
+
     @available(*, deprecated, message: "Please use `onPong { socket, data in /* … */ }` with the additional `data` parameter.")
-    @preconcurrency public func onPong(_ callback: @Sendable @escaping (WebSocket) -> ()) {
+    @preconcurrency public func onPong(_ callback: @Sendable @escaping (WebSocket) -> Void) {
         self.onPongCallback.value = { ws, _ in callback(ws) }
     }
 
-    public func onPing(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> ()) {
+    public func onPing(_ callback: @Sendable @escaping (WebSocket, ByteBuffer) -> Void) {
         self.onPingCallback.value = callback
     }
-    
+
     @available(*, deprecated, message: "Please use `onPing { socket, data in /* … */ }` with the additional `data` parameter.")
-    @preconcurrency public func onPing(_ callback: @Sendable @escaping (WebSocket) -> ()) {
+    @preconcurrency public func onPing(_ callback: @Sendable @escaping (WebSocket) -> Void) {
         self.onPingCallback.value = { ws, _ in callback(ws) }
     }
 
@@ -111,8 +115,7 @@ public final class WebSocket: Sendable {
 
     @inlinable
     public func send<S>(_ text: S, promise: EventLoopPromise<Void>? = nil)
-        where S: Collection, S.Element == Character
-    {
+        where S: Collection, S.Element == Character {
         let string = String(text)
         let buffer = channel.allocator.buffer(string: string)
         self.send(buffer, opcode: .text, fin: true, promise: promise)
@@ -143,8 +146,7 @@ public final class WebSocket: Sendable {
         fin: Bool = true,
         promise: EventLoopPromise<Void>? = nil
     )
-        where Data: DataProtocol
-    {
+        where Data: DataProtocol {
         if let byteBufferView = data as? ByteBufferView {
             // optimisation: converting from `ByteBufferView` to `ByteBuffer` doesn't allocate or copy any data
             send(ByteBuffer(byteBufferView), opcode: opcode, fin: fin, promise: promise)
@@ -342,13 +344,9 @@ public final class WebSocket: Sendable {
     }
 }
 
-extension WebSocket: BufferWritableMonitorDelegate {
-    func onBufferWritableChanged(amountQueued: Int) {
-        self._bufferedBytes.withLockedValue { $0 = amountQueued }
-    }
-    
-    public var bufferedBytes: Int {
-        self._bufferedBytes.withLockedValue { $0 }
+extension WebSocket {
+    public func getBufferedBytes() -> EventLoopFuture<Int> {
+        return channel.getOption(.bufferedWritableBytes)
     }
 }
 
diff --git a/Sources/WebSocketKit/WebSocketClient.swift b/Sources/WebSocketKit/WebSocketClient.swift
@@ -23,7 +23,7 @@ public final class WebSocketClient: Sendable {
 
     public struct Configuration: Sendable {
         public var tlsConfiguration: TLSConfiguration?
-        public var maxFrameSize: Int
+        public var maxFrameSize: Int = 1 << 24
 
         /// Defends against small payloads in frame aggregation.
         /// See `NIOWebSocketFrameAggregator` for details.
@@ -37,7 +37,7 @@ public final class WebSocketClient: Sendable {
 
         public init(
             tlsConfiguration: TLSConfiguration? = nil,
-            maxFrameSize: Int = 1 << 14
+            maxFrameSize: Int = 1 << 24
         ) {
             self.tlsConfiguration = tlsConfiguration
             self.maxFrameSize = maxFrameSize
@@ -71,11 +71,11 @@ public final class WebSocketClient: Sendable {
         path: String = "/",
         query: String? = nil,
         headers: HTTPHeaders = [:],
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         self.connect(scheme: scheme, host: host, port: port, path: path, query: query, headers: headers, proxy: nil, onUpgrade: onUpgrade)
     }
-    
+
     @preconcurrency
     public func connect(
         scheme: String,
@@ -85,11 +85,11 @@ public final class WebSocketClient: Sendable {
         query: String? = nil,
         headers: HTTPHeaders = [:],
         maxQueueSize: Int? = nil,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         self.connect(scheme: scheme, host: host, port: port, path: path, query: query, maxQueueSize: maxQueueSize, headers: headers, proxy: nil, onUpgrade: onUpgrade)
     }
-    
+
     /// Establish a WebSocket connection via a proxy server.
     ///
     /// - Parameters:
@@ -118,7 +118,7 @@ public final class WebSocketClient: Sendable {
         proxyPort: Int? = nil,
         proxyHeaders: HTTPHeaders = [:],
         proxyConnectDeadline: NIODeadline = NIODeadline.distantFuture,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         assert(["ws", "wss"].contains(scheme))
         let upgradePromise = self.group.any().makePromise(of: Void.self)
@@ -152,14 +152,14 @@ public final class WebSocketClient: Sendable {
                 let websocketUpgrader = NIOWebSocketClientUpgrader(
                     maxFrameSize: self.configuration.maxFrameSize,
                     automaticErrorHandling: true,
-                    upgradePipelineHandler: { channel, req in
+                    upgradePipelineHandler: { channel, _ in
                         return WebSocket.client(on: channel, maxQueueSize: maxQueueSize, config: .init(clientConfig: self.configuration), onUpgrade: onUpgrade)
                     }
                 )
 
                 let config: NIOHTTPClientUpgradeConfiguration = (
                     upgraders: [websocketUpgrader],
-                    completionHandler: { context in
+                    completionHandler: { _ in
                         upgradePromise.succeed(())
                         channel.pipeline.removeHandler(httpUpgradeRequestHandlerBox.value, promise: nil)
                     }
@@ -254,16 +254,24 @@ public final class WebSocketClient: Sendable {
 
         let connect = bootstrap.connect(host: proxy ?? host, port: proxyPort ?? port)
         connect.cascadeFailure(to: upgradePromise)
-        return connect.flatMap { channel in
+        return connect.flatMap { _ in
             return upgradePromise.futureResult
         }
     }
 
     @Sendable
     private func makeTLSHandler(tlsConfiguration: TLSConfiguration?, host: String) throws -> NIOSSLClientHandler {
+        var tlsConfig: TLSConfiguration
+        if self.configuration.tlsConfiguration == nil {
+            tlsConfig = .makeClientConfiguration()
+            tlsConfig.certificateVerification = .none
+        } else {
+            tlsConfig = self.configuration.tlsConfiguration!
+        }
         let context = try NIOSSLContext(
-            configuration: self.configuration.tlsConfiguration ?? .makeClientConfiguration()
+            configuration: tlsConfig
         )
+
         let tlsHandler: NIOSSLClientHandler
         do {
             tlsHandler = try NIOSSLClientHandler(context: context, serverHostname: host)
@@ -289,7 +297,7 @@ public final class WebSocketClient: Sendable {
             }
         }
     }
-    
+
     private static func makeBootstrap(on eventLoop: EventLoopGroup) -> NIOClientTCPBootstrapProtocol {
         #if canImport(Network)
         if let tsBootstrap = NIOTSConnectionBootstrap(validatingGroup: eventLoop) {
diff --git a/Sources/WebSocketKit/WebSocketHandler.swift b/Sources/WebSocketKit/WebSocketHandler.swift
@@ -36,7 +36,7 @@ extension WebSocket {
     @preconcurrency
     public static func client(
         on channel: Channel,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         return self.configure(on: channel, as: .client, with: Configuration(), onUpgrade: onUpgrade)
     }
@@ -52,7 +52,7 @@ extension WebSocket {
         on channel: Channel,
         maxQueueSize: Int? = nil,
         config: Configuration,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         return self.configure(on: channel, as: .client, maxQueueSize: maxQueueSize, with: config, onUpgrade: onUpgrade)
     }
@@ -65,7 +65,7 @@ extension WebSocket {
     @preconcurrency
     public static func server(
         on channel: Channel,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         return self.configure(on: channel, as: .server, with: Configuration(), onUpgrade: onUpgrade)
     }
@@ -80,7 +80,7 @@ extension WebSocket {
     public static func server(
         on channel: Channel,
         config: Configuration,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         return self.configure(on: channel, as: .server, with: config, onUpgrade: onUpgrade)
     }
@@ -90,7 +90,7 @@ extension WebSocket {
         as type: PeerType,
         maxQueueSize: Int? = nil,
         with config: Configuration,
-        onUpgrade: @Sendable @escaping (WebSocket) -> ()
+        onUpgrade: @Sendable @escaping (WebSocket) -> Void
     ) -> EventLoopFuture<Void> {
         let webSocket = WebSocket(channel: channel, type: type)
 
@@ -102,19 +102,10 @@ extension WebSocket {
             ),
             WebSocketHandler(webSocket: webSocket)
         ]).map { _ in
-            if let maxQueueSize = maxQueueSize {
-                channel.pipeline.addHandler(BufferWritableMonitorHandler(capacity: maxQueueSize, delegate: webSocket), position: .first).whenSuccess { _ in
-                    #if DEBUG
-                    print(channel.pipeline.debugDescription)
-                    #endif
-                    onUpgrade(webSocket)
-                }
-            } else {
-                #if DEBUG
-                print(channel.pipeline.debugDescription)
-                #endif
-                onUpgrade(webSocket)
-            }
+             #if DEBUG
+             print(channel.pipeline.debugDescription)
+             #endif
+            onUpgrade(webSocket)
         }.flatMapError { error in
             return channel.eventLoop.makeFailedFuture(error)
         }
