From 1020a98e55c033a6fe693d85c0ef0285f09e81f3 Mon Sep 17 00:00:00 2001
From: Petr Pulc <petrpulc@gmail.com>
Date: Mon, 22 Dec 2014 03:14:37 +0100
Subject: [PATCH 1/2] Better timeout handling on sockets

---
 lib/net/ldap/connection.rb | 60 +++++++++++++++++++++++++++-----------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/lib/net/ldap/connection.rb b/lib/net/ldap/connection.rb
index 0e3d7e05..112030f2 100644
--- a/lib/net/ldap/connection.rb
+++ b/lib/net/ldap/connection.rb
@@ -6,26 +6,52 @@ class Net::LDAP::Connection #:nodoc:
   LdapVersion = 3
   MaxSaslChallenges = 10
 
-  def initialize(server)
+  def initialize(server, timeout = 5)
     @instrumentation_service = server[:instrumentation_service]
+    
+    if server[:socket]
+      @conn = server[:socket]
+    else
+      addr = Socket.getaddrinfo(server[:host], nil)
+      sockaddr = Socket.pack_sockaddr_in(server[:port], addr[0][3])
+
+      Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0).tap do |socket|
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        begin
+          socket.connect_nonblock(sockaddr)
+        rescue IO::WaitWritable  
+          if IO.select(nil, [socket], nil, timeout)
+            begin
+              @conn = socket.connect_nonblock(sockaddr)
+            rescue Errno::EISCONN
+              # Good news everybody, the socket is connected!
+            rescue SocketError
+              raise Net::LDAP::LdapError, "No such address or other socket error."
+            rescue Errno::ECONNREFUSED
+              raise Net::LDAP::LdapError, "Server #{server[:host]} refused connection on port #{server[:port]}."
+            rescue Errno::EHOSTUNREACH => error
+              raise Net::LDAP::LdapError, "Host #{server[:host]} was unreachable (#{error.message})"
+            rescue
+              # An unexpected exception was raised - the connection is no good.
+              socket.close
+              raise
+            end
+          else
+            # IO.select returns nil when the socket is not ready before timeout 
+            # seconds have elapsed
+            socket.close
+            raise "Connection timeout"
+          end
+        end
+      end
+      
+      if server[:encryption]
+        setup_encryption server[:encryption]
+      end
 
-    begin
-      @conn = server[:socket] || TCPSocket.new(server[:host], server[:port])
-    rescue SocketError
-      raise Net::LDAP::LdapError, "No such address or other socket error."
-    rescue Errno::ECONNREFUSED
-      raise Net::LDAP::LdapError, "Server #{server[:host]} refused connection on port #{server[:port]}."
-    rescue Errno::EHOSTUNREACH => error
-      raise Net::LDAP::LdapError, "Host #{server[:host]} was unreachable (#{error.message})"
-    rescue Errno::ETIMEDOUT
-      raise Net::LDAP::LdapError, "Connection to #{server[:host]} timed out."
-    end
-
-    if server[:encryption]
-      setup_encryption server[:encryption]
+      yield self if block_given?
     end
-
-    yield self if block_given?
   end
 
   module GetbyteForSSLSocket

From 933d5259ff2ec21fe3a9d8c7c92fd15387cb1549 Mon Sep 17 00:00:00 2001
From: Petr Pulc <petrpulc@gmail.com>
Date: Mon, 22 Dec 2014 03:41:23 +0100
Subject: [PATCH 2/2] Minor fixes of cleanup

---
 lib/net/ldap/connection.rb | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/lib/net/ldap/connection.rb b/lib/net/ldap/connection.rb
index 112030f2..ee2cc8a3 100644
--- a/lib/net/ldap/connection.rb
+++ b/lib/net/ldap/connection.rb
@@ -16,14 +16,15 @@ def initialize(server, timeout = 5)
       sockaddr = Socket.pack_sockaddr_in(server[:port], addr[0][3])
 
       Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0).tap do |socket|
-        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        @conn = socket
+        @conn.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
 
         begin
-          socket.connect_nonblock(sockaddr)
+          @conn.connect_nonblock(sockaddr)
         rescue IO::WaitWritable  
           if IO.select(nil, [socket], nil, timeout)
             begin
-              @conn = socket.connect_nonblock(sockaddr)
+              @conn.connect_nonblock(sockaddr)
             rescue Errno::EISCONN
               # Good news everybody, the socket is connected!
             rescue SocketError
@@ -34,13 +35,11 @@ def initialize(server, timeout = 5)
               raise Net::LDAP::LdapError, "Host #{server[:host]} was unreachable (#{error.message})"
             rescue
               # An unexpected exception was raised - the connection is no good.
-              socket.close
               raise
             end
           else
             # IO.select returns nil when the socket is not ready before timeout 
             # seconds have elapsed
-            socket.close
             raise "Connection timeout"
           end
         end