Add support for passing IO objects to Websockets.open() (#1251)

* Add support for passing IO objects to Websockets.open()

This is handy if the caller already has the required stream/socket open.

* fixup! Add support for passing IO objects to Websockets.open()

Author: JamesWrigley

docs/src/guides/protocols.md

@@ -50,6 +50,17 @@ proxy-aware clients without forcing you through internal parser state. The
 [WebSockets API reference](../api/websockets.md) is the canonical home for the
 public docstrings.
 
+`HTTP.WebSockets.open` can also handshake over an already-connected `IO` instead
+of a URL. Pass any byte stream (a raw `TCPSocket`, a TLS stream, etc.) as the
+first argument, and optionally override the request-line `target`/`Host` header.
+This is handy when the transport is established out-of-band or for tunnelling
+WebSockets over a custom stream. The caller retains ownership of the `IO`, `open`
+never closes it.
+
+```julia
+ws = HTTP.WebSockets.open(io; target = "/echo", host = "example.com:80")
+```
+
 `HTTP.WebSockets.open` also accepts the client-side handshake timeout controls:
 
 - `connect_timeout`

src/http_core.jl

@@ -1056,6 +1056,27 @@ function defaultheader!(headers::Headers, item::Pair)
     return setheader(headers, item)
 end
 
+"""
+Read at least one byte into `b` (up to `nb`), blocking until data is available
+or EOF, and return the number of bytes read (`0` only at EOF).
+
+The WebSocket read loop relies on Reseau's "block until at least one byte or
+EOF" semantics, which `readbytes!(conn; all=false)` provides for Reseau conns.
+A stdlib `TCPSocket` with `all=false` instead returns `0` whenever no bytes are
+buffered *yet*, so for any other `IO` we reimplement the blocking contract on
+top of `eof`/`bytesavailable`. This lets `WebSockets.open(io)` accept an
+arbitrary caller-provided stream that bypasses the connection pool.
+"""
+function _blocking_readbytes!(conn::Union{TCP.Conn,TLS.Conn}, b::AbstractVector{UInt8}, nb::Integer=length(b))
+    readbytes!(conn, b, nb; all=false)
+end
+
+function _blocking_readbytes!(io::IO, b::AbstractVector{UInt8}, nb::Integer=length(b))
+    eof(io) && return 0
+    n = min(bytesavailable(io), Int(nb))
+    return readbytes!(io, b, n)
+end
+
 """
 Internal buffered reader that first drains already-read bytes before continuing
 from the underlying TCP or TLS connection.

src/http_websockets.jl

@@ -31,6 +31,7 @@ import ..CookieJar
 import ..COOKIEJAR
 import ..Cookies
 import .._ConnReader
+import .._blocking_readbytes!
 import .._USE_TRANSPORT_PROXY
 import .._close_conn!
 import .._client_for_request
@@ -442,8 +443,8 @@ function _ws_read_loop!(ws::WebSocket, buffer_bytes::Int=DEFAULT_READ_BUFFER_BYT
             # Read directly into the reusable `buf`. `readavailable` would
             # allocate a fresh Base.SZ_UNBUFFERED_IO (64KB) buffer on every
             # frame read (16× the bytes/RTT vs HTTP 1.x), driving GC pressure;
-            # `readbytes!(...; all=false)` does one socket read into `buf`.
-            n = readbytes!(ws.stream, buf, length(buf); all=false)
+            # `_blocking_readbytes!` does one blocking socket read into `buf`.
+            n = _blocking_readbytes!(ws.stream, buf, length(buf))
             n == 0 && break
             ws_on_incoming_data!(on_frame, ws.codec, buf, n)
             _flush_ws_output!(ws)
@@ -495,6 +496,42 @@ function _start_read_task!(ws::WebSocket, buffer_bytes::Int=DEFAULT_READ_BUFFER_
     return nothing
 end
 
+# Shared client-side handshake completion for both `open` forms: build the
+# client `WebSocket` over an already-established `stream`, record the handshake
+# messages, replay any early frames the server already sent (`buffered`), and
+# start the read task.
+function _finish_client_websocket(
+    stream,
+    close_transport!,
+    request::Request,
+    response::Response,
+    negotiated::Union{Nothing,String};
+    maxframesize::Integer,
+    maxfragmentation::Integer,
+    read_idle_timeout_ns::Integer=0,
+    buffered::Vector{UInt8}=UInt8[],
+    pmce::Union{Nothing,PMCEContext}=nothing,
+)::WebSocket
+    ws = WebSocket(
+        stream,
+        close_transport!;
+        subprotocol=negotiated,
+        maxframesize=maxframesize,
+        maxfragmentation=maxfragmentation,
+        read_idle_timeout_ns=read_idle_timeout_ns,
+        is_client=true,
+        pmce=pmce,
+    )
+    ws.handshake_request = request
+    ws.handshake_response = response
+    if !isempty(buffered)
+        ws_on_incoming_data!(frame -> _process_incoming_frame!(ws, frame), ws.codec, buffered)
+        _flush_ws_output!(ws)
+    end
+    _start_read_task!(ws)
+    return ws
+end
+
 """
     send(ws, message) -> Int
 
@@ -885,24 +922,18 @@ function _open_client_websocket(
                     return nothing
                 end
             end
-            ws = WebSocket(
+            return _finish_client_websocket(
                 _conn_stream(conn),
                 close_transport!,
-                subprotocol=negotiated,
+                send_request,
+                response,
+                negotiated;
                 maxframesize=maxframesize,
                 maxfragmentation=maxfragmentation,
                 read_idle_timeout_ns=_request_read_idle_timeout_ns(send_request),
-                is_client=true,
+                buffered=attempt.buffered,
                 pmce=pmce_ctx,
             )
-            ws.handshake_request = send_request
-            ws.handshake_response = response
-            if !isempty(attempt.buffered)
-                ws_on_incoming_data!(frame -> _process_incoming_frame!(ws, frame), ws.codec, attempt.buffered)
-                _flush_ws_output!(ws)
-            end
-            _start_read_task!(ws)
-            return ws
         end
         if !_is_redirect_status(response.status) || redirect_policy.max_redirects == 0
             owns_client && close(req_client)
@@ -1044,6 +1075,98 @@ function open(
     end
 end
 
+# Perform the websocket client handshake directly over a caller-provided `io`,
+# bypassing the connection pool and TLS/proxy machinery. The handshake request is
+# serialized and written to `io`, then the 101 response is parsed back off it. The
+# returned WebSocket's `close_transport!` is a no-op: the caller retains ownership
+# of `io` and is responsible for closing it.
+function _open_client_websocket_io(
+    io::IO;
+    target::AbstractString="/",
+    host::AbstractString="",
+    headers=Pair{String,String}[],
+    maxframesize::Integer=typemax(Int),
+    maxfragmentation::Integer=DEFAULT_MAX_FRAG,
+    subprotocols::AbstractVector{<:AbstractString}=String[],
+)::WebSocket
+    req_headers = _normalize_headers_input(headers)
+    isempty(host) || setheader(req_headers, "Host", host)
+    key = ws_random_handshake_key()
+    _apply_websocket_request_headers!(req_headers, key, subprotocols)
+    request = Request("GET", String(target); headers=req_headers, host=String(host), body=EmptyBody(), content_length=0)
+    expected_accept = ws_compute_accept_key(header(request.headers, "Sec-WebSocket-Key")::String)
+
+    # Serialize and write the handshake request to the caller's transport.
+    request_buf = IOBuffer()
+    write_request!(request_buf, request)
+    write(io, take!(request_buf))
+
+    incoming = _read_incoming_response(io, request)
+    @try_ignore body_close!(incoming.rawbody)
+    response = _streaming_response(incoming)
+    negotiated = _validate_websocket_upgrade!(response, expected_accept, subprotocols)
+
+    # No-op `close_transport!`: the caller retains ownership of `io`.
+    return _finish_client_websocket(
+        io,
+        () -> nothing,
+        request,
+        response,
+        negotiated;
+        maxframesize=maxframesize,
+        maxfragmentation=maxfragmentation,
+    )
+end
+
+"""
+    open(io::IO; target="/", host="", kwargs...) -> WebSocket
+    open(f, io::IO; target="/", host="", kwargs...) -> Any
+
+Perform the WebSocket client handshake directly over an already-connected `io`
+(a raw `TCPSocket`, a TLS stream, or any other byte `IO`) instead of dialing a
+new connection from a URL. This is useful when the transport is established
+out-of-band, for testing, or for tunnelling WebSockets over a custom stream.
+
+Because there is no URL to derive them from, the request-line `target` and `Host`
+header default to `"/"` and unset; override them with the `target` and `host`
+keyword arguments. The remaining `headers`, `subprotocols`, `maxframesize`, and
+`maxfragmentation` keywords match the URL-based [`open`](@ref); pool, TLS, proxy,
+redirect, cookie, and timeout options do not apply.
+
+`io` is **not** closed by `open` — the caller retains ownership of the underlying
+stream's lifetime. As with the URL form, calling `open` with a function closes
+the WebSocket with status code `1000` when `f` returns (the close frame is sent
+over `io`, but `io` itself stays open).
+"""
+open(io::IO; kwargs...) = _open_client_websocket_io(io; kwargs...)
+
+function open(
+    f::Function,
+    io::IO;
+    suppress_close_error::Bool=false,
+    kwargs...,
+)
+    ws = open(io; kwargs...)
+    try
+        return f(ws)
+    catch err
+        if err isa WebSocketError && isok(err)
+            return nothing
+        end
+        rethrow(err)
+    finally
+        if !isclosed(ws)
+            try
+                close(ws, CloseFrameBody(1000, ""))
+            catch err
+                if !(suppress_close_error && err isa WebSocketError)
+                    rethrow(err)
+                end
+            end
+        end
+    end
+end
+
 """
     Server(; network="tcp", address="127.0.0.1:0", handler, tls_config=nothing, ...)
 

test/Project.toml

@@ -5,4 +5,5 @@ HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reseau = "802f3686-a58f-41ce-bb0c-3c43c75bba36"
+Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/http_websocket_raw_io_tests.jl

@@ -0,0 +1,114 @@
+using Test
+using HTTP
+using Reseau
+using Sockets
+
+const W = HTTP.WebSockets
+
+@testset "HTTP.WebSockets.open over a raw IO" begin
+    # Echo server that also surfaces the request target the client handshook with.
+    server = W.listen!("127.0.0.1", 0) do ws
+        W.send(ws, "target:" * ws.handshake_request.target)
+        for msg in ws
+            W.send(ws, "echo: " * (msg isa String ? msg : String(msg)))
+        end
+    end
+
+    try
+        host, port = Reseau.HostResolvers.split_host_port(W.server_addr(server))
+        port = parse(Int, port)
+
+        @testset "echo round-trip over a connected TCPSocket" begin
+            sock = Sockets.connect(host, port)
+            got = String[]
+            W.open(sock; host="$host:$port") do ws
+                push!(got, W.receive(ws))
+                W.send(ws, "hello")
+                push!(got, W.receive(ws))
+                W.send(ws, "world")
+                push!(got, W.receive(ws))
+            end
+            @test got == ["target:/", "echo: hello", "echo: world"]
+            # open() must not close the caller-owned transport: only the caller
+            # may close `sock`, so it is still locally open here.
+            @test isopen(sock)
+            close(sock)
+        end
+
+        @testset "custom target keyword is sent in the request line" begin
+            sock = Sockets.connect(host, port)
+            target = nothing
+            W.open(sock; target="/ws/v1", host="$host:$port") do ws
+                target = W.receive(ws)
+            end
+            @test target == "target:/ws/v1"
+            close(sock)
+        end
+
+        @testset "non-function form returns a usable WebSocket" begin
+            sock = Sockets.connect(host, port)
+            ws = W.open(sock; host="$host:$port")
+            try
+                @test W.receive(ws) == "target:/"
+                W.send(ws, "hi")
+                @test W.receive(ws) == "echo: hi"
+            finally
+                close(ws)
+            end
+            @test isopen(sock)
+            close(sock)
+        end
+
+        @testset "round-trip over a Base.Pipe forwarded to the server" begin
+            # `open` runs the handshake over a duplex `Base.PipeEndpoint` (a
+            # non-socket `Base.IO` whose `all=false` reads return 0 before data
+            # arrives) rather than over the TCP socket directly. A pair of tasks
+            # forwards bytes between the pipe's far end and a normal TCP
+            # connection to the echo server above, so the whole exchange is
+            # driven over a pipe.
+            # A named pipe gives us a full-duplex `Base.PipeEndpoint`. The
+            # address format differs by platform: Windows uses a `\\.\pipe\...`
+            # name (not a filesystem path), other platforms a Unix-domain socket.
+            pipename = "http_jl_ws_pipe_$(getpid())"
+            sockpath = Sys.iswindows() ? "\\\\.\\pipe\\$pipename" : joinpath(mktempdir(), "$pipename.sock")
+            relay = Sockets.listen(sockpath)
+            client_pipe = Sockets.connect(sockpath)   # handed to open()
+            bridge = Sockets.accept(relay)            # forwarded to the server
+            close(relay)
+            tcp = Sockets.connect(host, port)
+
+            fwd_up = Threads.@spawn try
+                while !eof(bridge)
+                    write(tcp, readavailable(bridge))
+                end
+            finally
+                close(tcp)
+            end
+            fwd_down = Threads.@spawn try
+                while !eof(tcp)
+                    write(bridge, readavailable(tcp))
+                end
+            finally
+                close(bridge)
+            end
+
+            try
+                @test client_pipe isa Base.PipeEndpoint
+                got = String[]
+                W.open(client_pipe) do ws
+                    push!(got, W.receive(ws))
+                    W.send(ws, "hello")
+                    push!(got, W.receive(ws))
+                end
+
+                @test got == ["target:/", "echo: hello"]
+            finally
+                close(client_pipe)
+                wait(fwd_up)
+                wait(fwd_down)
+            end
+        end
+    finally
+        close(server)
+    end
+end

test/runtests.jl

@@ -106,6 +106,7 @@ test_files = [
     "http_websocket_client_tests.jl",
     "http_websocket_server_tests.jl",
     "http_websocket_integration_tests.jl",
+    "http_websocket_raw_io_tests.jl",
     "http_retry_tests.jl",
     "http_client_transport_tests.jl",
     "http_client_proxy_tests.jl",