Merge pull request #742 from jjiang-stripe/jjiang/ssl-proxy-headers

geemus · web-flow · commit 7c154c11df51 · 2021-04-24T15:50:56.000-05:00
Support additional headers for TLS proxies
diff --git a/lib/excon/connection.rb b/lib/excon/connection.rb
@@ -115,7 +115,7 @@ def request_call(datum)
           # we already have data from a middleware, so bail
           return datum
         else
-          socket.data = datum
+          socket(datum)
           # start with "METHOD /path"
           request = datum[:method].to_s.upcase + ' '
           if datum[:proxy] && datum[:scheme] != HTTPS
@@ -144,35 +144,25 @@ def request_call(datum)
           end
 
           # add headers to request
-          datum[:headers].each do |key, values|
-            if key.to_s.match(/[\r\n]/)
-              raise Excon::Errors::InvalidHeaderKey.new(key.to_s.inspect + ' contains forbidden "\r" or "\n"')
-            end
-            [values].flatten.each do |value|
-              if value.to_s.match(/[\r\n]/)
-                raise Excon::Errors::InvalidHeaderValue.new(value.to_s.inspect + ' contains forbidden "\r" or "\n"')
-              end
-              request << key.to_s << ': ' << value.to_s << CR_NL
-            end
-          end
+          request << Utils.headers_hash_to_s(datum[:headers])
 
           # add additional "\r\n" to indicate end of headers
           request << CR_NL
 
           if datum.has_key?(:request_block)
-            socket.write(request) # write out request + headers
+            socket(datum).write(request) # write out request + headers
             while true # write out body with chunked encoding
               chunk = datum[:request_block].call
               chunk = binary_encode(chunk)
               if chunk.length > 0
-                socket.write(chunk.length.to_s(16) << CR_NL << chunk << CR_NL)
+                socket(datum).write(chunk.length.to_s(16) << CR_NL << chunk << CR_NL)
               else
-                socket.write(String.new("0#{CR_NL}#{CR_NL}"))
+                socket(datum).write(String.new("0#{CR_NL}#{CR_NL}"))
                 break
               end
             end
           elsif body.nil?
-            socket.write(request) # write out request + headers
+            socket(datum).write(request) # write out request + headers
           else # write out body
             if body.respond_to?(:binmode) && !body.is_a?(StringIO)
               body.binmode
@@ -186,13 +176,13 @@ def request_call(datum)
             chunk = body.read([datum[:chunk_size] - request.length, 0].max)
             if chunk
               chunk = binary_encode(chunk)
-              socket.write(request << chunk)
+              socket(datum).write(request << chunk)
             else
-              socket.write(request) # write out request + headers
+              socket(datum).write(request) # write out request + headers
             end
 
             while (chunk = body.read(datum[:chunk_size]))
-              socket.write(chunk)
+              socket(datum).write(chunk)
             end
           end
         end
@@ -463,14 +453,14 @@ def response(datum={})
       end
     end
 
-    def socket
-      unix_proxy = @data[:proxy] ? @data[:proxy][:scheme] == UNIX : false
-      sockets[@socket_key] ||= if @data[:scheme] == UNIX || unix_proxy
-        Excon::UnixSocket.new(@data)
-      elsif @data[:ssl_uri_schemes].include?(@data[:scheme])
-        Excon::SSLSocket.new(@data)
+    def socket(datum = @data)
+      unix_proxy = datum[:proxy] ? datum[:proxy][:scheme] == UNIX : false
+      sockets[@socket_key] ||= if datum[:scheme] == UNIX || unix_proxy
+        Excon::UnixSocket.new(datum)
+      elsif datum[:ssl_uri_schemes].include?(datum[:scheme])
+        Excon::SSLSocket.new(datum)
       else
-        Excon::Socket.new(@data)
+        Excon::Socket.new(datum)
       end
     end
 
@@ -573,6 +563,9 @@ def setup_proxy
         if uri.user
           @data[:proxy][:user] = uri.user
         end
+        if @data[:ssl_proxy_headers] && !@data[:ssl_uri_schemes].include?(@data[:scheme])
+          raise ArgumentError, "The `:ssl_proxy_headers` parameter should only be used with HTTPS requests."
+        end
         if @data[:proxy][:scheme] == UNIX
           if @data[:proxy][:host]
             raise ArgumentError, "The `:host` parameter should not be set for `unix://` proxies.\n" +
diff --git a/lib/excon/constants.rb b/lib/excon/constants.rb
@@ -99,6 +99,7 @@ module Excon
     :ssl_version,
     :ssl_min_version,
     :ssl_max_version,
+    :ssl_proxy_headers,
     :ssl_uri_schemes,
     :tcp_nodelay,
     :thread_safe_sockets,
diff --git a/lib/excon/ssl_socket.rb b/lib/excon/ssl_socket.rb
@@ -104,6 +104,10 @@ def initialize(data = {})
 
         request += "Proxy-Connection: Keep-Alive#{Excon::CR_NL}"
 
+        if @data[:ssl_proxy_headers]
+          request << Utils.headers_hash_to_s(@data[:ssl_proxy_headers])
+        end
+
         request += Excon::CR_NL
 
         # write out the proxy setup request
diff --git a/lib/excon/utils.rb b/lib/excon/utils.rb
@@ -121,5 +121,22 @@ def unescape_form(str)
       str.gsub!(/\+/, ' ')
       str.gsub(ESCAPED) { $1.hex.chr }
     end
+
+    # Performs validation on the passed header hash and returns a string representation of the headers
+    def headers_hash_to_s(headers)
+      headers_str = String.new
+      headers.each do |key, values|
+        if key.to_s.match(/[\r\n]/)
+          raise Excon::Errors::InvalidHeaderKey.new(key.to_s.inspect + ' contains forbidden "\r" or "\n"')
+        end
+        [values].flatten.each do |value|
+          if value.to_s.match(/[\r\n]/)
+            raise Excon::Errors::InvalidHeaderValue.new(value.to_s.inspect + ' contains forbidden "\r" or "\n"')
+          end
+          headers_str << key.to_s << ': ' << value.to_s << CR_NL
+        end
+      end
+      headers_str
+    end
   end
 end
