arachni-rpc-em 0.1

data/lib/arachni/rpc/em/connection_utilities.rb

@@ -0,0 +1,44 @@
+=begin
+                  Arachni-RPC
+  Copyright (c) 2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+
+  This is free software; you can copy and distribute and modify
+  this program under the term of the GPL v2.0 License
+  (See LICENSE file for details)
+
+=end
+
+module Arachni
+module RPC
+module EM
+
+#
+# Helper methods to be included in EventMachine::Connection classes
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1
+#
+module ConnectionUtilities
+
+    #
+    # @return   [String]    IP address of the client
+    #
+    def peer_ip_addr
+        begin
+            if peername = get_peername
+                Socket.unpack_sockaddr_in( peername )[1]
+            else
+                'n/a'
+            end
+        rescue
+            'n/a'
+        end
+    end
+
+end
+
+end
+end
+end

data/lib/arachni/rpc/em/em.rb

@@ -0,0 +1,105 @@
+=begin
+                  Arachni-RPC
+  Copyright (c) 2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+
+  This is free software; you can copy and distribute and modify
+  this program under the term of the GPL v2.0 License
+  (See LICENSE file for details)
+
+=end
+
+module Arachni
+module RPC
+
+#
+# Provides some convenient methods for EventMachine's Reactor.
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1
+#
+module EM
+
+module Synchrony
+
+    def run( &block )
+        @@root_f = Fiber.new {
+            block.call
+        }.resume
+    end
+
+    extend self
+
+end
+
+    #
+    # Inits method variables for the Reactor tasks and its Mutex.
+    #
+    def init
+        @@reactor_tasks_mutex ||= Mutex.new
+        @@reactor_tasks ||= []
+    end
+
+    #
+    # Adds a block in the Reactor.
+    #
+    # @param    [Proc]    &block    block to be included in the Reactor loop
+    #
+    def add_to_reactor( &block )
+
+        self.init
+
+        # if we're already in the Reactor thread just run the block straight up.
+        if ::EM::reactor_thread?
+            block.call
+        else
+            @@reactor_tasks_mutex.lock
+            @@reactor_tasks << block
+
+            ensure_em_running!
+            @@reactor_tasks_mutex.unlock
+        end
+
+    end
+
+    #
+    # Blocks until the Reactor stops running
+    #
+    def block!
+        # beware of deadlocks, we can't join our own thread
+        ::EM.reactor_thread.join if ::EM.reactor_thread && !::EM::reactor_thread?
+    end
+
+    #
+    # Puts the Reactor in its own thread and runs it.
+    #
+    # It also runs all blocks sent to {#add_to_reactor}.
+    #
+    def ensure_em_running!
+        self.init
+
+        if !::EM::reactor_running?
+            q = Queue.new
+
+            Thread.new do
+                ::EM::run do
+
+                    ::EM.error_handler do |e|
+                        $stderr.puts "Exception raised during event loop: " +
+                        "#{e.message} (#{e.class})\n#{(e.backtrace ||
+                            [])[0..5].join("\n")}"
+                    end
+
+                    @@reactor_tasks.each { |task| task.call }
+                    q << true
+                end
+            end
+            q.pop
+        end
+    end
+
+    extend self
+end
+end
+end

data/lib/arachni/rpc/em/protocol.rb

@@ -0,0 +1,160 @@
+=begin
+                  Arachni-RPC
+  Copyright (c) 2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+
+  This is free software; you can copy and distribute and modify
+  this program under the term of the GPL v2.0 License
+  (See LICENSE file for details)
+
+=end
+
+module Arachni
+module RPC
+module EM
+
+#
+# Provides helper transport methods for {Message} transmission.
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1
+#
+module Protocol
+    include ::Arachni::RPC::EM::SSL
+
+    # send a maximum of 16kb of data per tick
+    MAX_CHUNK_SIZE = 1024 * 16
+
+    # become a server
+    def assume_server_role!
+        @role = :server
+    end
+
+    # become a client
+    def assume_client_role!
+        @role = :client
+    end
+
+    #
+    # Stub method, should be implemented by servers.
+    #
+    # @param    [Arachni::RPC::EM::Request]     request
+    #
+    def receive_request( request )
+        p request
+    end
+
+    #
+    # Stub method, should be implemented by clients.
+    #
+    # @param    [Arachni::RPC::EM::Response]    response
+    #
+    def receive_response( response )
+        p response
+    end
+
+    #
+    # Converts incoming hash objects to Requests or Response
+    # (depending on the assumed role) and calls receive_request() or receive_response()
+    # accordingly.
+    #
+    # @param    [Hash]      obj
+    #
+    def receive_object( obj )
+        if @role == :server
+            receive_request( Request.new( obj ) )
+        else
+            receive_response( Response.new( obj ) )
+        end
+    end
+
+    #
+    # Sends a message to the peer.
+    #
+    # @param    [Arachni::RPC::EM::Message]    msg
+    #
+    def send_message( msg )
+        ::EM.schedule {
+            send_object( msg.prepare_for_tx )
+        }
+    end
+    alias :send_request  :send_message
+    alias :send_response :send_message
+
+    #
+    # Receives data from the network.
+    #
+    # In this case the data will be chunks of a serialized object which
+    # will be buffered until the whole transmission has finished.
+    #
+    # It will then unresialize it and pass it to receive_object().
+    #
+    def receive_data( data )
+        #
+        # cut them out as soon as possible
+        #
+        # don't buffer any data from unverified peers if SSL peer
+        # veification has been enabled
+        #
+        if ssl_opts? && !verified_peer? && @role == :server
+            e = Arachni::RPC::Exceptions::SSLPeerVerificationFailed.new( 'Could not verify peer.' )
+            send_response Response.new( :obj => {
+                'exception' => e.to_s,
+                'backtrace' => e.backtrace,
+                'type'      => 'SSLPeerVerificationFailed'
+            })
+
+            log( :error, 'SSL', " Could not verify peer. ['#{peer_ip_addr}']." )
+            return
+        end
+
+        (@buf ||= '') << data
+
+        while @buf.size >= 4
+            if @buf.size >= 4 + ( size = @buf.unpack( 'N' ).first )
+                @buf.slice!( 0, 4 )
+                receive_object( serializer.load( @buf.slice!( 0, size ) ) )
+            else
+                break
+            end
+        end
+    end
+
+    #
+    # Sends a ruby object over the network
+    #
+    # Will split the object in chunks of MAX_CHUNK_SIZE and transmit one at a time.
+    #
+    def send_object( obj )
+        data = serializer.dump( obj )
+        packed = [data.bytesize, data].pack( 'Na*' )
+
+        while( packed )
+            if packed.bytesize > MAX_CHUNK_SIZE
+                send_data( packed.slice!( 0, MAX_CHUNK_SIZE ) )
+            else
+                send_data( packed )
+                break
+            end
+        end
+    end
+
+    #
+    # Returns the preferred serializer based on the 'serializer' option of the server.
+    #
+    # Defaults to <i>YAML</i>.
+    #
+    # @return   [Class]     serializer to be used
+    #
+    # @see http://eventmachine.rubyforge.org/EventMachine/Protocols/ObjectProtocol.html#M000369
+    #
+    def serializer
+        @opts[:serializer] ? @opts[:serializer] : YAML
+    end
+
+end
+
+end
+end
+end

data/lib/arachni/rpc/em/server.rb

@@ -0,0 +1,421 @@
+=begin
+                  Arachni-RPC
+  Copyright (c) 2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+
+  This is free software; you can copy and distribute and modify
+  this program under the term of the GPL v2.0 License
+  (See LICENSE file for details)
+
+=end
+
+module Arachni
+module RPC
+module EM
+
+#
+# EventMachine-based RPC server class.
+#
+# It's capable of:
+# - performing and handling a few thousands requests per second (depending on call size, network conditions and the like)
+# - TLS encryption
+# - asynchronous and synchronous requests
+# - handling asynchronous methods that require a block
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1
+#
+class Server
+
+    include ::Arachni::RPC::Exceptions
+
+    #
+    # Handles EventMachine's connection stuff.
+    #
+    # It's responsible for TLS, serializing, transmitting and receiving objects,
+    # as well as authenticating the client using the token.
+    #
+    # It also handles and forwards exceptions.
+    #
+    # @author: Tasos "Zapotek" Laskos
+    #                                      <tasos.laskos@gmail.com>
+    #                                      <zapotek@segfault.gr>
+    # @version: 0.1
+    #
+    class Proxy < EventMachine::Connection
+
+        include ::Arachni::RPC::EM::Protocol
+        include ::Arachni::RPC::Exceptions
+        include ::Arachni::RPC::EM::ConnectionUtilities
+
+        INACTIVITY_TIMEOUT = 10
+
+        attr_reader :request
+
+        def initialize( server )
+            super
+            @server = server
+            @opts   = server.opts
+
+            assume_server_role!
+
+            @id = nil
+            @request = nil
+
+            # do not tolerate long periods of
+            # inactivity in order to avoid zombie connections
+            set_comm_inactivity_timeout( INACTIVITY_TIMEOUT )
+        end
+
+        # starts TLS
+        def post_init
+            start_ssl
+        end
+
+        def unbind
+            end_ssl
+        end
+
+        def log( severity, progname, msg )
+            sev_sym = Logger.const_get( severity.to_s.upcase.to_sym )
+            @server.logger.add( sev_sym, msg, progname )
+        end
+
+        #
+        # Handles requests and sends back the responses.
+        #
+        # @param    [Arachni::RPC::EM::Request]     req
+        #
+        def receive_request( req )
+            @request = req
+
+            # the method call may block a little so tell EventMachine to
+            # stick it in its own thread.
+            # ::EM.defer( proc {
+                res  = Response.new
+                peer = peer_ip_addr
+
+                begin
+                    # token-based authentication
+                    authenticate!
+
+                    # grab the result of the method call
+                    res.merge!( @server.call( self ) )
+
+                # handle exceptions and convert them to a simple hash,
+                # ready to be passed to the client.
+                rescue Exception => e
+
+                    type = ''
+
+                    # if it's an RPC exception pass the type along as is
+                    if e.rpc_exception?
+                        type = e.class.name.split( ':' )[-1]
+
+                    # otherwise set it to a RemoteExeption
+                    else
+                        type = 'RemoteException'
+                    end
+
+                    res.obj = {
+                        'exception' => e.to_s,
+                        'backtrace' => e.backtrace,
+                        'type'      => type
+                    }
+
+                    msg = "#{e.to_s}\n#{e.backtrace.join( "\n" )}"
+                    @server.logger.error( 'Exception' ){ msg + " [on behalf of #{peer}]" }
+                end
+
+                # res
+            # }, proc {
+                # |res|
+
+                #
+                # pass the result of the RPC call back to the client
+                # along with the callback ID but *only* if it wan't async
+                # because server.call() will have already taken care of it
+                #
+                send_response( res ) if !res.async?
+            # })
+        end
+
+        #
+        # Authenticates the client based on the token in the request.
+        #
+        # It will raise an exception if the token doesn't check-out.
+        #
+        # @param    [String]    peer    IP address of the client
+        # @param    [Hash]      req     request
+        #
+        def authenticate!
+            if !valid_token?( @request.token )
+
+                msg = 'Token missing or invalid while calling: ' + @request.message
+
+                @server.logger.error( 'Authenticator' ){
+                    msg + " [on behalf of #{peer_ip_addr}]"
+                }
+
+                raise( InvalidToken.new( msg ) )
+            end
+        end
+
+        #
+        # Compares the authentication token in the param with the one of the server.
+        #
+        # @param    [String]    token
+        #
+        # @return   [Bool]
+        #
+        def valid_token?( token )
+            token == @server.token
+        end
+
+    end
+
+    attr_reader :token
+    attr_reader :opts
+    attr_reader :logger
+
+    #
+    # Starts EventMachine and the RPC server.
+    #
+    # opts example:
+    #
+    #    {
+    #        :host  => 'localhost',
+    #        :port  => 7331,
+    #
+    #        # optional authentication token, if it doesn't match the one
+    #        # set on the server-side you'll be getting exceptions.
+    #        :token => 'superdupersecret',
+    #
+    #        # optional serializer (defaults to YAML)
+    #        # see the 'serializer' method at:
+    #        # http://eventmachine.rubyforge.org/EventMachine/Protocols/ObjectProtocol.html#M000369
+    #        :serializer => Marshal,
+    #
+    #        #
+    #        # In order to enable peer verification one must first provide
+    #        # the following:
+    #        #
+    #        # SSL CA certificate
+    #        :ssl_ca     => cwd + '/../spec/pems/cacert.pem',
+    #        # SSL private key
+    #        :ssl_pkey   => cwd + '/../spec/pems/client/key.pem',
+    #        # SSL certificate
+    #        :ssl_cert   => cwd + '/../spec/pems/client/cert.pem'
+    #    }
+    #
+    # @param    [Hash]  opts
+    #
+    def initialize( opts )
+        @opts  = opts
+
+        if @opts[:ssl_pkey] && @opts[:ssl_cert]
+            if !File.exist?( @opts[:ssl_pkey] )
+                raise 'Could not find private key at: ' + @opts[:ssl_pkey]
+            end
+
+            if !File.exist?( @opts[:ssl_cert] )
+                raise 'Could not find certificate at: ' + @opts[:ssl_cert]
+            end
+        end
+
+        @token = @opts[:token]
+
+        @logger = ::Logger.new( STDOUT )
+        @logger.level = Logger::INFO
+
+        @host, @port = @opts[:host], @opts[:port]
+
+        clear_handlers
+    end
+
+    #
+    # This is a way to identify methods that pass their result to a block
+    # instead of simply returning them (which is the most usual operation of async methods.
+    #
+    # So no need to change your coding conventions to fit the RPC stuff,
+    # you can just decide dynamically based on the plethora of data which Ruby provides
+    # by its 'Method' class.
+    #
+    #    server.add_async_check {
+    #        |method|
+    #        #
+    #        # Must return 'true' for async and 'false' for sync.
+    #        #
+    #        # Very simple check here...
+    #        #
+    #        'async' ==  method.name.to_s.split( '_' )[0]
+    #    }
+    #
+    # @param    [Proc]  &block
+    #
+    def add_async_check( &block )
+        @async_checks << block
+    end
+
+    #
+    # Adds a handler by name:
+    #
+    #    server.add_handler( 'myclass', MyClass.new )
+    #
+    # @param    [String]    name    name via which to make the object available over RPC
+    # @param    [Object]    obj     object instance
+    #
+    def add_handler( name, obj )
+        @objects[name] = obj
+        @methods[name] = Set.new # no lookup overhead please :)
+        @async_methods[name] = Set.new
+
+        obj.class.public_instance_methods( false ).each {
+            |method|
+            @methods[name] << method.to_s
+            @async_methods[name] << method.to_s if async_check( obj.method( method ) )
+        }
+    end
+
+    #
+    # Clears all handlers and their associated information like methods
+    # and async check blocks.
+    #
+    def clear_handlers
+        @objects = {}
+        @methods = {}
+
+        @async_checks  = []
+        @async_methods = {}
+    end
+
+    #
+    # Runs the server and blocks.
+    #
+    def run
+        Arachni::RPC::EM.add_to_reactor {
+            start
+        }
+        Arachni::RPC::EM.block!
+    end
+
+    #
+    # Starts the server but does not block.
+    #
+    def start
+        @logger.info( 'System' ){ "RPC Server started." }
+        @logger.info( 'System' ){ "Listening on #{@host}:#{@port}" }
+
+        ::EM.start_server( @host, @port, Proxy, self )
+    end
+
+    def call( connection )
+
+        req = connection.request
+        peer_ip_addr = connection.peer_ip_addr
+
+        expr, args = req.message, req.args
+        meth_name, obj_name = parse_expr( expr )
+
+        log_call( peer_ip_addr, expr, *args )
+
+        if !object_exist?( obj_name )
+            msg = "Trying to access non-existent object '#{obj_name}'."
+            @logger.error( 'Call' ){ msg + " [on behalf of #{peer_ip_addr}]" }
+            raise( InvalidObject.new( msg ) )
+        end
+
+        if !public_method?( obj_name, meth_name )
+            msg = "Trying to access non-public method '#{meth_name}'."
+            @logger.error( 'Call' ){ msg + " [on behalf of #{peer_ip_addr}]" }
+            raise( InvalidMethod.new( msg ) )
+        end
+
+        # the proxy needs to know whether this is an async call because if it
+        # is we'll have already send the response.
+        res = Response.new
+        res.async! if async?( obj_name, meth_name )
+
+        if !res.async?
+            res.obj = @objects[obj_name].send( meth_name.to_sym, *args )
+        else
+            @objects[obj_name].send( meth_name.to_sym, *args ){
+                |obj|
+                res.obj = obj
+                connection.send_response( res )
+            }
+        end
+
+        return res
+    end
+
+    #
+    # @return   [TrueClass]
+    #
+    def alive?
+        return true
+    end
+
+    #
+    # Shuts down the server after 2 seconds
+    #
+    def shutdown
+        wait_for = 2
+
+        @logger.info( 'System' ){ "Shutting down in #{wait_for} seconds..." }
+
+        # don't die before returning
+        EventMachine::add_timer( wait_for ) { ::EM.stop }
+        return true
+    end
+
+    private
+
+    def async?( objname, method )
+        @async_methods[objname].include?( method )
+    end
+
+    def async_check( method )
+        @async_checks.each {
+            |check|
+            return true if check.call( method )
+        }
+        return false
+    end
+
+
+    def log_call( peer_ip_addr, expr, *args )
+        msg = "#{expr}"
+
+        # this should be in a @logger.debug call but it'll get out of sync
+        if @logger.level == Logger::DEBUG
+            cargs = args.map { |arg| arg.inspect }
+            msg += "( #{cargs.join( ', ' )} )"
+        end
+
+        msg += " [#{peer_ip_addr}]"
+
+        @logger.info( 'Call' ){ msg }
+    end
+
+    def parse_expr( expr )
+        parts = expr.to_s.split( '.' )
+
+        # method name, object name
+        [ parts.pop, parts.join( '.' ) ]
+    end
+
+    def object_exist?( obj_name )
+        @objects[obj_name] ? true : false
+    end
+
+    def public_method?( obj_name, method )
+        @methods[obj_name].include?( method )
+    end
+
+end
+
+end
+end
+end