Support custom CONNECT headers (both request and response) during HTTPS proxy requests

lib/http/chainable.rb

@@ -153,8 +153,10 @@ def via(*proxy)
       proxy_hash[:proxy_port]     = proxy[1] if proxy[1].is_a?(Integer)
       proxy_hash[:proxy_username] = proxy[2] if proxy[2].is_a?(String)
       proxy_hash[:proxy_password] = proxy[3] if proxy[3].is_a?(String)
+      proxy_hash[:proxy_headers]  = proxy[2] if proxy[2].is_a?(Hash)
+      proxy_hash[:proxy_headers]  = proxy[4] if proxy[4].is_a?(Hash)
 
-      raise(RequestError, "invalid HTTP proxy: #{proxy_hash}") unless [2, 4].include?(proxy_hash.keys.size)
+      raise(RequestError, "invalid HTTP proxy: #{proxy_hash}") unless (2..5).cover?(proxy_hash.keys.size)
 
       branch default_options.with_proxy(proxy_hash)
     end

lib/http/client.rb

@@ -68,12 +68,13 @@ def perform(req, options)
       end
 
       res = Response.new(
-        :status     => @connection.status_code,
-        :version    => @connection.http_version,
-        :headers    => @connection.headers,
-        :connection => @connection,
-        :encoding   => options.encoding,
-        :uri        => req.uri
+        :status        => @connection.status_code,
+        :version       => @connection.http_version,
+        :headers       => @connection.headers,
+        :proxy_headers => @connection.proxy_response_headers,
+        :connection    => @connection,
+        :encoding      => options.encoding,
+        :uri           => req.uri
       )
 
       @connection.finish_response if req.verb == :head

lib/http/connection.rb

@@ -19,6 +19,9 @@ class Connection
     # HTTP/1.1
     HTTP_1_1 = "1.1".freeze
 
+    # Returned after HTTP CONNECT (via proxy)
+    attr_reader :proxy_response_headers
+
     # @param [HTTP::Request] req
     # @param [HTTP::Options] options
     # @raise [HTTP::ConnectionError] when failed to connect
@@ -164,6 +167,7 @@ def send_proxy_connect_request(req)
       @pending_response = true
 
       read_headers!
+      @proxy_response_headers = @parser.headers
 
       if @parser.status_code != 200
         @failed_proxy_connect = true

lib/http/request.rb

@@ -106,7 +106,7 @@ def redirect(uri, verb = @verb)
 
     # Stream the request to a socket
     def stream(socket)
-      include_proxy_authorization_header if using_authenticated_proxy? && [email protected]?
+      include_proxy_headers if using_proxy? && [email protected]?
       Request::Writer.new(socket, body, headers, headline).stream
     end
 
@@ -117,7 +117,12 @@ def using_proxy?
 
     # Is this request using an authenticated proxy?
     def using_authenticated_proxy?
-      proxy && proxy.keys.size == 4
+      proxy && proxy.keys.size >= 4
+    end
+
+    def include_proxy_headers
+      headers.merge!(proxy[:proxy_headers]) if proxy.key?(:proxy_headers)
+      include_proxy_authorization_header if using_authenticated_proxy?
     end
 
     # Compute and add the Proxy-Authorization header
@@ -154,7 +159,7 @@ def proxy_connect_headers
       )
 
       connect_headers[Headers::PROXY_AUTHORIZATION] = proxy_authorization_header if using_authenticated_proxy?
-
+      connect_headers.merge!(proxy[:proxy_headers]) if proxy.key?(:proxy_headers)
       connect_headers
     end
 

lib/http/response.rb

@@ -23,20 +23,25 @@ class Response
     # @return [URI, nil]
     attr_reader :uri
 
+    # @return [Hash]
+    attr_reader :proxy_headers
+
     # Inits a new instance
     #
     # @option opts [Integer] :status Status code
     # @option opts [String] :version HTTP version
     # @option opts [Hash] :headers
+    # @option opts [Hash] :proxy_headers
     # @option opts [HTTP::Connection] :connection
     # @option opts [String] :encoding Encoding to use when reading body
     # @option opts [String] :body
     # @option opts [String] :uri
     def initialize(opts)
-      @version  = opts.fetch(:version)
-      @uri      = HTTP::URI.parse(opts.fetch(:uri)) if opts.include? :uri
-      @status   = HTTP::Response::Status.new(opts.fetch(:status))
-      @headers  = HTTP::Headers.coerce(opts[:headers] || {})
+      @version       = opts.fetch(:version)
+      @uri           = HTTP::URI.parse(opts.fetch(:uri)) if opts.include? :uri
+      @status        = HTTP::Response::Status.new(opts.fetch(:status))
+      @headers       = HTTP::Headers.coerce(opts[:headers] || {})
+      @proxy_headers = HTTP::Headers.coerce(opts[:proxy_headers] || {})
 
       if opts.include?(:connection)
         connection = opts.fetch(:connection)