proxy.rb in httpx-0.4.0

- old
+ new

@@ -4,15 +4,23 @@
 require "ipaddr"
 require "forwardable"
 
 module HTTPX
   module Plugins
+    #
+    # This plugin adds support for proxies. It ships with support for:
+    #
+    # * HTTP proxies
+    # * HTTPS proxies
+    # * Socks4/4a proxies
+    # * Socks5 proxies
+    #
     module Proxy
       Error = Class.new(Error)
-      class Parameters
-        extend Registry
+      PROXY_ERRORS = [TimeoutError, IOError, SystemCallError, Error].freeze
 
+      class Parameters
         attr_reader :uri, :username, :password
 
         def initialize(uri:, username: nil, password: nil)
           @uri = uri.is_a?(URI::Generic) ? uri : URI(uri)
           @username = username || @uri.user
@@ -24,130 +32,169 @@
         end
 
         def token_authentication
           Base64.strict_encode64("#{user}:#{password}")
         end
+
+        def ==(other)
+          if other.is_a?(Parameters)
+            @uri == other.uri &&
+              @username == other.username &&
+              @password == other.password
+          else
+            super
+          end
+        end
       end
 
+      class << self
+        def configure(klass, *)
+          klass.plugin(:"proxy/http")
+          klass.plugin(:"proxy/socks4")
+          klass.plugin(:"proxy/socks5")
+        end
+
+        def extra_options(options)
+          Class.new(options.class) do
+            def_option(:proxy) do |pr|
+              Hash[pr]
+            end
+          end.new(options)
+        end
+      end
+
       module InstanceMethods
         def with_proxy(*args)
           branch(default_options.with_proxy(*args))
         end
 
         private
 
-        def proxy_params(uri)
+        def proxy_uris(uri, options)
           @_proxy_uris ||= begin
-            uris = @options.proxy ? Array(@options.proxy[:uri]) : []
+            uris = options.proxy ? Array(options.proxy[:uri]) : []
             if uris.empty?
               uri = URI(uri).find_proxy
               uris << uri if uri
             end
             uris
           end
-          @options.proxy.merge(uri: @_proxy_uris.shift) unless @_proxy_uris.empty?
+          options.proxy.merge(uri: @_proxy_uris.first) unless @_proxy_uris.empty?
         end
 
-        def find_channel(request, **options)
+        def find_connection(request, connections, options)
+          return super unless options.respond_to?(:proxy)
+
           uri = URI(request.uri)
-          proxy = proxy_params(uri)
-          raise Error, "Failed to connect to proxy" unless proxy
-          @connection.find_channel(proxy) || build_channel(proxy, options)
-        end
+          next_proxy = proxy_uris(uri, options)
+          raise Error, "Failed to connect to proxy" unless next_proxy
 
-        def build_channel(proxy, options)
-          return super if proxy.is_a?(URI::Generic)
-          channel = build_proxy_channel(proxy, **options)
-          set_channel_callbacks(channel, options)
-          channel
+          proxy_options = options.merge(proxy: Parameters.new(**next_proxy))
+          connection = pool.find_connection(uri, proxy_options) || build_connection(uri, proxy_options)
+          unless connections.nil? || connections.include?(connection)
+            connections << connection
+            set_connection_callbacks(connection, options)
+          end
+          connection
         end
 
-        def build_proxy_channel(proxy, **options)
-          parameters = Parameters.new(**proxy)
-          uri = parameters.uri
-          log { "proxy: #{uri}" }
-          proxy_type = Parameters.registry(parameters.uri.scheme)
-          channel = proxy_type.new("tcp", uri, parameters, @options.merge(options), &method(:on_response))
-          @connection.__send__(:resolve_channel, channel)
-          channel
+        def build_connection(uri, options)
+          proxy = options.proxy
+          return super unless proxy
+
+          connection = options.connection_class.new("tcp", uri, options)
+          pool.init_connection(connection, options)
+          connection
         end
 
-        def fetch_response(request)
+        def fetch_response(request, connections, options)
           response = super
           if response.is_a?(ErrorResponse) &&
              # either it was a timeout error connecting, or it was a proxy error
-             (((response.error.is_a?(TimeoutError) || response.error.is_a?(IOError)) && request.state == :idle) ||
-              response.error.is_a?(Error)) &&
-             !@_proxy_uris.empty?
+             PROXY_ERRORS.any? { |ex| response.error.is_a?(ex) } && !@_proxy_uris.empty?
+            @_proxy_uris.shift
             log { "failed connecting to proxy, trying next..." }
-            channel = find_channel(request)
-            channel.send(request)
+            request.transition(:idle)
+            connection = find_connection(request, connections, options)
+            connections << connection unless connections.include?(connection)
+            connection.send(request)
             return
           end
           response
         end
       end
 
-      module OptionsMethods
-        def self.included(klass)
+      module ConnectionMethods
+        using URIExtensions
+
+        def initialize(*)
           super
-          klass.def_option(:proxy) do |pr|
-            Hash[pr]
-          end
+          return unless @options.proxy
+
+          # redefining the connection origin as the proxy's URI,
+          # as this will be used as the tcp peer ip.
+          @origin = URI(@options.proxy.uri.origin)
         end
-      end
 
-      def self.configure(klass, *)
-        klass.plugin(:"proxy/http")
-        klass.plugin(:"proxy/socks4")
-        klass.plugin(:"proxy/socks5")
-      end
-    end
-    register_plugin :proxy, Proxy
-  end
+        def match?(uri, options)
+          return super unless @options.proxy
 
-  class ProxyChannel < Channel
-    def initialize(type, uri, parameters, options, &blk)
-      super(type, uri, options, &blk)
-      @parameters = parameters
-    end
+          super && @options.proxy == options.proxy
+        end
 
-    def match?(*)
-      true
-    end
+        # should not coalesce connections here, as the IP is the IP of the proxy
+        def coalescable?(*)
+          return super unless @options.proxy
 
-    def send(request, **args)
-      @pending << [request, args]
-    end
+          false
+        end
 
-    def connecting?
-      super || @state == :connecting || @state == :connected
-    end
+        def send(request)
+          return super unless @options.proxy
+          return super unless connecting?
 
-    def to_io
-      case @state
-      when :idle
-        transition(:connecting)
-      when :connected
-        transition(:open)
-      end
-      @io.to_io
-    end
+          @pending << request
+        end
 
-    def call
-      super
-      case @state
-      when :connecting
-        consume
+        def connecting?
+          return super unless @options.proxy
+
+          super || @state == :connecting || @state == :connected
+        end
+
+        def to_io
+          return super unless @options.proxy
+
+          case @state
+          when :idle
+            transition(:connecting)
+          when :connected
+            transition(:open)
+          end
+          @io.to_io
+        end
+
+        def call
+          super
+          return unless @options.proxy
+
+          case @state
+          when :connecting
+            consume
+          end
+        end
+
+        def reset
+          return super unless @options.proxy
+
+          @state = :open
+          transition(:closing)
+          transition(:closed)
+          emit(:close)
+        end
       end
     end
-
-    def reset
-      @state = :open
-      transition(:closing)
-      transition(:closed)
-      emit(:close)
-    end
+    register_plugin :proxy, Proxy
   end
 
   class ProxySSL < SSL
     def initialize(tcp, request_uri, options)
       @io = tcp.to_io