lib/arachni/reactor/connection.rb in arachni-reactor-0.1.0.beta5 vs lib/arachni/reactor/connection.rb in arachni-reactor-0.1.0

@@ -5,20 +5,19 @@
     web site for more information on licensing and terms of use.
 
 =end
 
 require_relative 'connection/error'
-require_relative 'connection/peer_info'
 require_relative 'connection/callbacks'
+require_relative 'connection/peer_info'
 require_relative 'connection/tls'
 
 module Arachni
 class Reactor
 
 # @author Tasos "Zapotek" Laskos <[e-mail]>
 class Connection
-    include PeerInfo
     include Callbacks
 
     # Maximum amount of data to be written or read at a time.
     #
     # We set this to the same max block size as the OpenSSL buffers because more
@@ -174,10 +173,96 @@
         on_close reason
         close_without_callback
         nil
     end
 
+    # Accepts a new client connection.
+    #
+    # @return   [Connection, nil]
+    #   New connection or `nil` if the socket isn't ready to accept new
+    #   connections yet.
+    #
+    # @private
+    def accept
+        return if !(accepted = socket_accept)
+
+        connection = @server_handler.call
+        connection.configure socket: accepted, role: :server
+        @reactor.attach connection
+        connection
+    end
+
+    # @param    [Socket]    socket
+    #   Ruby `Socket` associated with this connection.
+    # @param    [Symbol]    role
+    #   `:server` or `:client`.
+    # @param    [Block]    server_handler
+    #   Block that generates a handler as specified in {Reactor#listen}.
+    #
+    # @private
+    def configure( options = {} )
+        @socket         = options[:socket]
+        @role           = options[:role]
+        @host           = options[:host]
+        @port           = options[:port]
+        @server_handler = options[:server_handler]
+
+        # If we're a server without a handler then we're an accepted connection.
+        if unix? || role == :server
+            @connected = true
+            on_connect
+        end
+
+        nil
+    end
+
+    def connected?
+        !!@connected
+    end
+
+    # @private
+    def _connect
+        return true if unix? || connected?
+
+        begin
+            Error.translate do
+                socket.connect_nonblock( Socket.sockaddr_in( @port, @host ) )
+            end
+        # Already connected. :)
+        rescue Errno::EISCONN
+        end
+
+        @connected = true
+        on_connect
+
+        true
+    rescue IO::WaitReadable, IO::WaitWritable, Errno::EINPROGRESS
+    rescue Error => e
+        close e
+    end
+
+    # @note If this is a server {#listener?} it will delegate to {#accept}.
+    # @note If this is a normal socket it will read {BLOCK_SIZE} amount of data.
+    #   and pass it to {#on_read}.
+    #
+    # Processes a `read` event for this connection.
+    #
+    # @private
+    def _read
+        return _connect if !listener? && !connected?
+        return accept   if listener?
+
+        Error.translate do
+            on_read @socket.read_nonblock( BLOCK_SIZE )
+        end
+
+    # Not ready to read or write yet, we'll catch it on future Reactor ticks.
+    rescue IO::WaitReadable, IO::WaitWritable
+    rescue Error => e
+        close e
+    end
+
     # @note Will call {#on_write} every time any of the buffer is consumed,
     #   can be multiple times when performing partial writes.
     # @note Will call {#on_flush} once all of the buffer has been consumed.
     #
     # Processes a `write` event for this connection.
@@ -188,16 +273,18 @@
     # @return   [Integer]
     #   Amount of the buffer consumed.
     #
     # @private
     def _write
+        return _connect if !connected?
+
         chunk = write_buffer.slice( 0, BLOCK_SIZE )
         total_written = 0
 
         begin
             Error.translate do
-                # Send out the buffer, **all** of it, or at least try to.
+                # Send out the chunk, **all** of it, or at least try to.
                 loop do
                     total_written += written = @socket.write_nonblock( chunk )
                     write_buffer.slice!( 0, written )
 
                     # Call #on_write every time any of the buffer is consumed.
@@ -218,63 +305,9 @@
         end
 
         total_written
     rescue Error => e
         close e
-    end
-
-    # @note If this is a server {#listener?} it will delegate to {#accept}.
-    # @note If this is a normal socket it will read {BLOCK_SIZE} amount of data.
-    #   and pass it to {#on_read}.
-    #
-    # Processes a `read` event for this connection.
-    #
-    # @private
-    def _read
-        return accept if listener?
-
-        Error.translate do
-            on_read @socket.read_nonblock( BLOCK_SIZE )
-        end
-
-    # Not ready to read or write yet, we'll catch it on future Reactor ticks.
-    rescue IO::WaitReadable, IO::WaitWritable
-    rescue Error => e
-        close e
-    end
-
-    # Accepts a new client connection.
-    #
-    # @return   [Connection, nil]
-    #   New connection or `nil` if the socket isn't ready to accept new
-    #   connections yet.
-    #
-    # @private
-    def accept
-        return if !(accepted = socket_accept)
-
-        connection = @server_handler.call
-        connection.configure accepted, :server
-        @reactor.attach connection
-        connection
-    end
-
-    # @param    [Socket]    socket
-    #   Ruby `Socket` associated with this connection.
-    # @param    [Symbol]    role
-    #   `:server` or `:client`.
-    # @param    [Block]    server_handler
-    #   Block that generates a handler as specified in {Reactor#listen}.
-    #
-    # @private
-    def configure( socket, role, server_handler = nil )
-        @socket         = socket
-        @role           = role
-        @server_handler = server_handler
-
-        on_connect
-
-        nil
     end
 
     private
 
     def write_buffer