lib/bunny/ssl_socket.rb in bunny-0.9.0.pre13 vs lib/bunny/ssl_socket.rb in bunny-0.9.0.rc1

@@ -2,23 +2,47 @@
 
 module Bunny
   begin
     require "openssl"
 
+    # TLS-enabled TCP socket that implements convenience
+    # methods found in Bunny::Socket.
     class SSLSocket < OpenSSL::SSL::SSLSocket
+
+      # IO::WaitReadable is 1.9+ only
+      READ_RETRY_EXCEPTION_CLASSES = [Errno::EAGAIN, Errno::EWOULDBLOCK]
+      READ_RETRY_EXCEPTION_CLASSES << IO::WaitReadable if IO.const_defined?(:WaitReadable)
+
+
+      # Reads given number of bytes with an optional timeout
+      #
+      # @param [Integer] count How many bytes to read
+      # @param [Integer] timeout Timeout
+      #
+      # @return [String] Data read from the socket
+      # @api public
       def read_fully(count, timeout = nil)
-        return nil if @eof
+        return nil if @__bunny_socket_eof_flag__
 
         value = ''
         begin
           loop do
             value << read_nonblock(count - value.bytesize)
             break if value.bytesize >= count
           end
-        rescue EOFError
-          @eof = true
-        rescue Errno::EAGAIN, Errno::EWOULDBLOCK, OpenSSL::SSL::SSLError => e
-          puts e.inspect
+        rescue EOFError => e
+          @__bunny_socket_eof_flag__ = true
+        rescue OpenSSL::SSL::SSLError => e
+          if e.message == "read would block"
+            if IO.select([self], nil, nil, timeout)
+              retry
+            else
+              raise Timeout::Error, "IO timeout when reading #{count} bytes"
+            end
+          else
+            raise e
+          end
+        rescue *READ_RETRY_EXCEPTION_CLASSES => e
           if IO.select([self], nil, nil, timeout)
             retry
           else
             raise Timeout::Error, "IO timeout when reading #{count} bytes"
           end