socks4.rb in httpx-0.4.0

- old
+ new

@@ -8,62 +8,67 @@
     module Proxy
       module Socks4
         VERSION = 4
         CONNECT = 1
         GRANTED = 90
+        PROTOCOLS = %w[socks4 socks4a].freeze
 
         Error = Class.new(Error)
 
-        class Socks4ProxyChannel < ProxyChannel
+        module ConnectionMethods
           private
 
-          def proxy_connect
-            @parser = SocksParser.new(@write_buffer, @options)
-            @parser.once(:packet, &method(:on_packet))
-          end
-
-          def on_packet(packet)
-            _version, status, _port, _ip = packet.unpack("CCnN")
-            if status == GRANTED
-              req, _ = @pending.first
-              request_uri = req.uri
-              @io = ProxySSL.new(@io, request_uri, @options) if request_uri.scheme == "https"
-              transition(:connected)
-              throw(:called)
-            else
-              on_socks_error("socks error: #{status}")
-            end
-          end
-
           def transition(nextstate)
+            return super unless @options.proxy && PROTOCOLS.include?(@options.proxy.uri.scheme)
+
             case nextstate
             when :connecting
               return unless @state == :idle
+
               @io.connect
               return unless @io.connected?
+
               req, _ = @pending.first
               return unless req
+
               request_uri = req.uri
-              @write_buffer << Packet.connect(@parameters, request_uri)
-              proxy_connect
+              @write_buffer << Packet.connect(@options.proxy, request_uri)
+              __socks4_proxy_connect
             when :connected
               return unless @state == :connecting
+
               @parser = nil
             end
             log(level: 1, label: "SOCKS4: ") { "#{nextstate}: #{@write_buffer.to_s.inspect}" } unless nextstate == :open
             super
           end
 
-          def on_socks_error(message)
+          def __socks4_proxy_connect
+            @parser = SocksParser.new(@write_buffer, @options)
+            @parser.once(:packet, &method(:__socks4_on_packet))
+          end
+
+          def __socks4_on_packet(packet)
+            _version, status, _port, _ip = packet.unpack("CCnN")
+            if status == GRANTED
+              req, _ = @pending.first
+              request_uri = req.uri
+              @io = ProxySSL.new(@io, request_uri, @options) if request_uri.scheme == "https"
+              transition(:connected)
+              throw(:called)
+            else
+              on_socks4_error("socks error: #{status}")
+            end
+          end
+
+          def on_socks4_error(message)
             ex = Error.new(message)
             ex.set_backtrace(caller)
             on_error(ex)
             throw(:called)
           end
         end
-        Parameters.register("socks4", Socks4ProxyChannel)
-        Parameters.register("socks4a", Socks4ProxyChannel)
 
         class SocksParser
           include Callbacks
 
           def initialize(buffer, options)
@@ -90,9 +95,10 @@
           def connect(parameters, uri)
             packet = [VERSION, CONNECT, uri.port].pack("CCn")
             begin
               ip = IPAddr.new(uri.host)
               raise Error, "Socks4 connection to #{ip} not supported" unless ip.ipv4?
+
               packet << [ip.to_i].pack("N")
             rescue IPAddr::InvalidAddressError
               if parameters.uri.scheme == "socks4"
                 # resolv defaults to IPv4, and socks4 doesn't support IPv6 otherwise
                 ip = IPAddr.new(Resolv.getaddress(uri.host))