aikido-zen 0.1.0.alpha4-x86_64-linux → 0.1.1-x86_64-linux

data/lib/aikido/zen/agent/heartbeats_manager.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  # Handles scheduling the heartbeats we send to the Aikido servers, managing
+  # runtime changes to the heartbeat interval.
+  class Agent::HeartbeatsManager
+    def initialize(worker:, settings: Aikido::Zen.runtime_settings, config: Aikido::Zen.config)
+      @settings = settings
+      @config = config
+      @worker = worker
+
+      @timer = nil
+    end
+
+    # @return [Boolean]
+    def running?
+      !!@timer&.running?
+    end
+
+    # @return [Boolean] whether the currently running heartbeat matches the
+    #   expected interval in the runtime settings.
+    def stale_settings?
+      running? && @timer.execution_interval != @settings.heartbeat_interval
+    end
+
+    # Sets up the the timer to run the given block at the appropriate interval.
+    # Re-entrant, and does nothing if already running.
+    #
+    # @return [void]
+    def start(&task)
+      return if running?
+
+      if @settings.heartbeat_interval&.nonzero?
+        @config.logger.debug "Scheduling heartbeats every #{@settings.heartbeat_interval} seconds"
+        @timer = @worker.every(@settings.heartbeat_interval, run_now: false, &task)
+      else
+        @config.logger.warn(format("Heartbeat could not be set up (interval: %p)", @settings.heartbeat_interval))
+      end
+    end
+
+    # Cleans up the timer.
+    #
+    # @return [void]
+    def stop
+      return unless running?
+
+      @timer.shutdown
+      @timer = nil
+    end
+
+    # Resets the timer to start with any new settings, if needed.
+    #
+    # @return [void]
+    def restart(&task)
+      stop
+      start(&task)
+    end
+
+    # @api private
+    #
+    # @return [Integer] the current delay between events.
+    def interval
+      @settings.heartbeat_interval
+    end
+  end
+end

data/lib/aikido/zen/agent.rb

@@ -2,7 +2,6 @@
 
 require "concurrent"
 require_relative "event"
-require_relative "stats"
 require_relative "config"
 require_relative "system_info"
 
@@ -10,31 +9,89 @@ module Aikido::Zen
   # Handles the background processes that communicate with the Aikido servers,
   # including managing the runtime settings that keep the app protected.
   class Agent
-    # @return [Aikido::Zen::Stats] the statistics collected by the agent.
-    attr_reader :stats
+    # Initialize and start an agent instance.
+    #
+    # @return [Aikido::Zen::Agent]
+    def self.start(**opts)
+      new(**opts).tap(&:start!)
+    end
 
     def initialize(
-      stats: Aikido::Zen::Stats.new,
       config: Aikido::Zen.config,
-      info: Aikido::Zen.system_info,
-      api_client: Aikido::Zen::APIClient.new
+      collector: Aikido::Zen.collector,
+      worker: Aikido::Zen::Worker.new(config: config),
+      api_client: Aikido::Zen::APIClient.new(config: config)
     )
       @started_at = nil
 
-      @stats = stats
-      @info = info
       @config = config
+      @worker = worker
       @api_client = api_client
-      @timer_tasks = []
-      @delayed_tasks = []
-      @reporting_pool = nil
-      @heartbeats = nil
+      @collector = collector
     end
 
     def started?
       !!@started_at
     end
 
+    def start!
+      @config.logger.info "Starting Aikido agent"
+
+      raise Aikido::ZenError, "Aikido Agent already started!" if started?
+      @started_at = Time.now.utc
+      @collector.start(at: @started_at)
+
+      if @config.blocking_mode?
+        @config.logger.info "Requests identified as attacks will be blocked"
+      else
+        @config.logger.warn "Non-blocking mode enabled! No requests will be blocked."
+      end
+
+      if @api_client.can_make_requests?
+        @config.logger.info "API Token set! Reporting has been enabled."
+      else
+        @config.logger.warn "No API Token set! Reporting has been disabled."
+        return
+      end
+
+      at_exit { stop! if started? }
+
+      report(Events::Started.new(time: @started_at)) do |response|
+        Aikido::Zen.runtime_settings.update_from_json(response)
+        @config.logger.info "Updated runtime settings."
+      rescue => err
+        @config.logger.error(err.message)
+      end
+
+      poll_for_setting_updates
+
+      @worker.delay(@config.initial_heartbeat_delay) do
+        send_heartbeat if @collector.stats.any?
+      end
+    end
+
+    # Clean up any ongoing threads, and reset the state. Called automatically
+    # when the process exits.
+    #
+    # @return [void]
+    def stop!
+      @config.logger.info "Stopping Aikido agent"
+      @started_at = nil
+      @worker.shutdown
+    end
+
+    # Respond to the runtime settings changing after being fetched from the
+    # Aikido servers.
+    #
+    # @return [void]
+    def updated_settings!
+      if !heartbeats.running?
+        heartbeats.start { send_heartbeat }
+      elsif heartbeats.stale_settings?
+        heartbeats.restart { send_heartbeat }
+      end
+    end
+
     # Given an Attack, report it to the Aikido server, and/or block the request
     # depending on configuration.
     #
@@ -49,8 +106,8 @@ module Aikido::Zen
       @config.logger.error("[ATTACK DETECTED] #{attack.log_message}")
       report(Events::Attack.new(attack: attack)) if @api_client.can_make_requests?
 
-      stats.add_attack(attack, being_blocked: @config.blocking_mode?)
-      raise attack if @config.blocking_mode?
+      @collector.track_attack(attack)
+      raise attack if attack.blocked?
     end
 
     # Asynchronously reports an Event of any kind to the Aikido dashboard. If
@@ -62,7 +119,7 @@ module Aikido::Zen
     #
     # @return [void]
     def report(event)
-      reporting_pool.post do
+      @worker.perform do
         response = @api_client.report(event)
         yield response if response && block_given?
       rescue Aikido::Zen::APIError, Aikido::Zen::NetworkError => err
@@ -70,53 +127,35 @@ module Aikido::Zen
       end
     end
 
-    def start!
-      @config.logger.info "Starting Aikido agent"
-
-      raise Aikido::ZenError, "Aikido Agent already started!" if started?
-      @started_at = Time.now.utc
-
-      stats.start(@started_at)
-
-      if @config.blocking_mode?
-        @config.logger.info "Requests identified as attacks will be blocked"
-      else
-        @config.logger.warn "Non-blocking mode enabled! No requests will be blocked."
-      end
-
-      if @api_client.can_make_requests?
-        @config.logger.info "API Token set! Reporting has been enabled."
-      else
-        @config.logger.warn "No API Token set! Reporting has been disabled."
-        return
-      end
-
-      # Subscribe to firewall setting changes so we can correctly re-configure
-      # the heartbeat process.
-      Aikido::Zen.runtime_settings.add_observer(self, :setup_heartbeat)
-
-      at_exit { stop! if started? }
-
-      report(Events::Started.new(time: @started_at)) do |response|
-        Aikido::Zen.runtime_settings.update_from_json(response)
-        @config.logger.info "Updated runtime settings."
-      end
-
-      poll_for_setting_updates
-    end
-
-    def send_heartbeat
+    # @api private
+    #
+    # Atomically flushes all the stats stored by the agent, and sends a
+    # heartbeat event. Scheduled to run automatically on a recurring schedule
+    # when reporting is enabled.
+    #
+    # @param at [Time] the event time. Defaults to now.
+    # @return [void]
+    # @see Aikido::Zen::RuntimeSettings#heartbeat_interval
+    def send_heartbeat(at: Time.now.utc)
       return unless @api_client.can_make_requests?
 
-      flushed_stats = stats.reset
-      report(Events::Heartbeat.new(stats: flushed_stats)) do |response|
+      event = @collector.flush(at: at)
+
+      report(event) do |response|
         Aikido::Zen.runtime_settings.update_from_json(response)
         @config.logger.info "Updated runtime settings after heartbeat"
       end
     end
 
+    # @api private
+    #
+    # Sets up the timer task that polls the Aikido Runtime API for updates to
+    # the runtime settings every minute.
+    #
+    # @return [void]
+    # @see Aikido::Zen::RuntimeSettings
     def poll_for_setting_updates
-      timer_task(every: @config.polling_interval) do
+      @worker.every(@config.polling_interval) do
         if @api_client.should_fetch_settings?
           Aikido::Zen.runtime_settings.update_from_json(@api_client.fetch_settings)
           @config.logger.info "Updated runtime settings after polling"
@@ -124,64 +163,13 @@ module Aikido::Zen
       end
     end
 
-    def setup_heartbeat(settings)
-      return unless @api_client.can_make_requests?
-
-      # If the desired interval changed, then clear the current heartbeat timer
-      # and set up a new one.
-      if @heartbeats&.running? && @heartbeats.execution_interval != settings.heartbeat_interval
-        @heartbeats.shutdown
-        @heartbeats = nil
-        setup_heartbeat(settings)
-
-      # If the heartbeat timer isn't running but we know how often it should run, schedule it.
-      elsif !@heartbeats&.running? && settings.heartbeat_interval&.nonzero?
-        @config.logger.debug "Scheduling heartbeats every #{settings.heartbeat_interval} seconds"
-        @heartbeats = timer_task(every: settings.heartbeat_interval, run_now: false) do
-          send_heartbeat
-        end
-
-      elsif !@heartbeats&.running?
-        @config.logger.debug(format("Heartbeat could not be setup (interval: %p)", settings.heartbeat_interval))
-      end
-
-      # If the server hasn't received any stats, we want to also run a one-off
-      # heartbeat request in a minute.
-      if !settings.received_any_stats
-        delay(@config.initial_heartbeat_delay) { send_heartbeat if stats.any? }
-      end
-    end
-
-    def stop!
-      @started_at = nil
-
-      @config.logger.info "Stopping Aikido agent"
-
-      @timer_tasks.each { |task| task.shutdown }
-      @delayed_tasks.each { |task| task.cancel if task.pending? }
-
-      @reporting_pool&.shutdown
-      @reporting_pool&.wait_for_termination(30)
-    end
-
-    private def reporting_pool
-      @reporting_pool ||= Concurrent::SingleThreadExecutor.new
-    end
-
-    private def delay(delay, &task)
-      Concurrent::ScheduledTask.execute(delay, executor: reporting_pool, &task)
-        .tap { |task| @delayed_tasks << task }
-    end
-
-    private def timer_task(every:, **opts, &block)
-      Concurrent::TimerTask.execute(
-        run_now: true,
-        interval_type: :fixed_rate,
-        execution_interval: every,
-        executor: reporting_pool,
-        **opts,
-        &block
-      ).tap { |task| @timer_tasks << task }
+    private def heartbeats
+      @heartbeats ||= Aikido::Zen::Agent::HeartbeatsManager.new(
+        config: @config,
+        worker: @worker
+      )
     end
   end
 end
+
+require_relative "agent/heartbeats_manager"

data/lib/aikido/zen/collector/hosts.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative "../capped_collections"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Keeps track of the hostnames to which the app has made outbound HTTP
+  # requests.
+  class Collector::Hosts < Aikido::Zen::CappedSet
+    def initialize(config = Aikido::Zen.config)
+      super(config.max_outbound_connections)
+    end
+  end
+end

data/lib/aikido/zen/collector/routes.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require_relative "../request/schema/empty_schema"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Keeps track of the visited routes.
+  class Collector::Routes
+    def initialize(config = Aikido::Zen.config)
+      @config = config
+      @visits = Hash.new { |h, k| h[k] = Record.new }
+    end
+
+    # @param request [Aikido::Zen::Request].
+    # @return [self]
+    def add(request)
+      @visits[request.route].increment(request) unless request.route.nil?
+      self
+    end
+
+    def as_json
+      @visits.map do |route, record|
+        {
+          method: route.verb,
+          path: route.path,
+          hits: record.hits,
+          apispec: record.schema.as_json
+        }.compact
+      end
+    end
+
+    # @api private
+    def [](route)
+      @visits[route]
+    end
+
+    # @api private
+    def empty?
+      @visits.empty?
+    end
+
+    # @api private
+    Record = Struct.new(:hits, :schema, :samples) do
+      def initialize(config = Aikido::Zen.config)
+        super(0, Aikido::Zen::Request::Schema::EMPTY_SCHEMA, 0)
+        @config = config
+      end
+
+      def increment(request)
+        self.hits += 1
+
+        if sample_schema?
+          self.samples += 1
+          self.schema |= request.schema
+        end
+      end
+
+      private def sample_schema?
+        samples < @config.api_schema_max_samples
+      end
+    end
+  end
+end

data/lib/aikido/zen/{stats → collector}/sink_stats.rb

@@ -6,7 +6,7 @@ module Aikido::Zen
   # @api private
   #
   # Tracks data specific to a single Sink.
-  class Stats::SinkStats
+  class Collector::SinkStats
     # @return [Integer] number of total calls to Sink#scan.
     attr_accessor :scans
 

data/lib/aikido/zen/collector/stats.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+require_relative "../capped_collections"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Tracks information about how the Aikido Agent is used in the app.
+  class Collector::Stats
+    # @!visibility private
+    attr_reader :started_at, :ended_at, :requests, :aborted_requests, :sinks
+
+    # @!visibility private
+    attr_writer :ended_at
+
+    def initialize(config = Aikido::Zen.config)
+      super()
+      @config = config
+      @sinks = Hash.new { |h, k| h[k] = Collector::SinkStats.new(k, @config) }
+      @started_at = @ended_at = nil
+      @requests = 0
+      @aborted_requests = 0
+    end
+
+    # @return [Boolean]
+    def empty?
+      @requests.zero? && @sinks.empty?
+    end
+
+    # @return [Boolean]
+    def any?
+      !empty?
+    end
+
+    # Track the timestamp we start tracking this series of stats.
+    #
+    # @param at [Time]
+    # @return [self]
+    def start(at = Time.now.utc)
+      @started_at = at
+      self
+    end
+
+    # Sets the end time for these stats block, freezes it to avoid any more
+    # writing to them, and compresses the timing stats in anticipation of
+    # sending these to the Aikido servers.
+    #
+    # @param at [Time] the time at which we're resetting, which is set as the
+    #   ending time for the returned copy.
+    # @return [self]
+    def flush(at: Time.now.utc)
+      # Make sure the timing stats are compressed before copying, since we
+      # need these compressed when we serialize this for the API.
+      @sinks.each_value { |sink| sink.compress_timings(at: at) }
+      @ended_at = at
+      freeze
+    end
+
+    # @return [self]
+    def add_request
+      @requests += 1
+      self
+    end
+
+    # @param scan [Aikido::Zen::Scan]
+    # @return [self]
+    def add_scan(scan)
+      stats = @sinks[scan.sink.name]
+      stats.scans += 1
+      stats.errors += 1 if scan.errors?
+      stats.add_timing(scan.duration)
+      self
+    end
+
+    # @param attack [Aikido::Zen::Attack]
+    # @param being_blocked [Boolean] whether the Agent blocked the
+    #   request where this Attack happened or not.
+    # @return [self]
+    def add_attack(attack, being_blocked:)
+      stats = @sinks[attack.sink.name]
+      stats.attacks += 1
+      stats.blocked_attacks += 1 if being_blocked
+      self
+    end
+
+    def as_json
+      total_attacks, total_blocked = aggregate_attacks_from_sinks
+      {
+        startedAt: @started_at.to_i * 1000,
+        endedAt: (@ended_at.to_i * 1000 if @ended_at),
+        sinks: @sinks.transform_values(&:as_json),
+        requests: {
+          total: @requests,
+          aborted: @aborted_requests,
+          attacksDetected: {
+            total: total_attacks,
+            blocked: total_blocked
+          }
+        }
+      }
+    end
+
+    private def aggregate_attacks_from_sinks
+      @sinks.each_value.reduce([0, 0]) { |(attacks, blocked), stats|
+        [attacks + stats.attacks, blocked + stats.blocked_attacks]
+      }
+    end
+  end
+end
+
+require_relative "sink_stats"

data/lib/aikido/zen/{stats → collector}/users.rb

@@ -2,11 +2,15 @@
 
 require_relative "../capped_collections"
 
-class Aikido::Zen::Stats
+module Aikido::Zen
   # @api private
   #
   # Keeps track of the users that were seen by the app.
-  class Users < Aikido::Zen::CappedMap
+  class Collector::Users < Aikido::Zen::CappedMap
+    def initialize(config = Aikido::Zen.config)
+      super(config.max_users_tracked)
+    end
+
     def add(actor)
       if key?(actor.id)
         self[actor.id].update

data/lib/aikido/zen/collector.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  # Handles collecting all the runtime statistics to report back to the Aikido
+  # servers.
+  class Collector
+    def initialize(config: Aikido::Zen.config)
+      @config = config
+
+      @stats = Concurrent::AtomicReference.new(Stats.new(@config))
+      @users = Concurrent::AtomicReference.new(Users.new(@config))
+      @hosts = Concurrent::AtomicReference.new(Hosts.new(@config))
+      @routes = Concurrent::AtomicReference.new(Routes.new(@config))
+    end
+
+    # Flush all the stats into a Heartbeat event that can be reported back to
+    # the Aikido servers.
+    #
+    # @param at [Time] the time at which stats collection stopped and the start
+    #   of the new stats collection period. Defaults to now.
+    # @return [Aikido::Zen::Events::Heartbeat]
+    def flush(at: Time.now.utc)
+      stats = @stats.get_and_set(Stats.new(@config))
+      users = @users.get_and_set(Users.new(@config))
+      hosts = @hosts.get_and_set(Hosts.new(@config))
+      routes = @routes.get_and_set(Routes.new(@config))
+
+      start(at: at)
+      stats = stats.flush(at: at)
+
+      Events::Heartbeat.new(stats: stats, users: users, hosts: hosts, routes: routes)
+    end
+
+    # Sets the start time for this collection period.
+    #
+    # @param at [Time] defaults to now.
+    # @return [void]
+    def start(at: Time.now.utc)
+      synchronize(@stats) { |stats| stats.start(at) }
+    end
+
+    # Track stats about the request, record the visited endpoint, and if
+    # enabled, the API schema for this endpoint.
+    #
+    # @param request [Aikido::Zen::Request]
+    # @return [void]
+    def track_request(request)
+      synchronize(@stats) { |stats| stats.add_request }
+      synchronize(@routes) { |routes| routes.add(request) if request.route }
+    end
+
+    # Track stats about a scan performed by one of our sinks.
+    #
+    # @param scan [Aikido::Zen::Scan]
+    # @return [void]
+    def track_scan(scan)
+      synchronize(@stats) { |stats| stats.add_scan(scan) }
+    end
+
+    # Track stats about an attack detected by our scanners.
+    #
+    # @param attack [Aikido::Zen::Attack]
+    # @return [void]
+    def track_attack(attack)
+      synchronize(@stats) do |stats|
+        stats.add_attack(attack, being_blocked: attack.blocked?)
+      end
+    end
+
+    # Track an HTTP connections to an external host.
+    #
+    # @param connection [Aikido::Zen::OutboundConnection]
+    # @return [void]
+    def track_outbound(connection)
+      synchronize(@hosts) { |hosts| hosts.add(connection) }
+    end
+
+    # Track the user reported by the developer to be behind this request.
+    #
+    # @param actor [Aikido::Zen::Actor]
+    # @return [void]
+    def track_user(actor)
+      synchronize(@users) { |users| users.add(actor) }
+    end
+
+    # @api private
+    def routes
+      @routes.get
+    end
+
+    # @api private
+    def users
+      @users.get
+    end
+
+    # @api private
+    def hosts
+      @hosts.get
+    end
+
+    # @api private
+    def stats
+      @stats.get
+    end
+
+    # Atomically modify an object's state within a block, ensuring it's safe
+    # from other threads.
+    private def synchronize(object)
+      object.update { |obj| obj.tap { yield obj } }
+    end
+  end
+end
+
+require_relative "collector/stats"
+require_relative "collector/users"
+require_relative "collector/hosts"
+require_relative "collector/routes"