Add support for timeouts on read and write operations

Author: John Wedoff

Contributors.rdoc

@@ -23,3 +23,4 @@ Contributions since:
 * Cody Cutrer (ccutrer)
 * WoodsBagotAndreMarquesLee
 * Rufus Post (mynameisrufus)
+* Akamai Technologies, Inc. (jwedoff)

lib/net/ber/ber_parser.rb

@@ -136,7 +136,7 @@ def parse_ber_object(syntax, id, data)
   # invalid BER length case. Because the "lengthlength" value was not used
   # inside of #read_ber, we no longer return it.
   def read_ber_length
-    n = getbyte_nonblock
+    n = ber_timeout_getbyte
 
     if n <= 0x7f
       n
@@ -146,7 +146,7 @@ def read_ber_length
       raise Net::BER::BerError, "Invalid BER length 0xFF detected."
     else
       v = 0
-      read_ber_nonblock(n & 0x7f).each_byte do |b|
+      ber_timeout_read(n & 0x7f).each_byte do |b|
         v = (v << 8) + b
       end
       v
@@ -177,45 +177,126 @@ def read_ber(syntax = nil)
       raise Net::BER::BerError,
             "Indeterminite BER content length not implemented."
     end
-    data = read_ber_nonblock(content_length)
+    data = ber_timeout_read(content_length)
 
     parse_ber_object(syntax, id, data)
   end
 
   # Internal: Returns the BER message ID or nil.
   def read_ber_id
-    getbyte_nonblock
+    ber_timeout_getbyte
   end
   private :read_ber_id
 
+  # Internal: specify the BER socket read timeouts, nil by default (no timeout).
+  attr_accessor :ber_io_deadline
+  private :ber_io_deadline
+
+  ##
+  # sets a timeout of timeout seconds for read_ber and ber_timeout_write operations in the provided block the proin the future for if there is not already a earlier deadline set
+  def with_timeout(timeout)
+    timeout = timeout.to_f
+    # don't change deadline if run without timeout
+    return yield if timeout <= 0
+    # clear deadline if it is not in the future
+    self.ber_io_deadline = nil unless ber_io_timeout&.send(:>, 0)
+    new_deadline = Time.now + timeout
+    # don't add deadline if current deadline is shorter
+    return yield if ber_io_deadline && ber_io_deadline < new_deadline
+    old_deadline = ber_io_deadline
+    begin
+      self.ber_io_deadline = new_deadline
+      yield
+    ensure
+      self.ber_io_deadline = old_deadline
+    end
+  end
+
+  # seconds until ber_io_deadline
+  def ber_io_timeout
+    ber_io_deadline ? ber_io_deadline - Time.now : nil
+  end
+  private :ber_io_timeout
+
+  def read_select!
+    return if IO.select([self], nil, nil, ber_io_timeout)
+    raise Net::LDAP::LdapError, "Timed out reading from the socket"
+  end
+  private :read_select!
+
+  def write_select!
+    return if IO.select(nil, [self], nil, ber_io_timeout)
+    raise Net::LDAP::LdapError, "Timed out reading from the socket"
+  end
+  private :write_select!
+
   # Internal: Replaces `getbyte` with nonblocking implementation.
-  def getbyte_nonblock
+  def ber_timeout_getbyte
     begin
       read_nonblock(1).ord
     rescue IO::WaitReadable
-      if IO.select([self], nil, nil, read_ber_timeout)
-        read_nonblock(1).ord
-      else
-        raise Net::LDAP::LdapError, "Timed out reading from the socket"
-      end
+      read_select!
+      retry
+    rescue IO::WaitWritable
+      write_select!
+      retry
+    rescue EOFError
+      # nothing to read on the socket (StringIO)
+      nil
     end
-  rescue EOFError
-    # nothing to read on the socket (StringIO)
-    nil
   end
-  private :getbyte_nonblock
+  private :ber_timeout_getbyte
 
   # Internal: Read `len` bytes, respecting timeout.
-  def read_ber_nonblock(len)
+  def ber_timeout_read(len)
+    buffer ||= ''.force_encoding(Encoding::ASCII_8BIT)
     begin
-      read_nonblock(len)
-    rescue IO::WaitReadable
-      if IO.select([self], nil, nil, read_ber_timeout)
-        read_nonblock(len)
-      else
-        raise Net::LDAP::LdapError, "Timed out reading from the socket"
+      read_nonblock(len, buffer)
+      return buffer if buffer.bytesize >= len
+    rescue IO::WaitReadable, IO::WaitWritable
+      buffer.clear
+    rescue EOFError
+      # nothing to read on the socket (StringIO)
+      nil
+    end
+    block ||= ''.force_encoding(Encoding::ASCII_8BIT)
+    len -= buffer.bytesize
+    loop do
+      begin
+        read_nonblock(len, block)
+      rescue IO::WaitReadable
+        read_select!
+        retry
+      rescue IO::WaitWritable
+        write_select!
+        retry
+      rescue EOFError
+        return buffer.empty? ? nil : buffer
+      end
+      buffer << block
+      len -= block.bytesize
+      return buffer if len <= 0
+    end
+  end
+  private :ber_timeout_read
+
+  ##
+  # Writes val as a plain write would, but respecting the dealine set by with_timeout
+  def ber_timeout_write(val)
+    total_written = 0
+    while 0 < val.bytesize
+      begin
+        written = write_nonblock(val)
+      rescue IO::WaitReadable
+        read_select!
+        retry
+      rescue IO::WaitWritable
+        write_select!
+        retry
       end
+      total_written += written
+      val = val.byteslice(written..-1)
     end
+    total_written
   end
-  private :read_ber_nonblock
 end

lib/net/ldap.rb

@@ -553,6 +553,7 @@ def initialize(args = {})
     @force_no_page = args[:force_no_page] || DefaultForceNoPage
     @encryption = normalize_encryption(args[:encryption]) # may be nil
     @connect_timeout = args[:connect_timeout]
+    @io_timeout = args[:io_timeout]
 
     if pr = @auth[:password] and pr.respond_to?(:call)
       @auth[:password] = pr.call
@@ -1293,14 +1294,19 @@ def connection=(connection)
   # result from that, and :use_connection: will not yield at all. If not
   # the return value is whatever is returned from the block.
   def use_connection(args)
+    timeout_args = args.slice(:io_timeout).values
     if @open_connection
-      yield @open_connection
+      @open_connection.with_timeout(*timeout_args) do
+        yield(@open_connection)
+      end
     else
       begin
         conn = new_connection
-        result = conn.bind(args[:auth] || @auth)
-        return result unless result.result_code == Net::LDAP::ResultCodeSuccess
-        yield conn
+        conn.with_timeout(*timeout_args) do
+          result = conn.bind(args[:auth] || @auth)
+          return result unless result.result_code == Net::LDAP::ResultCodeSuccess
+          yield(conn)
+        end
       ensure
         conn.close if conn
       end
@@ -1315,7 +1321,8 @@ def new_connection
       :hosts                   => @hosts,
       :encryption              => @encryption,
       :instrumentation_service => @instrumentation_service,
-      :connect_timeout         => @connect_timeout
+      :connect_timeout         => @connect_timeout,
+      :io_timeout              => @io_timeout
 
     # Force connect to see if there's a connection error
     connection.socket