Return close code from WebSocketClient.run/connect (#58)

hummingbird-project · Apr 18, 2024 · 1b9bdfc · 1b9bdfc
1 parent 505d324
commit 1b9bdfc
Show file tree

Hide file tree

Showing 10 changed files with 261 additions and 111 deletions.
diff --git a/Sources/HummingbirdWSClient/Client/ClientChannel.swift b/Sources/HummingbirdWSClient/Client/ClientChannel.swift
@@ -19,6 +19,7 @@ import NIOCore
 @_documentation(visibility: internal)
 public protocol ClientConnectionChannel: Sendable {
     associatedtype Value: Sendable
+    associatedtype Result
 
     /// Setup child channel
     /// - Parameters:
@@ -31,5 +32,5 @@ public protocol ClientConnectionChannel: Sendable {
     /// - Parameters:
     ///   - value: Object to process input/output on child channel
     ///   - logger: Logger to use while processing messages
-    func handle(value: Value, logger: Logger) async throws
+    func handle(value: Value, logger: Logger) async throws -> Result
 }
diff --git a/Sources/HummingbirdWSClient/Client/ClientConnection.swift b/Sources/HummingbirdWSClient/Client/ClientConnection.swift
@@ -19,6 +19,7 @@ import NIOPosix
 import Network
 import NIOTransportServices
 #endif
+import NIOWebSocket
 
 /// A generic client connection to a server.
 ///
@@ -86,12 +87,12 @@ public struct ClientConnection<ClientChannel: ClientConnectionChannel>: Sendable
     }
     #endif
 
-    public func run() async throws {
+    public func run() async throws -> ClientChannel.Result {
         let channelResult = try await self.makeClient(
             clientChannel: self.clientChannel,
             address: self.address
         )
-        try await self.clientChannel.handle(value: channelResult, logger: self.logger)
+        return try await self.clientChannel.handle(value: channelResult, logger: self.logger)
     }
 
     /// Connect to server

diff --git a/Sources/HummingbirdWSClient/Client/TLSClientChannel.swift b/Sources/HummingbirdWSClient/Client/TLSClientChannel.swift
@@ -15,11 +15,13 @@
 import Logging
 import NIOCore
 import NIOSSL
+import NIOWebSocket
 
 /// Sets up client channel to use TLS before accessing base channel setup
 @_documentation(visibility: internal)
 public struct TLSClientChannel<BaseChannel: ClientConnectionChannel>: ClientConnectionChannel {
     public typealias Value = BaseChannel.Value
+    public typealias Result = BaseChannel.Result
 
     ///  Initialize TLSChannel
     /// - Parameters:
@@ -51,7 +53,7 @@ public struct TLSClientChannel<BaseChannel: ClientConnectionChannel>: ClientConn
     /// - Parameters:
     ///   - value: Object to process input/output on child channel
     ///   - logger: Logger to use while processing messages
-    public func handle(value: BaseChannel.Value, logger: Logging.Logger) async throws {
+    public func handle(value: BaseChannel.Value, logger: Logging.Logger) async throws -> Result {
         try await self.baseChannel.handle(value: value, logger: logger)
     }
 

diff --git a/Sources/HummingbirdWSClient/WebSocketClient.swift b/Sources/HummingbirdWSClient/WebSocketClient.swift
@@ -19,6 +19,7 @@ import NIOCore
 import NIOPosix
 import NIOSSL
 import NIOTransportServices
+import NIOWebSocket
 
 /// WebSocket client
 ///
@@ -111,8 +112,9 @@ public struct WebSocketClient {
     }
     #endif
 
-    ///  Connect and run handler
-    public func run() async throws {
+    /// Connect and run handler
+    /// - Returns: WebSocket close frame details if server returned any
+    @discardableResult public func run() async throws -> WebSocketCloseFrame? {
         guard let host = url.host else { throw WebSocketClientError.invalidURL }
         let requiresTLS = self.url.scheme == .wss || self.url.scheme == .https
         let port = self.url.port ?? (requiresTLS ? 443 : 80)
@@ -129,7 +131,7 @@ public struct WebSocketClient {
                     eventLoopGroup: self.eventLoopGroup,
                     logger: self.logger
                 )
-                try await client.run()
+                return try await client.run()
 
             #if canImport(Network)
             case .ts(let tlsOptions):
@@ -140,7 +142,7 @@ public struct WebSocketClient {
                     eventLoopGroup: self.eventLoopGroup,
                     logger: self.logger
                 )
-                try await client.run()
+                return try await client.run()
 
             #endif
             case .none:
@@ -158,7 +160,7 @@ public struct WebSocketClient {
                     eventLoopGroup: self.eventLoopGroup,
                     logger: self.logger
                 )
-                try await client.run()
+                return try await client.run()
             }
         } else {
             let client = try ClientConnection(
@@ -171,7 +173,7 @@ public struct WebSocketClient {
                 eventLoopGroup: self.eventLoopGroup,
                 logger: self.logger
             )
-            try await client.run()
+            return try await client.run()
         }
     }
 }
@@ -186,14 +188,15 @@ extension WebSocketClient {
     ///   - eventLoopGroup: EventLoopGroup to run WebSocket client on
     ///   - logger: Logger
     ///   - process: Closure handling webSocket
-    public static func connect(
+    /// - Returns: WebSocket close frame details if server returned any
+    @discardableResult public static func connect(
         url: String,
         configuration: WebSocketClientConfiguration = .init(),
         tlsConfiguration: TLSConfiguration? = nil,
         eventLoopGroup: EventLoopGroup = MultiThreadedEventLoopGroup.singleton,
         logger: Logger,
         handler: @escaping WebSocketDataHandler<BasicWebSocketContext>
-    ) async throws {
+    ) async throws -> WebSocketCloseFrame? {
         let ws = self.init(
             url: url,
             configuration: configuration,
@@ -202,7 +205,7 @@ extension WebSocketClient {
             logger: logger,
             handler: handler
         )
-        try await ws.run()
+        return try await ws.run()
     }
 
     #if canImport(Network)
@@ -215,14 +218,15 @@ extension WebSocketClient {
     ///   - eventLoopGroup: EventLoopGroup to run WebSocket client on
     ///   - logger: Logger
     ///   - process: WebSocket data handler
+    /// - Returns: WebSocket close frame details if server returned any
     public static func connect(
         url: String,
         configuration: WebSocketClientConfiguration = .init(),
         transportServicesTLSOptions: TSTLSOptions,
         eventLoopGroup: NIOTSEventLoopGroup = NIOTSEventLoopGroup.singleton,
         logger: Logger,
         handler: @escaping WebSocketDataHandler<BasicWebSocketContext>
-    ) async throws {
+    ) async throws -> WebSocketCloseFrame? {
         let ws = self.init(
             url: url,
             configuration: configuration,
@@ -231,7 +235,7 @@ extension WebSocketClient {
             logger: logger,
             handler: handler
         )
-        try await ws.run()
+        return try await ws.run()
     }
     #endif
 }
diff --git a/Sources/HummingbirdWSClient/WebSocketClientChannel.swift b/Sources/HummingbirdWSClient/WebSocketClientChannel.swift
@@ -94,10 +94,10 @@ struct WebSocketClientChannel: ClientConnectionChannel {
         }
     }
 
-    func handle(value: Value, logger: Logger) async throws {
+    func handle(value: Value, logger: Logger) async throws -> WebSocketCloseFrame? {
         switch try await value.get() {
         case .websocket(let webSocketChannel, let extensions):
-            await WebSocketHandler.handle(
+            return try await WebSocketHandler.handle(
                 type: .client,
                 configuration: .init(
                     extensions: extensions,