http.rb in httpx-0.20.0

- old
+ new

@@ -4,10 +4,46 @@
 
 module HTTPX
   module Plugins
     module Proxy
       module HTTP
+        module InstanceMethods
+          def with_proxy_basic_auth(opts)
+            with(proxy: opts.merge(scheme: "basic"))
+          end
+
+          def with_proxy_digest_auth(opts)
+            with(proxy: opts.merge(scheme: "digest"))
+          end
+
+          def with_proxy_ntlm_auth(opts)
+            with(proxy: opts.merge(scheme: "ntlm"))
+          end
+
+          def fetch_response(request, connections, options)
+            response = super
+
+            if response &&
+               response.status == 407 &&
+               !request.headers.key?("proxy-authorization") &&
+               response.headers.key?("proxy-authenticate")
+
+              connection = find_connection(request, connections, options)
+
+              if connection.options.proxy.can_authenticate?(response.headers["proxy-authenticate"])
+                request.transition(:idle)
+                request.headers["proxy-authorization"] =
+                  connection.options.proxy.authenticate(request, response.headers["proxy-authenticate"])
+                connection.send(request)
+                return
+              end
+            end
+
+            response
+          end
+        end
+
         module ConnectionMethods
           def connecting?
             super || @state == :connecting || @state == :connected
           end
 
@@ -21,15 +57,31 @@
               return unless @state == :idle
 
               @io.connect
               return unless @io.connected?
 
-              @parser = registry(@io.protocol).new(@write_buffer, @options.merge(max_concurrent_requests: 1))
-              @parser.extend(ProxyParser)
-              @parser.once(:response, &method(:__http_on_connect))
-              @parser.on(:close) { transition(:closing) }
-              __http_proxy_connect
+              @parser || begin
+                @parser = registry(@io.protocol).new(@write_buffer, @options.merge(max_concurrent_requests: 1))
+                parser = @parser
+                parser.extend(ProxyParser)
+                parser.on(:response, &method(:__http_on_connect))
+                parser.on(:close) { transition(:closing) }
+                parser.on(:reset) do
+                  if parser.empty?
+                    reset
+                  else
+                    transition(:closing)
+                    transition(:closed)
+                    emit(:reset)
+
+                    parser.reset if @parser
+                    transition(:idle)
+                    transition(:connecting)
+                  end
+                end
+                __http_proxy_connect(parser)
+              end
               return if @state == :connected
             when :connected
               return unless @state == :idle || @state == :connecting
 
               case @state
@@ -42,33 +94,41 @@
               end
             end
             super
           end
 
-          def __http_proxy_connect
+          def __http_proxy_connect(parser)
             req = @pending.first
-            # if the first request after CONNECT is to an https address, it is assumed that
-            # all requests in the queue are not only ALL HTTPS, but they also share the certificate,
-            # and therefore, will share the connection.
-            #
-            if req.uri.scheme == "https"
+            if req && req.uri.scheme == "https"
+              # if the first request after CONNECT is to an https address, it is assumed that
+              # all requests in the queue are not only ALL HTTPS, but they also share the certificate,
+              # and therefore, will share the connection.
+              #
               connect_request = ConnectRequest.new(req.uri, @options)
               @inflight += 1
               parser.send(connect_request)
             else
               handle_transition(:connected)
             end
           end
 
-          def __http_on_connect(_, response)
+          def __http_on_connect(request, response)
             @inflight -= 1
             if response.status == 200
               req = @pending.first
               request_uri = req.uri
               @io = ProxySSL.new(@io, request_uri, @options)
               transition(:connected)
               throw(:called)
+            elsif response.status == 407 &&
+                  !request.headers.key?("proxy-authorization") &&
+                  @options.proxy.can_authenticate?(response.headers["proxy-authenticate"])
+
+              request.transition(:idle)
+              request.headers["proxy-authorization"] = @options.proxy.authenticate(request, response.headers["proxy-authenticate"])
+              @parser.send(request)
+              @inflight += 1
             else
               pending = @pending + @parser.pending
               while (req = pending.shift)
                 req.emit(:response, response)
               end
@@ -86,10 +146,13 @@
 
           def set_protocol_headers(request)
             extra_headers = super
 
             proxy_params = @options.proxy
-            extra_headers["proxy-authorization"] = "Basic #{proxy_params.token_authentication}" if proxy_params.authenticated?
+            if proxy_params.scheme == "basic"
+              # opt for basic auth
+              extra_headers["proxy-authorization"] = proxy_params.authenticate(extra_headers)
+            end
             extra_headers["proxy-connection"] = extra_headers.delete("connection") if extra_headers.key?("connection")
             extra_headers
           end
         end
 