Skip to content

Expose upgrade response headers to WebSocket client onUpgrade #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
awaisdev5765 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Expose upgrade response headers to WebSocket client onUpgrade #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 96 additions & 0 deletions Sources/WebSocketKit/Concurrency/WebSocket+Concurrency.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,4 +173,100 @@ extension WebSocket {
}
).get()
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - to: URL for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
@preconcurrency public static func connect(
to url: String,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) async -> ()
) async throws {
return try await self.connect(
to: url,
headers: headers,
configuration: configuration,
on: eventLoopGroup,
onUpgradeWithResponse: { ws, head in
Task {
await onUpgradeWithResponse(ws, head)
}
}
).get()
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - to: URL for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
@preconcurrency public static func connect(
to url: URL,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) async -> ()
) async throws {
return try await self.connect(
to: url,
headers: headers,
configuration: configuration,
on: eventLoopGroup,
onUpgradeWithResponse: { ws, head in
Task {
await onUpgradeWithResponse(ws, head)
}
}
).get()
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - scheme: Scheme component of the URI for the WebSocket server.
/// - host: Host component of the URI for the WebSocket server.
/// - port: Port on which to connect to the WebSocket server.
/// - path: Path component of the URI for the WebSocket server.
/// - query: Query component of the URI for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
@preconcurrency public static func connect(
scheme: String = "ws",
host: String,
port: Int = 80,
path: String = "/",
query: String? = nil,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) async -> ()
) async throws {
return try await self.connect(
scheme: scheme,
host: host,
port: port,
path: path,
query: query,
headers: headers,
configuration: configuration,
on: eventLoopGroup,
onUpgradeWithResponse: { ws, head in
Task {
await onUpgradeWithResponse(ws, head)
}
}
).get()
}
}
156 changes: 153 additions & 3 deletions Sources/WebSocketKit/WebSocket+Connect.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,35 @@ extension WebSocket {
)
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - url: URL for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgrade: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the WebSocket server is established.
@preconcurrency
public static func connect(
to url: String,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
guard let url = URL(string: url) else {
return eventLoopGroup.any().makeFailedFuture(WebSocketClient.Error.invalidURL)
}
return self.connect(
to: url,
headers: headers,
configuration: configuration,
on: eventLoopGroup,
onUpgradeWithResponse: onUpgradeWithResponse
)
}

/// Establish a WebSocket connection.
///
/// - Parameters:
Expand Down Expand Up @@ -63,6 +92,37 @@ extension WebSocket {
)
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - url: URL for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the WebSocket server is established.
@preconcurrency
public static func connect(
to url: URL,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
let scheme = url.scheme ?? "ws"
return self.connect(
scheme: scheme,
host: url.host ?? "localhost",
port: url.port ?? (scheme == "wss" ? 443 : 80),
path: url.path,
query: url.query,
headers: headers,
configuration: configuration,
on: eventLoopGroup,
onUpgradeWithResponse: onUpgradeWithResponse
)
}

/// Establish a WebSocket connection.
///
/// - Parameters:
Expand Down Expand Up @@ -102,6 +162,45 @@ extension WebSocket {
)
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - scheme: Scheme component of the URI for the WebSocket server.
/// - host: Host component of the URI for the WebSocket server.
/// - port: Port on which to connect to the WebSocket server.
/// - path: Path component of the URI for the WebSocket server.
/// - query: Query component of the URI for the WebSocket server.
/// - headers: Headers to send to the WebSocket server.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the WebSocket server is established.
@preconcurrency
public static func connect(
scheme: String = "ws",
host: String,
port: Int = 80,
path: String = "/",
query: String? = nil,
headers: HTTPHeaders = [:],
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
return WebSocketClient(
eventLoopGroupProvider: .shared(eventLoopGroup),
configuration: configuration
).connect(
scheme: scheme,
host: host,
port: port,
path: path,
query: query,
headers: headers,
onUpgradeWithResponse: onUpgradeWithResponse
)
}

/// Establish a WebSocket connection via a proxy server.
///
/// - Parameters:
Expand Down Expand Up @@ -153,8 +252,59 @@ extension WebSocket {
)
}

/// Establish a WebSocket connection via a proxy server, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - scheme: Scheme component of the URI for the origin server.
/// - host: Host component of the URI for the origin server.
/// - port: Port on which to connect to the origin server.
/// - path: Path component of the URI for the origin server.
/// - query: Query component of the URI for the origin server.
/// - headers: Headers to send to the origin server.
/// - proxy: Host component of the URI for the proxy server.
/// - proxyPort: Port on which to connect to the proxy server.
/// - proxyHeaders: Headers to send to the proxy server.
/// - proxyConnectDeadline: Deadline for establishing the proxy connection.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the origin server is established.
@preconcurrency
public static func connect(
scheme: String = "ws",
host: String,
port: Int = 80,
path: String = "/",
query: String? = nil,
headers: HTTPHeaders = [:],
proxy: String?,
proxyPort: Int? = nil,
proxyHeaders: HTTPHeaders = [:],
proxyConnectDeadline: NIODeadline = NIODeadline.distantFuture,
configuration: WebSocketClient.Configuration = .init(),
on eventLoopGroup: any EventLoopGroup,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
return WebSocketClient(
eventLoopGroupProvider: .shared(eventLoopGroup),
configuration: configuration
).connect(
scheme: scheme,
host: host,
port: port,
path: path,
query: query,
headers: headers,
proxy: proxy,
proxyPort: proxyPort,
proxyHeaders: proxyHeaders,
proxyConnectDeadline: proxyConnectDeadline,
onUpgradeWithResponse: onUpgradeWithResponse
)
}

/// Description
/// Establish a WebSocket connection via a proxy server, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - url: URL for the origin server.
/// - headers: Headers to send to the origin server.
Expand All @@ -164,10 +314,10 @@ extension WebSocket {
/// - proxyConnectDeadline: Deadline for establishing the proxy connection.
/// - configuration: Configuration for the WebSocket client.
/// - eventLoopGroup: Event loop group to be used by the WebSocket client.
/// - onUpgrade: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the origin server is established.
@preconcurrency
public static func connect(
public static func connect(
to url: String,
headers: HTTPHeaders = [:],
proxy: String?,
Expand Down
70 changes: 69 additions & 1 deletion Sources/WebSocketKit/WebSocketClient.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,33 @@ public final class WebSocketClient: Sendable {
self.connect(scheme: scheme, host: host, port: port, path: path, query: query, headers: headers, proxy: nil, onUpgrade: onUpgrade)
}

/// Establish a WebSocket connection, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - scheme: Scheme component of the URI for the origin server.
/// - host: Host component of the URI for the origin server.
/// - port: Port on which to connect to the origin server.
/// - path: Path component of the URI for the origin server.
/// - query: Query component of the URI for the origin server.
/// - headers: Headers to send to the origin server.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the origin server is established.
@preconcurrency
public func connect(
scheme: String,
host: String,
port: Int,
path: String = "/",
query: String? = nil,
headers: HTTPHeaders = [:],
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
self.connect(
scheme: scheme, host: host, port: port, path: path, query: query, headers: headers, proxy: nil,
onUpgradeWithResponse: onUpgradeWithResponse
)
}

/// Establish a WebSocket connection via a proxy server.
///
/// - Parameters:
Expand Down Expand Up @@ -104,6 +131,43 @@ public final class WebSocketClient: Sendable {
proxyHeaders: HTTPHeaders = [:],
proxyConnectDeadline: NIODeadline = NIODeadline.distantFuture,
onUpgrade: @Sendable @escaping (WebSocket) -> ()
) -> EventLoopFuture<Void> {
self.connect(
scheme: scheme, host: host, port: port, path: path, query: query, headers: headers,
proxy: proxy, proxyPort: proxyPort, proxyHeaders: proxyHeaders,
proxyConnectDeadline: proxyConnectDeadline,
onUpgradeWithResponse: { ws, _ in onUpgrade(ws) }
)
}

/// Establish a WebSocket connection via a proxy server, receiving the server's upgrade response headers.
///
/// - Parameters:
/// - scheme: Scheme component of the URI for the origin server.
/// - host: Host component of the URI for the origin server.
/// - port: Port on which to connect to the origin server.
/// - path: Path component of the URI for the origin server.
/// - query: Query component of the URI for the origin server.
/// - headers: Headers to send to the origin server.
/// - proxy: Host component of the URI for the proxy server.
/// - proxyPort: Port on which to connect to the proxy server.
/// - proxyHeaders: Headers to send to the proxy server.
/// - proxyConnectDeadline: Deadline for establishing the proxy connection.
/// - onUpgradeWithResponse: An escaping closure to be executed after the upgrade is completed by `NIOWebSocketClientUpgrader`, invoked with the 101 Switching Protocols response head.
/// - Returns: A future which completes when the connection to the origin server is established.
@preconcurrency
public func connect(
scheme: String,
host: String,
port: Int,
path: String = "/",
query: String? = nil,
headers: HTTPHeaders = [:],
proxy: String?,
proxyPort: Int? = nil,
proxyHeaders: HTTPHeaders = [:],
proxyConnectDeadline: NIODeadline = NIODeadline.distantFuture,
onUpgradeWithResponse: @Sendable @escaping (WebSocket, HTTPResponseHead) -> ()
) -> EventLoopFuture<Void> {
assert(["ws", "wss"].contains(scheme))
let upgradePromise = self.group.any().makePromise(of: Void.self)
Expand Down Expand Up @@ -140,7 +204,11 @@ public final class WebSocketClient: Sendable {
maxFrameSize: self.configuration.maxFrameSize,
automaticErrorHandling: true,
upgradePipelineHandler: { channel, req in
return WebSocket.client(on: channel, config: .init(clientConfig: self.configuration), onUpgrade: onUpgrade)
return WebSocket.client(
on: channel,
config: .init(clientConfig: self.configuration),
onUpgrade: { ws in onUpgradeWithResponse(ws, req) }
)
}
)

Expand Down
Loading