Connection pools support keepalive through a proxyserver

gustafsson · gustafsson · commit 002647af29a2 · 2023-05-15T12:15:33.000+02:00
diff --git a/src/Connections.jl b/src/Connections.jl
@@ -25,6 +25,7 @@ using Sockets, LoggingExtras, NetworkOptions
 using MbedTLS: SSLConfig, SSLContext, setup!, associate!, hostname!, handshake!
 using MbedTLS, OpenSSL, ConcurrentUtilities
 using ..IOExtras, ..Conditions, ..Exceptions
+using URIs: URI
 
 const nolimit = typemax(Int)
 
@@ -73,6 +74,7 @@ Fields:
 mutable struct Connection{IO_t <: IO} <: IO
     host::String
     port::String
+    proxy::Union{Nothing,String}
     idle_timeout::Int
     require_ssl_verification::Bool
     keepalive::Bool
@@ -98,20 +100,21 @@ request parameters of host and port, and if ssl verification is required, if kee
 and if an existing Connection was already created with the exact.
 same parameters, we can re-use it (as long as it's not already being used, obviously).
 """
-connectionkey(x::Connection) = (x.host, x.port, x.require_ssl_verification, x.keepalive, x.clientconnection)
+connectionkey(x::Connection) = (x.host, x.port, x.proxy, x.require_ssl_verification, x.keepalive, x.clientconnection)
 
-const ConnectionKeyType = Tuple{AbstractString, AbstractString, Bool, Bool, Bool}
+const ConnectionKeyType = Tuple{AbstractString, AbstractString, Union{Nothing,AbstractString}, Bool, Bool, Bool}
 
-Connection(host::AbstractString, port::AbstractString,
+Connection(host::AbstractString, port::AbstractString, proxy::Union{Nothing,AbstractString},
            idle_timeout::Int,
            require_ssl_verification::Bool, keepalive::Bool, io::T, client=true) where {T}=
-    Connection{T}(host, port, idle_timeout,
+    Connection{T}(host, port, proxy, idle_timeout,
                 require_ssl_verification, keepalive,
                 safe_getpeername(io)..., localport(io),
-                io, client, PipeBuffer(), time(), false, false, IOBuffer(), nothing)
+                io, client,
+                PipeBuffer(), time(), false, false, IOBuffer(), nothing)
 
 Connection(io; require_ssl_verification::Bool=true, keepalive::Bool=true) =
-    Connection("", "", 0, require_ssl_verification, keepalive, io, false)
+    Connection("", "", nothing, 0, require_ssl_verification, keepalive, io, false)
 
 getrawstream(c::Connection) = c.io
 
@@ -432,30 +435,59 @@ end
 Find a reusable `Connection` in the `pool`,
 or create a new `Connection` if required.
 """
-function newconnection(::Type{T},
-                       host::AbstractString,
-                       port::AbstractString;
+function newconnection(wrapconnection::Function,
+                       url::URI;
+                       proxy::Union{Nothing, AbstractString}=nothing,
+                       socket_type::Type,
+                       socket_type_tls::Type,
                        pool::Union{Nothing, Pool}=nothing,
                        connection_limit=nothing,
                        forcenew::Bool=false,
                        idle_timeout=typemax(Int),
-                       require_ssl_verification::Bool=NetworkOptions.verify_host(host, "SSL"),
+                       require_ssl_verification::Bool=NetworkOptions.verify_host(url.host, "SSL"),
                        keepalive::Bool=true,
-                       kw...) where {T <: IO}
+                       kw...)
+    IOType = sockettype(url, socket_type, socket_type_tls)
+
     connection_limit_warning(connection_limit)
-    return acquire(
-            getpool(pool, T),
-            (host, port, require_ssl_verification, keepalive, true);
+
+    key = (url.host, url.port, proxy, require_ssl_verification, keepalive, true)
+
+    acquire(
+            getpool(pool, IOType),
+            key;
             forcenew=forcenew,
             isvalid=c->connection_isvalid(c, Int(idle_timeout))) do
-        Connection(host, port,
+
+        if proxy !== nothing
+            url = URI(proxy)
+        end
+
+        innerIOType = sockettype(url, socket_type, socket_type_tls)
+
+        io = Connection(url.host, url.port, proxy,
             idle_timeout, require_ssl_verification, keepalive,
-            getconnection(T, host, port;
-                require_ssl_verification=require_ssl_verification, keepalive=keepalive, kw...)
+            getconnection(innerIOType, url.host, url.port;
+                require_ssl_verification, keepalive, kw...)
         )
+
+        try
+            io = wrapconnection(io)
+        catch ex
+            @try Base.IOError close(io)
+            rethrow(ex)
+        end
+
+        if connectionkey(io) != key
+            throw(ErrorException(string("Connection error ", (;expected = connectionkey(io), key))))
+        end
+
+        io
     end
 end
 
+sockettype(url::URI, tcp, tls) = url.scheme in ("wss", "https") ? tls : tcp
+
 function releaseconnection(c::Connection{T}, reuse; pool::Union{Nothing, Pool}=nothing, kw...) where {T}
     c.timestamp = time()
     release(getpool(pool, T), connectionkey(c), reuse ? c : nothing)
@@ -615,7 +647,6 @@ end
 
 function sslupgrade(::Type{IOType}, c::Connection{T},
                     host::AbstractString;
-                    pool::Union{Nothing, Pool}=nothing,
                     require_ssl_verification::Bool=NetworkOptions.verify_host(host, "SSL"),
                     keepalive::Bool=true,
                     readtimeout::Int=0,
@@ -630,12 +661,9 @@ function sslupgrade(::Type{IOType}, c::Connection{T},
     else
         sslconnection(IOType, c.io, host; require_ssl_verification=require_ssl_verification, keepalive=keepalive, kw...)
     end
+
     # success, now we turn it into a new Connection
-    conn = Connection(host, "", 0, require_ssl_verification, keepalive, tls)
-    # release the "old" one, but don't return the connection since we're hijacking the socket
-    release(getpool(pool, T), connectionkey(c))
-    # and return the new one
-    return acquire(() -> conn, getpool(pool, IOType), connectionkey(conn); forcenew=true)
+    Connection(host,"", "", 0, require_ssl_verification, keepalive, tls)
 end
 
 function Base.show(io::IO, c::Connection)
diff --git a/src/clientlayers/ConnectionRequest.jl b/src/clientlayers/ConnectionRequest.jl
@@ -57,27 +57,37 @@ Otherwise leave it open so that it can be reused.
 function connectionlayer(handler)
     return function connections(req; proxy=getproxy(req.url.scheme, req.url.host), socket_type::Type=TCPSocket, socket_type_tls::Type=SOCKET_TYPE_TLS[], readtimeout::Int=0, logerrors::Bool=false, logtag=nothing, kw...)
         local io, stream
+        url = req.url
         if proxy !== nothing
-            target_url = req.url
-            url = URI(proxy)
-            if target_url.scheme == "http"
-                req.target = string(target_url)
-            end
-
             userinfo = unescapeuri(url.userinfo)
             if !isempty(userinfo) && !hasheader(req.headers, "Proxy-Authorization")
                 @debugv 1 "Adding Proxy-Authorization: Basic header."
                 setheader(req.headers, "Proxy-Authorization" => "Basic $(base64encode(userinfo))")
             end
-        else
-            url = target_url = req.url
+
+            proxyauth = header(req, "Proxy-Authorization")
+            if url.scheme in ("https", "wss", "ws")
+                req.headers = filter(x->x.first != "Proxy-Authorization", req.headers)
+            end
+
+            if url.scheme == "http"
+                req.target = string(url)
+            end
         end
 
-        IOType = sockettype(url, socket_type, socket_type_tls)
         start_time = time()
         try
-            io = newconnection(IOType, url.host, url.port; readtimeout=readtimeout, kw...)
+            io = newconnection(url; proxy, socket_type, socket_type_tls, readtimeout, kw...) do io
+                if isnothing(proxy)
+                    io
+                else
+                    proxy_tunnel(io, url; proxyauth, proxy, socket_type_tls, readtimeout, kw...)
+                end
+            end
         catch e
+            if e isa StatusError && !get(kw, :status_exception, true)
+                return e.response
+            end
             if logerrors
                 err = current_exceptions_to_string(CapturedException(e, catch_backtrace()))
                 @error err type=Symbol("HTTP.ConnectError") method=req.method url=req.url context=req.context logtag=logtag
@@ -88,33 +98,8 @@ function connectionlayer(handler)
             req.context[:connect_duration_ms] = get(req.context, :connect_duration_ms, 0.0) +  (time() - start_time) * 1000
         end
 
-        shouldreuse = !(target_url.scheme in ("ws", "wss"))
+        shouldreuse = !(url.scheme in ("ws", "wss"))
         try
-            if proxy !== nothing && target_url.scheme in ("https", "wss", "ws")
-                shouldreuse = false
-                # tunnel request
-                if target_url.scheme in ("https", "wss")
-                    target_url = URI(target_url, port=443)
-                elseif target_url.scheme in ("ws", ) && target_url.port == ""
-                    target_url = URI(target_url, port=80) # if there is no port info, connect_tunnel will fail
-                end
-                r = if readtimeout > 0
-                    try_with_timeout(readtimeout) do _
-                        connect_tunnel(io, target_url, req)
-                    end
-                else
-                    connect_tunnel(io, target_url, req)
-                end
-                if r.status != 200
-                    close(io)
-                    return r
-                end
-                if target_url.scheme in ("https", "wss")
-                    io = Connections.sslupgrade(socket_type_tls, io, target_url.host; readtimeout=readtimeout, kw...)
-                end
-                req.headers = filter(x->x.first != "Proxy-Authorization", req.headers)
-            end
-
             stream = Stream(req.response, io)
             return handler(stream; readtimeout=readtimeout, logerrors=logerrors, logtag=logtag, kw...)
         catch e
@@ -151,13 +136,42 @@ function connectionlayer(handler)
     end
 end
 
-sockettype(url::URI, tcp, tls) = url.scheme in ("wss", "https") ? tls : tcp
+function proxy_tunnel(io, url; proxyauth, socket_type_tls, readtimeout, proxy, kw...)
+    if url.scheme in ("https", "wss", "ws")
+        # tunnel request
+
+        if (port = url.port; isempty(port))
+            # if there is no port info, connect_tunnel will fail
+            # The pool entry still has port = ""
+            if url.scheme in ("https", "wss")
+                port = 443
+            elseif scheme in ("ws", )
+                port = 80
+            end
+        end
 
-function connect_tunnel(io, target_url, req)
-    target = "$(URIs.hoststring(target_url.host)):$(target_url.port)"
+        if readtimeout > 0
+            try_with_timeout(readtimeout) do _
+                connect_tunnel(io, url.host, port, proxyauth)
+            end
+        else
+            connect_tunnel(io, url.host, port, proxyauth)
+        end
+
+        if url.scheme in ("https", "wss")
+            io = Connections.sslupgrade(socket_type_tls, io, url.host; readtimeout, kw...)
+            io.port = url.port
+            io.proxy = proxy
+        end
+    end
+    io
+end
+
+function connect_tunnel(io, host, port, auth)
+    target = "$(URIs.hoststring(host)):$(port)"
     @debugv 1 "📡  CONNECT HTTPS tunnel to $target"
     headers = Dict("Host" => target)
-    if (auth = header(req, "Proxy-Authorization"); !isempty(auth))
+    if !isempty(auth)
         headers["Proxy-Authorization"] = auth
     end
     request = Request("CONNECT", target, headers)
@@ -166,7 +180,10 @@ function connect_tunnel(io, target_url, req)
     # @debugv 2 "connect_tunnel: reading headers"
     readheaders(io, request.response)
     # @debugv 2 "connect_tunnel: done reading headers"
-    return request.response
+    if request.response.status != 200
+        throw(StatusError(request.response.status, 
+              request.method, request.target, request.response))
+    end
 end
 
 end # module ConnectionRequest
diff --git a/test/client.jl b/test/client.jl
@@ -553,6 +553,58 @@ end
     end
 end
 
+@testset "HTTP CONNECT Proxy pool" begin
+    # Stores the http request passed by the client
+    messages = []
+    streams = Set()
+
+    function forwardstream(src, dst)
+        while isopen(dst) && isopen(src) && !eof(src)
+            buff = readavailable(src)
+            !isempty(buff) && isopen(dst) && write(dst, buff)
+        end
+
+        close(src)
+        close(dst)
+    end
+
+    # Simple implementation of a proxy server
+    proxy = HTTP.listen!(IPv4(0), 8082; stream = true) do http::HTTP.Stream
+        push!(messages, http.message)
+        host, port = split(http.message.target, ":")
+        targetstream = connect(host, parse(Int, port))
+        HTTP.setstatus(http, 200)
+        HTTP.startwrite(http)
+        up = @async forwardstream(http.stream.io, targetstream)
+        down = @async forwardstream(targetstream, http.stream.io)
+        push!(streams, targetstream)
+        wait(up)
+        wait(down)
+        delete!(streams, targetstream)
+    end
+
+    try
+        # Make the HTTP request
+        r1 = HTTP.get("https://example.com:443"; proxy="http://localhost:8082", retry=false, status_exception=true)
+        @test length(messages) == 1
+        @test first(messages).method == "CONNECT"
+        @test length(streams) == 1 && isopen(first(streams)) # still alive
+
+        # Make another request
+        # This should reuse the connection pool and not make another request to the proxy
+        empty!(messages)
+        r2 = HTTP.get("https://example.com:443"; proxy="http://localhost:8082", retry=false, status_exception=true)
+        @test isempty(messages)
+        @test r1.body == r2.body # no new message to the proxy yet successfully get the same response from the target server
+        @test length(streams) == 1 && isopen(first(streams)) # still only one stream alive
+    finally
+        close.(streams)
+        close(proxy)
+        HTTP.Connections.closeall()
+        wait(proxy)
+    end
+end
+
 @testset "Retry with request/response body streams" begin
     shouldfail = Ref(true)
     status = Ref(200)
