arachni-rpc 0.1.3 → 0.2.0

data/lib/arachni/rpc/client/handler.rb

@@ -0,0 +1,167 @@
+=begin
+
+    This file is part of the Arachni-RPC project and may be subject to
+    redistribution and commercial restrictions. Please see the Arachni-RPC EM
+    web site for more information on licensing and terms of use.
+
+=end
+
+module Arachni
+module RPC
+class Client
+
+# Transmits {Request} objects and calls callbacks once an {Response} is received.
+#
+# @author Tasos "Zapotek" Laskos <tasos.laskos@arachni-scanner.com>
+class Handler < Reactor::Connection
+    include Protocol
+
+    # Default amount of tries for failed requests.
+    DEFAULT_TRIES = 9
+
+    # @return   [Symbol]    Status of the connection, can be:
+    #
+    # * `:idle` -- Just initialized.
+    # * `:ready` -- A connection has been established.
+    # * `:pending` -- Sending request and awaiting response.
+    # * `:done` -- Response received and callback invoked -- ready to be reused.
+    # * `:closed` -- Connection closed.
+    attr_reader :status
+
+    # @return   [Exceptions::ConnectionError]
+    attr_reader :error
+
+    # Prepares an RPC connection and sets {#status} to `:idle`.
+    #
+    # @param    [Hash]  opts
+    # @option   opts    [Integer]   :max_retries    (9)
+    #   Default amount of tries for failed requests.
+    #
+    # @option   opts    [Client]   :base
+    #   Client instance needed to {Client#push_connection push} ourselves
+    #   back to its connection pool once we're done and we're ready to be reused.
+    def initialize( opts )
+        @opts = opts.dup
+
+        @max_retries = @opts[:max_retries] || DEFAULT_TRIES
+        @client      = @opts[:client]
+
+        @opts[:tries] ||= 0
+        @tries ||= @opts[:tries]
+
+        @status  = :idle
+        @request = nil
+    end
+
+    # Sends an RPC request (i.e. performs an RPC call) and sets {#status}
+    # to `:pending`.
+    #
+    # @param    [Request]      req
+    def send_request( req )
+        @request = req
+        @status  = :pending
+        super( req )
+    end
+
+    # @note Pushes itself to the client's connection pool to be re-used.
+    #
+    # Handles responses to RPC requests, calls its callback and sets {#status}
+    # to `:done`.
+    #
+    # @param    [Arachni::RPC::Response]    res
+    def receive_response( res )
+        if res.exception?
+            res.obj = Exceptions.from_response( res )
+        end
+
+        @request.callback.call( res.obj ) if @request.callback
+    ensure
+        @request = nil # Help the GC out.
+        @error   = nil # Help the GC out.
+        @status  = :done
+
+        @opts[:tries] = @tries = 0
+        @client.push_connection self
+    end
+
+    # Handles closed connections, cleans up the SSL session, retries (if
+    # necessary) and sets {#status} to `:closed`.
+    #
+    # @private
+    def on_close( reason )
+         if @request
+             # If there is a request and a callback and the callback hasn't yet be
+             # called (i.e. not done) then we got here by error so retry.
+             if @request && @request.callback && !done?
+                 if retry?
+                     retry_request
+                 else
+                     @error = e = Exceptions::ConnectionError.new( "Connection closed [#{reason}]" )
+                     @request.callback.call( e )
+                     @client.connection_failed self
+                 end
+
+                 return
+             end
+         else
+             @error = reason
+             @client.connection_failed self
+         end
+
+        close_without_retry
+    end
+
+    # @note If `true`, the connection can be re-used.
+    #
+    # @return   [Boolean]
+    #   `true` when the connection is done, `false` otherwise.
+    def done?
+        @status == :done
+    end
+
+    # Closes the connection without triggering a retry operation and sets
+    # {#status} to `:closed`.
+    def close_without_retry
+        @request = nil
+        @status  = :closed
+        close_without_callback
+    end
+
+    private
+
+    # Converts incoming hash objects to {Response} objects and calls
+    # {#receive_response}.
+    #
+    # @param    [Hash]      obj
+    def receive_object( obj )
+        receive_response( Response.new( obj ) )
+    end
+
+    def retry_request
+        opts = @opts.dup
+        opts[:tries] += 1
+
+        req = @request.dup
+
+        # The connection will be detached soon, keep a separate reference to
+        # the reactor.
+        reactor = @reactor
+
+        @tries += 1
+        reactor.delay( 0.2 ) do
+            address = opts[:socket] ? opts[:socket] : [opts[:host], opts[:port]]
+            reactor.connect( *[address, self.class, opts ].flatten ).send_request( req )
+        end
+
+        close_without_retry
+    end
+
+    def retry?
+        @tries < @max_retries
+    end
+
+end
+
+end
+end
+end

data/lib/arachni/rpc/exceptions.rb

@@ -6,7 +6,6 @@
 
 =end
 
-#
 # RPC Exceptions have methods that help identify them based on type.
 #
 # So in order to allow evaluations like:
@@ -16,11 +15,10 @@
 # to be possible on all objects these helper methods need to be available for
 # all objects.
 #
-# By default they'll return false, individual RPC Exceptions will overwrite them to
-# return true when applicable.
-#
-# @author: Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+# By default they'll return false, individual RPC Exceptions will overwrite them
+# to return true when applicable.
 #
+# @author Tasos "Zapotek" Laskos <tasos.laskos@arachni-scanner.com>
 class Object
 
     # @return   [Bool]  false
@@ -48,6 +46,11 @@ class Object
         false
     end
 
+    # @return   [Bool]  true
+    def rpc_ssl_error?
+        false
+    end
+
     # @return   [Bool]  false
     def rpc_exception?
         false
@@ -59,114 +62,87 @@ module Arachni
 module RPC
 module Exceptions
 
-    #
     # Returns an exception based on the response object.
     #
     # @param    [Arachni::RPC::Response]    response
     #
     # @return   [Exception]
-    #
     def self.from_response( response )
-        obj = response.obj
-        klass = Arachni::RPC::Exceptions.const_get( obj['type'].to_sym )
-        e = klass.new( obj['exception'] )
-        e.set_backtrace( obj['backtrace'] )
+        exception = response.exception
+        klass = Arachni::RPC::Exceptions.const_get( exception['type'].to_sym )
+        e = klass.new( exception['message'] )
+        e.set_backtrace( exception['backtrace'] )
         e
     end
 
     class Base < ::RuntimeError
-        #
+
         # @return   [Bool]  true
-        #
         def rpc_exception?
             true
         end
     end
 
-    #
     # Signifies an abruptly terminated connection.
     #
     # Look for network or SSL errors or a dead server or a mistyped server address/port.
-    #
     class ConnectionError < Base
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_connection_error?
             true
         end
     end
 
-    #
     # Signifies an exception that occurred on the server-side.
     #
     # Look errors on the remote method and review the server output for more details.
-    #
     class RemoteException < Base
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_remote_exception?
             true
         end
     end
 
-    #
     # An invalid object has been called.
     #
     # Make sure that there is a server-side handler for the object you called.
-    #
     class InvalidObject < Base
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_invalid_object_error?
             true
         end
 
     end
 
-    #
     # An invalid method has been called.
     #
     # Occurs when a remote method doesn't exist or isn't public.
-    #
     class InvalidMethod < Base
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_invalid_method_error?
             true
         end
 
     end
 
-    #
     # Signifies an authentication token mismatch between the client and the server.
-    #
     class InvalidToken  < Base
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_invalid_token_error?
             true
         end
 
     end
 
-    #
     # Signifies an authentication token mismatch between the client and the server.
-    #
-    class SSLPeerVerificationFailed  < ConnectionError
+    class SSLPeerVerificationFailed < ConnectionError
 
-        #
         # @return   [Bool]  true
-        #
         def rpc_ssl_error?
             true
         end

data/lib/arachni/rpc/message.rb

@@ -9,27 +9,22 @@
 module Arachni
 module RPC
 
-#
 # Represents an RPC message, serves as the basis for {Request} and {Response}.
 #
-# @author: Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
-#
+# @author Tasos "Zapotek" Laskos <tasos.laskos@arachni-scanner.com>
 class Message
 
-    #
-    # @param    [Hash]   opts   sets instance attributes
-    #
+    # @param    [Hash]   opts
+    #   Sets instance attributes.
     def initialize( opts = {} )
-        opts.each_pair { |k, v| instance_variable_set( "@#{k}".to_sym, v ) }
+        opts.each_pair { |k, v| send( "#{k}=".to_sym, v ) }
     end
 
-    #
     # Merges the attributes of another message with self.
     #
     # (The param doesn't *really* have to be a message, any object will do.)
     #
     # @param    [Message]   message
-    #
     def merge!( message )
         message.instance_variables.each do |var|
             val = message.instance_variable_get( var )
@@ -37,14 +32,12 @@ class Message
         end
     end
 
-    #
-    # Prepares the message for transmission (i.e. converts the message to a Hash).
+    # Prepares the message for transmission (i.e. converts the message to a `Hash`).
     #
     # Attributes that should not be included can be skipped by implementing
     # {#transmit?} and returning the appropriate value.
     #
     # @return   [Hash]
-    #
     def prepare_for_tx
         instance_variables.inject({}) do |h, k|
             h[normalize( k )] = instance_variable_get( k ) if transmit?( k )
@@ -52,11 +45,10 @@ class Message
         end
     end
 
-    #
     # Decides which attributes should be skipped by {#prepare_for_tx}.
     #
-    # @param    [Symbol]    attr    attribute symbol (i.e. :@token)
-    #
+    # @param    [Symbol]    attr
+    #   Instance variable symbol (i.e. `:@token`).
     def transmit?( attr )
         true
     end

data/lib/arachni/rpc/protocol.rb

@@ -0,0 +1,103 @@
+=begin
+
+    This file is part of the Arachni-RPC project and may be subject to
+    redistribution and commercial restrictions. Please see the Arachni-RPC EM
+    web site for more information on licensing and terms of use.
+
+=end
+
+module Arachni
+module RPC
+
+# Provides helper transport methods for {Message} transmission.
+#
+# @author Tasos "Zapotek" Laskos <tasos.laskos@arachni-scanner.com>
+module Protocol
+    include Reactor::Connection::TLS
+
+    # Initializes an SSL session once the connection has been established and
+    # sets {#status} to `:ready`.
+    #
+    # @private
+    def on_connect
+        start_tls(
+            ca:          @opts[:ssl_ca],
+            private_key: @opts[:ssl_pkey],
+            certificate: @opts[:ssl_cert]
+        )
+
+        @status = :ready
+    end
+
+    # @param    [Message]    msg
+    #   Message to send to the peer.
+    def send_message( msg )
+        send_object( msg.prepare_for_tx )
+    end
+    alias :send_request  :send_message
+    alias :send_response :send_message
+
+    # Receives data from the network.
+    #
+    # Rhe data will be chunks of a serialized object which will be buffered
+    # until the whole transmission has finished.
+    #
+    # It will then unserialize it and pass it to {#receive_object}.
+    def on_read( data )
+        (@buf ||= '') << data
+
+        while @buf.size >= 4
+            if @buf.size >= 4 + ( size = @buf.unpack( 'N' ).first )
+                @buf.slice!( 0, 4 )
+                receive_object( unserialize( @buf.slice!( 0, size ) ) )
+            else
+                break
+            end
+        end
+    end
+
+    private
+
+    # Stub method, should be implemented by servers.
+    #
+    # @param    [Request]     request
+    # @abstract
+    def receive_request( request )
+        p request
+    end
+
+    # Stub method, should be implemented by clients.
+    #
+    # @param    [Response]    response
+    # @abstract
+    def receive_response( response )
+        p response
+    end
+
+    #   Object to send.
+    def send_object( obj )
+        data = serialize( obj )
+        write [data.bytesize, data].pack( 'Na*' )
+    end
+
+    # Returns the preferred serializer based on the `serializer` option of the
+    # server.
+    #
+    # @return   [.load, .dump]
+    #   Serializer to be used (Defaults to `YAML`).
+    def serializer
+        @opts[:serializer] || YAML
+    end
+
+    def serialize( obj )
+        serializer.dump obj
+    end
+
+    def unserialize( obj )
+        serializer.load( obj )
+    end
+
+end
+
+end
+end