Non-blocking Socket Connect

lib/net/ldap/connection.rb

@@ -1,3 +1,5 @@
+require 'timeout'
+
 # This is a private class used internally by the library. It should not
 # be called by user code.
 class Net::LDAP::Connection #:nodoc:
@@ -10,7 +12,7 @@ def initialize(server)
     @instrumentation_service = server[:instrumentation_service]
 
     begin
-      @conn = server[:socket] || TCPSocket.new(server[:host], server[:port])
+      @conn = server[:socket] || connect(server)
     rescue SocketError
       raise Net::LDAP::LdapError, "No such address or other socket error."
     rescue Errno::ECONNREFUSED
@@ -28,6 +30,54 @@ def initialize(server)
     yield self if block_given?
   end
 
+  # Internal: Connect to the host and port.
+  #
+  # Accepted options:
+  # - host: the hostname or IP address String of the server to connect to.
+  # - port: the port Integer to connect to.
+  # - timeout: the Integer seconds to wait before timing out, or `nil` for no
+  #            timeout (the default).
+  #
+  # This connects to the specified server using non-blocking socket methods.
+  # This allows us to specify our own timeout to prevent blocking.
+  #
+  # Returns a connected Socket object.
+  def connect(options)
+    host, port = options[:host], options[:port]
+    timeout    = options[:timeout]
+
+    addr = Socket.getaddrinfo(host, nil)
+    ip   = addr[0][3]
+    sockaddr = Socket.sockaddr_in(port, ip)
+
+    socket = Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0)
+    socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+    begin
+      # start connection attempt
+      socket.connect_nonblock(sockaddr)
+    rescue IO::WaitWritable
+      if IO.select(nil, [socket], nil, timeout)
+        # validate connection succeeded
+        begin
+          socket.connect_nonblock(sockaddr)
+        rescue Errno::EISCONN
+          # already connected, no action necessary
+        rescue Exception => e
+          # unexpected error
+          socket.close
+          raise Net::LDAP::LdapError, "Connection to #{host}:#{port} (#{ip}) failed: (#{e.class}) #{e.message}"
+        end
+      else
+        # connection timeout
+        socket.close
+        raise Net::LDAP::LdapError, "Connection to #{host}:#{port} (#{ip}) timed out."
+      end
+    end
+
+    socket
+  end
+
   module GetbyteForSSLSocket
     def getbyte
       getc.ord

test/integration/test_open.rb

@@ -1,6 +1,6 @@
 require_relative '../test_helper'
 
-class TestBindIntegration < LDAPIntegrationTestCase
+class TestOpenIntegration < LDAPIntegrationTestCase
   def test_binds_without_open
     events = @service.subscribe "bind.net_ldap_connection"
 

test/test_ldap_connection.rb

@@ -8,20 +8,40 @@ def test_unresponsive_host
   end
 
   def test_blocked_port
-    flexmock(TCPSocket).should_receive(:new).and_raise(SocketError)
+    flexmock(Socket).should_receive(:new).and_raise(SocketError)
     assert_raise Net::LDAP::LdapError do
       Net::LDAP::Connection.new(:host => 'test.mocked.com', :port => 636)
     end
   end
 
   def test_raises_unknown_exceptions
     error = Class.new(StandardError)
-    flexmock(TCPSocket).should_receive(:new).and_raise(error)
+    flexmock(Socket).should_receive(:new).and_raise(error)
     assert_raise error do
       Net::LDAP::Connection.new(:host => 'test.mocked.com', :port => 636)
     end
   end
 
+  def test_connection_timeout
+    socket = flexmock(:socket)
+    flexmock(Socket).should_receive(:new).and_return(socket)
+
+    # standard behavior
+    socket.should_receive(:setsockopt)
+    socket.should_receive(:close)
+
+    # indicate connection happening in background (wait to write)
+    socket.should_receive(:connect_nonblock).
+           and_raise(Errno::EINPROGRESS.new.extend(IO::WaitWritable))
+
+    # time out waiting to write
+    flexmock(IO).should_receive(:select).and_return(nil)
+
+    assert_raise Net::LDAP::LdapError do
+      Net::LDAP::Connection.new(:host => "123.123.123.123", :port => 389, :timeout => 1)
+    end
+  end
+
   def test_modify_ops_delete
     args = { :operations => [ [ :delete, "mail" ] ] }
     result = Net::LDAP::Connection.modify_ops(args[:operations])
@@ -259,8 +279,7 @@ class TestLDAPConnectionErrors < Test::Unit::TestCase
   def setup
     @tcp_socket = flexmock(:connection)
     @tcp_socket.should_receive(:write)
-    flexmock(TCPSocket).should_receive(:new).and_return(@tcp_socket)
-    @connection = Net::LDAP::Connection.new(:host => 'test.mocked.com', :port => 636)
+    @connection = Net::LDAP::Connection.new(:socket => @tcp_socket)
   end
 
   def test_error_failed_operation
@@ -288,12 +307,10 @@ class TestLDAPConnectionInstrumentation < Test::Unit::TestCase
   def setup
     @tcp_socket = flexmock(:connection)
     @tcp_socket.should_receive(:write)
-    flexmock(TCPSocket).should_receive(:new).and_return(@tcp_socket)
 
     @service = MockInstrumentationService.new
     @connection = Net::LDAP::Connection.new \
-      :host => 'test.mocked.com',
-      :port => 636,
+      :socket => @tcp_socket,
       :instrumentation_service => @service
   end