cuboid 0.0.0 → 0.0.3

data/lib/cuboid/rpc/server/dispatcher.rb

@@ -0,0 +1,317 @@
+module Cuboid
+
+lib = Options.paths.lib
+require lib + 'processes/instances'
+require lib + 'rpc/client'
+require lib + 'rpc/server/base'
+require lib + 'rpc/server/instance'
+require lib + 'rpc/server/output'
+
+module RPC
+class Server
+
+# Dispatches RPC Instances on demand and allows for extensive process monitoring.
+#
+# The process goes something like this:
+#
+# * A client issues a {#dispatch} call.
+# * The Dispatcher spawns and returns Instance info to the client (url, auth token, etc.).
+# * The client connects to the Instance using that info.
+#
+# Once the client finishes using the RPC Instance it *must* shut it down
+# otherwise the system will be eaten away by zombie processes.
+#
+# @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+class Dispatcher
+    require Options.paths.lib + 'rpc/server/dispatcher/node'
+    require Options.paths.lib + 'rpc/server/dispatcher/service'
+
+    include Utilities
+    include UI::Output
+
+    SERVICE_NAMESPACE = Service
+
+    def initialize( options = Options.instance )
+        @options = options
+
+        @options.snapshot.path ||= @options.paths.snapshots
+
+        @server = Base.new( @options.rpc.to_server_options )
+        @server.logger.level = @options.datastore.log_level if @options.datastore.log_level
+
+        @server.add_async_check do |method|
+            # methods that expect a block are async
+            method.parameters.flatten.include? :block
+        end
+
+        Options.dispatcher.url = @url = @server.url
+
+        prep_logging
+
+        print_status 'Starting the RPC Server...'
+
+        @server.add_handler( 'dispatcher', self )
+
+        # trap interrupts and exit cleanly when required
+        trap_interrupts { shutdown }
+
+        @instances = []
+
+        Cuboid::Application.application.dispatcher_services.each do |name, service|
+            @server.add_handler( name.to_s, service.new( @options, self ) )
+        end
+
+        @node = Node.new( @options, @server, @logfile )
+        @server.add_handler( 'node', @node )
+
+        run
+    end
+
+    def services
+        Cuboid::Application.application.dispatcher_services.keys
+    end
+
+    # @return   [TrueClass]
+    #   true
+    def alive?
+        @server.alive?
+    end
+
+    # @return   [String, nil]
+    #   Depending on availability:
+    #
+    #   * URL of the least burdened Dispatcher. If not a grid member it will
+    #       return this Dispatcher's URL.
+    #   * `nil` if all nodes are at max utilization.
+    def preferred( &block )
+        if !@node.grid_member?
+            block.call( self.utilization == 1 ? nil : @url )
+            return
+        end
+
+        pick_utilization = proc do |url, utilization|
+            (utilization == 1 || utilization.rpc_exception?) ?
+                nil : [url, utilization]
+        end
+
+        each = proc do |neighbour, iter|
+            connect_to_peer( neighbour ).utilization do |utilization|
+                iter.return pick_utilization.call( neighbour, utilization )
+            end
+        end
+
+        after = proc do |nodes|
+            nodes << pick_utilization.call( @url, self.utilization )
+            nodes.compact!
+
+            # All nodes are at max utilization, pass.
+            if nodes.empty?
+                block.call
+                next
+            end
+
+            block.call nodes.sort_by { |_, score| score }[0][0]
+        end
+
+        Arachni::Reactor.global.create_iterator( @node.neighbours ).map( each, after )
+    end
+
+    # Dispatches an {Instance}.
+    #
+    # @param    [String]  owner
+    #   An owner to assign to the {Instance}.
+    # @param    [Hash]    helpers
+    #   Hash of helper data to be added to the instance info.
+    # @param    [Boolean]    load_balance
+    #   Return an {Instance} from the least burdened {Dispatcher} (when in Grid mode)
+    #   or from this one directly?
+    #
+    # @return   [Hash, nil]
+    #   Depending on availability:
+    #
+    #   * `Hash`: Connection and proc info.
+    #   * `nil`: Max utilization, wait for one of the instances to finish and retry.
+    def dispatch( options = {}, &block )
+        options      = options.my_symbolize_keys
+        owner        = options[:owner]
+        helpers      = options[:helpers] || {}
+        load_balance = options[:load_balance].nil? ? true : options[:load_balance]
+
+        if load_balance && @node.grid_member?
+            preferred do |url|
+                if !url
+                    block.call
+                    next
+                end
+
+                connect_to_peer( url ).dispatch( options.merge(
+                      helpers: helpers.merge( via: @url ),
+                      load_balance: false
+                    ),
+                    &block
+                )
+            end
+            return
+        end
+
+        if System.max_utilization?
+            block.call
+            return
+        end
+
+        spawn_instance do |info|
+            info['owner']   = owner
+            info['helpers'] = helpers
+
+            @instances << info
+
+            block.call info
+        end
+    end
+
+    # Returns proc info for a given pid
+    #
+    # @param    [Fixnum]      pid
+    #
+    # @return   [Hash]
+    def instance( pid )
+        @instances.each do |i|
+            next if i['pid'] != pid
+            i = i.dup
+
+            now = Time.now
+
+            i['now']   = now.to_s
+            i['age']   = now - Time.parse( i['birthdate'] )
+            i['alive'] = Cuboid::Processes::Manager.alive?( pid )
+
+            return i
+        end
+
+        nil
+    end
+
+    # @return   [Array<Hash>]
+    #   Returns info for all instances.
+    def instances
+        @instances.map { |i| instance( i['pid'] ) }.compact
+    end
+
+    # @return   [Array<Hash>]
+    #   Returns info for all running (alive) instances.
+    #
+    # @see #instances
+    def running_instances
+        instances.select { |i| i['alive'] }
+    end
+
+    # @return   [Array<Hash>]
+    #   Returns info for all finished (dead) instances.
+    #
+    # @see #instances
+    def finished_instances
+        instances.reject { |i| i['alive'] }
+    end
+
+    # @return   [Float]
+    #   Workload score for this Dispatcher, calculated using {System#utilization}.
+    #
+    #   * `0.0` => No utilization.
+    #   * `1.0` => Max utilization.
+    #
+    #   Lower is better.
+    def utilization
+        System.utilization
+    end
+
+    # @return   [Hash]
+    #   Returns server stats regarding the instances and pool.
+    def statistics
+        {
+            'utilization'         => utilization,
+            'running_instances'   => running_instances,
+            'finished_instances'  => finished_instances,
+            'consumed_pids'       => @instances.map { |i| i['pid'] }.compact,
+            'snapshots'           => Dir.glob( "#{@options.snapshot.path}*.#{Snapshot::EXTENSION}" ),
+            'node'                => @node.info
+        }
+    end
+
+    # @return   [String]
+    #   Contents of the log file
+    def log
+        IO.read prep_logging
+    end
+
+    # @private
+    def pid
+        Process.pid
+    end
+
+    private
+
+    def trap_interrupts( &block )
+        %w(QUIT INT).each do |signal|
+            trap( signal, &block || Proc.new{ } ) if Signal.list.has_key?( signal )
+        end
+    end
+
+    # Starts the dispatcher's server
+    def run
+        Arachni::Reactor.global.on_error do |_, e|
+            print_error "Reactor: #{e}"
+
+            e.backtrace.each do |l|
+                print_error "Reactor: #{l}"
+            end
+        end
+
+        print_status 'Ready'
+        @server.start
+    rescue => e
+        print_exception e
+
+        $stderr.puts "Could not start server, for details see: #{@logfile}"
+        exit 1
+    end
+
+    def shutdown
+        Thread.new do
+            print_status 'Shutting down...'
+            Arachni::Reactor.global.stop
+        end
+    end
+
+    def spawn_instance( options = {}, &block )
+        Processes::Instances.spawn( options.merge(
+            address:     @server.address,
+            port_range:  Options.dispatcher.instance_port_range,
+            token:       Utilities.generate_token,
+            application: Options.paths.application
+        )) do |client|
+            block.call(
+                'token'       => client.token,
+                'url'         => client.url,
+                'pid'         => client.pid,
+                'birthdate'   => Time.now.to_s,
+                'application' => Options.paths.application
+            )
+        end
+    end
+
+    def prep_logging
+        # reroute all output to a logfile
+        @logfile ||= reroute_to_file( @options.paths.logs +
+            "/Dispatcher - #{Process.pid}-#{@options.rpc.server_port}.log" )
+    end
+
+    def connect_to_peer( url )
+        @rpc_clients ||= {}
+        @rpc_clients[url] ||= Client::Dispatcher.new( url )
+    end
+
+end
+
+end
+end
+end

data/lib/cuboid/rpc/server/instance.rb

@@ -0,0 +1,338 @@
+require 'ostruct'
+
+module Cuboid
+lib = Options.paths.lib
+
+require lib + 'processes/manager'
+
+require lib + 'rpc/client/instance'
+
+require lib + 'rpc/server/base'
+require lib + 'rpc/server/active_options'
+require lib + 'rpc/server/output'
+require lib + 'rpc/server/application_wrapper'
+
+module RPC
+class Server
+
+# @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+class Instance
+    include UI::Output
+    include Utilities
+
+    [
+        :suspend!, :suspended?, :snapshot_path,
+        :pause!,   :paused?,
+        :abort!,   :resume!,
+        :running?, :status
+    ].each do |m|
+        define_method m do
+            @application.send m
+        end
+    end
+
+    private :error_logfile
+    public  :error_logfile
+
+    # Initializes the RPC interface and the framework.
+    #
+    # @param    [Options]    options
+    # @param    [String]    token
+    #   Authentication token.
+    def initialize( options, token )
+        @options = options
+        @token   = token
+
+        @application    = Server::ApplicationWrapper.new(
+          Cuboid::Application.application
+        )
+        @active_options = Server::ActiveOptions.new
+
+        @server = Base.new( @options.rpc.to_server_options, token )
+
+        if @options.datastore.log_level
+            @server.logger.level = @options.datastore.log_level
+        end
+
+        @options.datastore.token = token
+
+        if @options.output.reroute_to_logfile
+            reroute_to_file "#{@options.paths.logs}/Instance - #{Process.pid}" <<
+                                "-#{@options.rpc.server_port}.log"
+        else
+            reroute_to_file false
+        end
+
+        set_error_logfile "#{@options.paths.logs}/Instance - #{Process.pid}" <<
+                              "-#{@options.rpc.server_port}.error.log"
+
+        set_handlers( @server )
+
+        # trap interrupts and exit cleanly when required
+        %w(QUIT INT).each do |signal|
+            next if !Signal.list.has_key?( signal )
+            trap( signal ){ shutdown if !@options.datastore.do_not_trap }
+        end
+
+        Arachni::Reactor.global.run do
+            _run
+        end
+    end
+
+    def application
+        Application.application.to_s
+    end
+
+    # @return   [String, nil]
+    #   Scheduler URL to which this Instance is attached, `nil` if not attached.
+    def scheduler_url
+        @options.scheduler.url
+    end
+
+    # @return   [String, nil]
+    #   Dispatcher URL that provided this Instance, `nil` if not provided by a
+    #   Dispatcher.
+    def dispatcher_url
+        @options.dispatcher.url
+    end
+
+    # @param (see Cuboid::Application#restore)
+    # @return (see Cuboid::Application#restore)
+    #
+    # @see #suspend
+    # @see #snapshot_path
+    def restore!( snapshot )
+        @application.restore!( snapshot ).run
+        true
+    end
+
+    # @return   [true]
+    def alive?
+        @server.alive?
+    end
+
+    # @return   [Bool]
+    #   `true` if the scan is initializing or running, `false` otherwise.
+    def busy?
+        @run_initializing || @application.busy?
+    end
+
+    # Cleans up and returns the report.
+    #
+    # @return  [Hash]
+    #
+    # @see #report
+    def abort_and_generate_report
+        @application.clean_up
+        generate_report
+    end
+
+    # @return [Hash]
+    #   {Report#to_rpc_data}
+    def generate_report
+        @application.generate_report.to_rpc_data
+    end
+
+    # # Recommended usage
+    #
+    #   Please request from the method only the things you are going to actually
+    #   use, otherwise you'll just be wasting bandwidth.
+    #   In addition, ask to **not** be served data you already have, like
+    #   error messages.
+    #
+    #   To be kept completely up to date on the progress of a scan (i.e. receive
+    #   new issues and error messages asap) in an efficient manner, you will need
+    #   to keep track of the error messages you already have and explicitly tell
+    #   the method to not send the same data back to you on subsequent calls.
+    #
+    # ## Retrieving errors (`:errors` option) without duplicate data
+    #
+    #   This is done by telling the method how many error messages you already
+    #   have and you will be served the errors from the error-log that are past
+    #   that line.
+    #   So, if you were to use a loop to get fresh progress data it would look
+    #   like so:
+    #
+    #     error_cnt = 0
+    #     i = 0
+    #     while sleep 1
+    #         # Test method, triggers an error log...
+    #         instance.error_test "BOOM! #{i+=1}"
+    #
+    #         # Only request errors we don't already have
+    #         errors = instance.progress( with: { errors: error_cnt } )[:errors]
+    #         error_cnt += errors.size
+    #
+    #         # You will only see new errors
+    #         puts errors.join("\n")
+    #     end
+    #
+    # @param  [Hash]  options
+    #   Options about what progress data to retrieve and return.
+    # @option options [Array<Symbol, Hash>]  :with
+    #   Specify data to include:
+    #
+    #   * :errors -- Errors and the line offset to use for {#errors}.
+    #     Pass as a hash, like: `{ errors: 10 }`
+    # @option options [Array<Symbol, Hash>]  :without
+    #   Specify data to exclude:
+    #
+    #   * :statistics -- Don't include runtime statistics.
+    #
+    # @return [Hash]
+    #   * `statistics` -- General runtime statistics (merged when part of Grid)
+    #       (enabled by default)
+    #   * `status` -- {#status}
+    #   * `busy` -- {#busy?}
+    #   * `errors` -- {#errors} (disabled by default)
+    def progress( options = {} )
+        progress_handler( options.merge( as_hash: true ) )
+    end
+
+    # Configures and runs a job.
+    def run( options = nil )
+        # If the instance isn't clean bail out now.
+        return false if busy? || @called
+
+        if !@application.valid_options?( options )
+            fail ArgumentError, 'Invalid options!'
+        end
+
+        # There may be follow-up/retry calls by the client in cases of network
+        # errors (after the request has reached us) so we need to keep minimal
+        # track of state in order to bail out on subsequent calls.
+        @called = @run_initializing = true
+
+        @active_options.set( application: options )
+
+        @application.run
+        @run_initializing = false
+
+        true
+    end
+
+    # Makes the server go bye-bye...Lights out!
+    def shutdown( &block )
+        if @shutdown
+            block.call if block_given?
+            return
+        end
+        @shutdown = true
+
+        print_status 'Shutting down...'
+
+        # We're shutting down services so we need to use a concurrent way but
+        # without going through the Reactor.
+        Thread.new do
+            @server.shutdown
+            block.call true if block_given?
+        end
+
+        true
+    end
+
+    def errors( starting_line = 0 )
+        @application.errors( starting_line )
+    end
+
+    # @private
+    def error_test( str )
+        @application.error_test( str )
+    end
+
+    # @private
+    def consumed_pids
+        [Process.pid]
+    end
+
+    def self.parse_progress_opts( options, key )
+        parsed = {}
+        [options.delete( key ) || options.delete( key.to_s )].compact.each do |w|
+            case w
+                when Array
+                    w.compact.flatten.each do |q|
+                        case q
+                            when String, Symbol
+                                parsed[q.to_sym] = nil
+
+                            when Hash
+                                parsed.merge!( q.my_symbolize_keys )
+                        end
+                    end
+
+                when String, Symbol
+                    parsed[w.to_sym] = nil
+
+                when Hash
+                    parsed.merge!( w.my_symbolize_keys )
+            end
+        end
+
+        parsed
+    end
+
+    private
+
+    def progress_handler( options = {}, &block )
+        with    = self.class.parse_progress_opts( options, :with )
+        without = self.class.parse_progress_opts( options, :without )
+
+        options = {
+            as_hash:    options[:as_hash],
+            statistics: !without.include?( :statistics )
+        }
+
+        if with[:errors]
+            options[:errors] = with[:errors]
+        end
+
+        @application.progress( options ) do |data|
+            data[:busy] = busy?
+            block.call( data )
+        end
+    end
+
+    # Starts  RPC service.
+    def _run
+        Arachni::Reactor.global.on_error do |_, e|
+            print_error "Reactor: #{e}"
+
+            e.backtrace.each do |l|
+                print_error "Reactor: #{l}"
+            end
+        end
+
+        print_status 'Starting the server...'
+        @server.start
+    end
+
+    # Outputs the Engine banner.
+    #
+    # Displays version number, author details etc.
+    def banner
+        puts BANNER
+        puts
+        puts
+    end
+
+    # @param    [Base]  server
+    #   Prepares all the RPC handlers for the given `server`.
+    def set_handlers( server )
+        server.add_async_check do |method|
+            # methods that expect a block are async
+            method.parameters.flatten.include? :block
+        end
+
+        server.add_handler( 'instance', self )
+        server.add_handler( 'options',  @active_options )
+
+        Cuboid::Application.application.instance_services.each do |name, service|
+            server.add_handler( name.to_s, service.new )
+        end
+    end
+
+end
+
+end
+end
+end