diff --git a/lib/socksify.rb b/lib/socksify.rb index f9f1980..b3f3369 100644 --- a/lib/socksify.rb +++ b/lib/socksify.rb @@ -151,7 +151,7 @@ def peer_host alias :initialize_tcp :initialize # See http://tools.ietf.org/html/rfc1928 - def initialize(host=nil, port=0, local_host=nil, local_port=nil) + def initialize(host=nil, port=0, local_host=nil, local_port=nil, **kwargs) if host.is_a?(SOCKSConnectionPeerAddress) socks_peer = host socks_server = socks_peer.socks_server @@ -166,12 +166,12 @@ def initialize(host=nil, port=0, local_host=nil, local_port=nil) if socks_server and socks_port and not socks_ignores.include?(host) Socksify::debug_notice "Connecting to SOCKS server #{socks_server}:#{socks_port}" - initialize_tcp socks_server, socks_port + initialize_tcp socks_server, socks_port, **kwargs socks_authenticate unless @@socks_version =~ /^4/ if host - socks_connect(host, port) + socks_connect(host, port, kwargs[:connect_timeout]) end else Socksify::debug_notice "Connecting directly to #{host}:#{port}" @@ -216,7 +216,7 @@ def socks_authenticate end # Connect - def socks_connect(host, port) + def socks_connect(host, port, connect_timeout = nil) Socksify::debug_debug "Sending destination address" write TCPSOCKSSocket.socks_version Socksify::debug_debug TCPSOCKSSocket.socks_version.unpack "H*" @@ -252,15 +252,31 @@ def socks_connect(host, port) end write [port].pack('n') if @@socks_version == "5" - socks_receive_reply + socks_receive_reply(connect_timeout) Socksify::debug_notice "Connected to #{host}:#{port} over SOCKS" end # returns [bind_addr: String, bind_port: Fixnum] - def socks_receive_reply + def socks_receive_reply(connect_timeout = nil) Socksify::debug_debug "Waiting for SOCKS reply" if @@socks_version == "5" - connect_reply = recv(4) + connect_reply = '' + + if connect_timeout + begin + connect_reply = recv_nonblock(4) + rescue IO::WaitReadable + if IO.select([self], [],[self], connect_timeout) + Socksify::debug_debug("retrying recv_nonblock") + retry + else # IO.select returns nil when timeout is hit + Socksify::debug_debug("connect_timeout hit") + raise SOCKSError.new("Timed out wating for connection. (connect_timeout: #{connect_timeout})") + end + end + else + connect_reply = recv(4) + end Socksify::debug_debug connect_reply.unpack "H*" if connect_reply[0..0] != "\005" raise SOCKSError.new("SOCKS version #{connect_reply[0..0]} is not 5")