Merge pull request #57 from lizzypy/fix-tcp-socket-with-connect-timeout

Fix tcp socket with connect timeout

Author: MatzFan

lib/socksify/tcpsocket.rb

@@ -10,16 +10,16 @@ class TCPSocket
 
   # See http://tools.ietf.org/html/rfc1928
   # rubocop:disable Metrics/ParameterLists
-  def initialize(host = nil, port = nil, local_host = nil, local_port = nil)
+  def initialize(host = nil, port = nil, local_host = nil, local_port = nil, **kwargs)
     socks_peer = host if host.is_a?(SOCKSConnectionPeerAddress)
     socks_server = set_socks_server(socks_peer)
     socks_port = set_socks_port(socks_peer)
     socks_ignores = set_socks_ignores(socks_peer)
     host = socks_peer.peer_host if socks_peer
     if socks_server && socks_port && !socks_ignores.include?(host)
-      make_socks_connection(host, port, socks_server, socks_port)
+      make_socks_connection(host, port, socks_server, socks_port, **kwargs)
     else
-      make_direct_connection(host, port, local_host, local_port)
+      make_direct_connection(host, port, local_host, local_port, **kwargs)
     end
   end
   # rubocop:enable Metrics/ParameterLists
@@ -57,16 +57,16 @@ def set_socks_ignores(socks_peer = nil)
     socks_peer ? [] : self.class.socks_ignores
   end
 
-  def make_socks_connection(host, port, socks_server, socks_port)
+  def make_socks_connection(host, port, socks_server, socks_port, **kwargs)
     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/
     socks_connect(host, port) if host
   end
 
-  def make_direct_connection(host, port, local_host, local_port)
+  def make_direct_connection(host, port, local_host, local_port, **kwargs)
     Socksify.debug_notice "Connecting directly to #{host}:#{port}"
-    initialize_tcp host, port, local_host, local_port
+    initialize_tcp host, port, local_host, local_port, **kwargs
     Socksify.debug_debug "Connected to #{host}:#{port}"
   end
 end

test/test_tcpsocket.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require_relative 'test_helper'
+
+class TCPSocketTest < Minitest::Test
+  include HelperMethods
+
+  if RUBY_VERSION.to_f >= 3.0
+    def test_tcp_socket_direct_connection_with_connection_timeout
+      disable_socks
+
+      socket = TCPSocket.new('127.0.0.1', 9050, connect_timeout: 0.1)
+
+      assert !socket.closed?
+    end
+
+    def test_tcp_socket_socks_connection_with_connection_timeout
+      enable_socks
+
+      # leave off the host because we don't need to worry about connecting to socks
+      socket = TCPSocket.new(connect_timeout: 0.1)
+
+      assert !socket.closed?
+    end
+  end
+
+  def test_tcp_socket_direct_connection_with_connection_timeout_no_kwargs
+    disable_socks
+
+    socket = TCPSocket.new('127.0.0.1', 9050)
+
+    assert !socket.closed?
+  end
+end