aikido-zen 1.0.1.beta.5-arm64-darwin → 1.0.2-arm64-darwin

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

@@ -0,0 +1,238 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  class Collector
+    class Event
+      @@registry = {}
+
+      # @api protected
+      def self.register(type)
+        const_set(:TYPE, type)
+        @@registry[type] = self
+      end
+
+      def self.from_json(data)
+        type = data[:type]
+        subclass = @@registry[type]
+        subclass.from_json(data)
+      end
+
+      attr_reader :type
+
+      def initialize
+        @type = self.class::TYPE
+      end
+
+      def as_json
+        {
+          type: @type
+        }
+      end
+
+      def handle(collector)
+        raise NotImplementedError, "implement in subclasses"
+      end
+    end
+
+    # @api private
+    module Events
+      class TrackRequest < Event
+        register "track_request"
+
+        def self.from_json(data)
+          new
+        end
+
+        def handle(collector)
+          collector.handle_track_request
+        end
+
+        def inspect
+          "#<#{self.class.name}>"
+        end
+      end
+
+      class TrackAttackWave < Event
+        register "track_attack_wave"
+
+        def self.from_json(data)
+          new(
+            being_blocked: data[:being_blocked]
+          )
+        end
+
+        def initialize(being_blocked:)
+          super()
+          @being_blocked = being_blocked
+        end
+
+        def as_json
+          super.update({
+            being_blocked: @being_blocked
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_attack_wave(being_blocked: @being_blocked)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@being_blocked}>"
+        end
+      end
+
+      class TrackScan < Event
+        register "track_scan"
+
+        def self.from_json(data)
+          new(
+            data[:sink_name],
+            data[:duration],
+            has_errors: data[:has_errors]
+          )
+        end
+
+        def initialize(sink_name, duration, has_errors:)
+          super()
+          @sink_name = sink_name
+          @duration = duration
+          @has_errors = has_errors
+        end
+
+        def as_json
+          super.update({
+            sink_name: @sink_name,
+            duration: @duration,
+            has_errors: @has_errors
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_scan(@sink_name, @duration, has_errors: @has_errors)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@sink_name} #{format "%0.6f", @duration} #{@has_errors}>"
+        end
+      end
+
+      class TrackAttack < Event
+        register "track_attack"
+
+        def self.from_json(data)
+          new(
+            data[:sink_name],
+            being_blocked: data[:being_blocked]
+          )
+        end
+
+        def initialize(sink_name, being_blocked:)
+          super()
+          @sink_name = sink_name
+          @being_blocked = being_blocked
+        end
+
+        def as_json
+          super.update({
+            sink_name: @sink_name,
+            being_blocked: @being_blocked
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_attack(@sink_name, being_blocked: @being_blocked)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@sink_name} #{@being_blocked}>"
+        end
+      end
+
+      class TrackUser < Event
+        register "track_user"
+
+        def self.from_json(data)
+          new(Aikido::Zen::Actor.from_json(data[:actor]))
+        end
+
+        def initialize(actor)
+          super()
+          @actor = actor
+        end
+
+        def as_json
+          super.update({
+            actor: @actor.as_json
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_user(@actor)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@actor.id} #{@actor.name}>"
+        end
+      end
+
+      class TrackOutbound < Event
+        register "track_outbound"
+
+        def self.from_json(data)
+          new(OutboundConnection.from_json(data[:connection]))
+        end
+
+        def initialize(connection)
+          super()
+          @connection = connection
+        end
+
+        def as_json
+          super.update({
+            connection: @connection.as_json
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_outbound(@connection)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@connection.host}:#{@connection.port}>"
+        end
+      end
+
+      class TrackRoute < Event
+        register "track_route"
+
+        def self.from_json(data)
+          new(
+            Route.from_json(data[:route]),
+            Request::Schema.from_json(data[:schema])
+          )
+        end
+
+        def initialize(route, schema)
+          super()
+          @route = route
+          @schema = schema
+        end
+
+        def as_json
+          super.update({
+            route: @route.as_json,
+            schema: @schema.as_json
+          })
+        end
+
+        def handle(collector)
+          collector.handle_track_route(@route, @schema)
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@route.verb} #{@route.path.inspect}>"
+        end
+      end
+    end
+  end
+end

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

@@ -7,9 +7,24 @@ module Aikido::Zen
   #
   # Keeps track of the hostnames to which the app has made outbound HTTP
   # requests.
-  class Collector::Hosts < Aikido::Zen::CappedSet
+  class Collector::Hosts < Aikido::Zen::CappedMap
     def initialize(config = Aikido::Zen.config)
       super(config.max_outbound_connections)
     end
+
+    # @param host [Aikido::Zen::OutboundConnection]
+    # @return [void]
+    def add(host)
+      self[host] ||= host
+      self[host].hit
+    end
+
+    def each(&blk)
+      each_value(&blk)
+    end
+
+    def as_json
+      map(&:as_json)
+    end
   end
 end

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

@@ -7,6 +7,8 @@ module Aikido::Zen
   #
   # Keeps track of the visited routes.
   class Collector::Routes
+    # @api private
+    # Visible for testing.
     attr_reader :visits
 
     def initialize(config = Aikido::Zen.config)
@@ -14,11 +16,11 @@ module Aikido::Zen
       @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
+    # @param route [Aikido::Zen::Route] the route of the request
+    # @param schema [Aikido::Zen::Request::Schema] the schema for the request
+    # @return [void]
+    def add(route, schema)
+      @visits[route].increment(schema) unless route.nil?
     end
 
     def as_json
@@ -33,6 +35,7 @@ module Aikido::Zen
     end
 
     # @api private
+    # Visible for testing.
     def [](route)
       @visits[route]
     end
@@ -49,16 +52,18 @@ module Aikido::Zen
         @config = config
       end
 
-      def increment(request)
+      def increment(schema)
         self.hits += 1
 
         if sample_schema?
           self.samples += 1
-          self.schema |= request.schema
+          self.schema |= schema
         end
       end
 
-      private def sample_schema?
+      private
+
+      def sample_schema?
         samples < @config.api_schema_max_samples
       end
     end

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

@@ -8,7 +8,7 @@ module Aikido::Zen
   # 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
+    attr_reader :started_at, :ended_at, :requests, :aborted_requests, :attack_waves, :blocked_attack_waves, :sinks
 
     # @!visibility private
     attr_writer :ended_at
@@ -16,10 +16,13 @@ module Aikido::Zen
     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
+      @attack_waves = 0
+      @blocked_attack_waves = 0
+      @sinks = Hash.new { |h, k| h[k] = Collector::SinkStats.new(k, @config) }
     end
 
     # @return [Boolean]
@@ -35,10 +38,9 @@ module Aikido::Zen
     # Track the timestamp we start tracking this series of stats.
     #
     # @param at [Time]
-    # @return [self]
+    # @return [void]
     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
@@ -47,7 +49,7 @@ module Aikido::Zen
     #
     # @param at [Time] the time at which we're resetting, which is set as the
     #   ending time for the returned copy.
-    # @return [self]
+    # @return [void]
     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.
@@ -56,31 +58,36 @@ module Aikido::Zen
       freeze
     end
 
-    # @return [self]
+    # @return [void]
     def add_request
       @requests += 1
-      self
     end
 
-    # @param scan [Aikido::Zen::Scan]
-    # @return [self]
-    def add_scan(scan)
-      stats = @sinks[scan.sink.name]
+    # @param being_blocked [Boolean] whether the Agent blocked the request
+    # @return [void]
+    def add_attack_wave(being_blocked:)
+      @attack_waves += 1
+      @blocked_attack_waves += 1 if being_blocked
+    end
+
+    # @param sink_name [String] the name of the sink
+    # @param duration [Float] the length the scan in seconds
+    # @param has_errors [Boolean] whether errors occurred during the scan
+    # @return [void]
+    def add_scan(sink_name, duration, has_errors:)
+      stats = @sinks[sink_name]
       stats.scans += 1
-      stats.errors += 1 if scan.errors?
-      stats.add_timing(scan.duration)
-      self
+      stats.errors += 1 if has_errors
+      stats.add_timing(duration)
     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]
+    # @param sink_name [String] the name of the sink
+    # @param being_blocked [Boolean] whether the Agent blocked the request
+    # @return [void]
+    def add_attack(sink_name, being_blocked:)
+      stats = @sinks[sink_name]
       stats.attacks += 1
       stats.blocked_attacks += 1 if being_blocked
-      self
     end
 
     def as_json
@@ -88,13 +95,17 @@ module Aikido::Zen
       {
         startedAt: @started_at.to_i * 1000,
         endedAt: (@ended_at.to_i * 1000 if @ended_at),
-        sinks: @sinks.transform_values(&:as_json),
+        operations: @sinks.transform_values(&:as_json),
         requests: {
           total: @requests,
           aborted: @aborted_requests,
           attacksDetected: {
             total: total_attacks,
             blocked: total_blocked
+          },
+          attackWaves: {
+            total: @attack_waves,
+            blocked: @blocked_attack_waves
           }
         }
       }

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

@@ -11,16 +11,18 @@ module Aikido::Zen
       super(config.max_users_tracked)
     end
 
+    # @param actor [Aikido::Zen::Actor]
+    # @return [void]
     def add(actor)
       if key?(actor.id)
-        self[actor.id].update
+        self[actor.id] |= actor
       else
         self[actor.id] = actor
       end
     end
 
-    def each(&b)
-      each_value(&b)
+    def each(&blk)
+      each_value(&blk)
     end
 
     def as_json

data/lib/aikido/zen/collector.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "collector/event"
+
 module Aikido::Zen
   # Handles collecting all the runtime statistics to report back to the Aikido
   # servers.
@@ -7,14 +9,44 @@ module Aikido::Zen
     def initialize(config: Aikido::Zen.config)
       @config = config
 
+      @events = Queue.new
+
       @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))
-      @heartbeats = Queue.new
+
       @middleware_installed = Concurrent::AtomicBoolean.new
     end
 
+    # Add an event, to be handled in the collector in current process or sent
+    # to a collector in another process.
+    #
+    # @param event [Aikido::Zen::Collector::Event] the event to add
+    # @return [void]
+    def add_event(event)
+      @events << event
+    end
+
+    # @return [Boolean] whether the collector has any events
+    def has_events?
+      !@events.empty?
+    end
+
+    # Flush all events.
+    #
+    # @return [Array<Aikido::Zen::Collector::Event>]
+    def flush_events
+      Array.new(@events.size) { @events.pop }
+    end
+
+    # Handle the events in the queue.
+    #
+    # @return [void]
+    def handle
+      flush_events.each { |event| event.handle(self) }
+    end
+
     # Flush all the stats into a Heartbeat event that can be reported back to
     # the Aikido servers.
     #
@@ -22,6 +54,8 @@ module Aikido::Zen
     #   of the new stats collection period. Defaults to now.
     # @return [Aikido::Zen::Events::Heartbeat]
     def flush(at: Time.now.utc)
+      handle
+
       stats = @stats.get_and_set(Stats.new(@config))
       users = @users.get_and_set(Users.new(@config))
       hosts = @hosts.get_and_set(Hosts.new(@config))
@@ -30,21 +64,11 @@ module Aikido::Zen
       start(at: at)
       stats = stats.flush(at: at)
 
-      Events::Heartbeat.new(
+      Aikido::Zen::Events::Heartbeat.new(
         stats: stats, users: users, hosts: hosts, routes: routes, middleware_installed: middleware_installed?
       )
     end
 
-    # Put heartbeats coming from child processes into the internal queue.
-    def push_heartbeat(heartbeat)
-      @heartbeats << heartbeat
-    end
-
-    # Drains into an array all the queued heartbeats
-    def flush_heartbeats
-      Array.new(@heartbeats.size) { @heartbeats.pop }
-    end
-
     # Sets the start time for this collection period.
     #
     # @param at [Time] defaults to now.
@@ -55,16 +79,27 @@ module Aikido::Zen
 
     # Track stats about the requests
     #
-    # @param request [Aikido::Zen::Request]
     # @return [void]
-    def track_request(*)
+    def track_request
+      add_event(Events::TrackRequest.new)
+    end
+
+    def handle_track_request
       synchronize(@stats) { |stats| stats.add_request }
     end
 
-    #  Record the visited endpoint, and if enabled, the API schema for this endpoint.
-    # @param request [Aikido::Zen::Request]
-    def track_route(request)
-      synchronize(@routes) { |routes| routes.add(request) if request.route }
+    # Track stats about an attack detected by our scanners.
+    #
+    # @param attack [Aikido::Zen::Events::AttackWave]
+    # @return [void]
+    def track_attack_wave(being_blocked:)
+      add_event(Events::TrackAttackWave.new(being_blocked: being_blocked))
+    end
+
+    def handle_track_attack_wave(being_blocked:)
+      synchronize(@stats) do |stats|
+        stats.add_attack_wave(being_blocked: being_blocked)
+      end
     end
 
     # Track stats about a scan performed by one of our sinks.
@@ -72,7 +107,13 @@ module Aikido::Zen
     # @param scan [Aikido::Zen::Scan]
     # @return [void]
     def track_scan(scan)
-      synchronize(@stats) { |stats| stats.add_scan(scan) }
+      add_event(Events::TrackScan.new(scan.sink.name, scan.duration, has_errors: scan.errors?))
+    end
+
+    def handle_track_scan(sink_name, duration, has_errors:)
+      synchronize(@stats) do |stats|
+        stats.add_scan(sink_name, duration, has_errors: has_errors)
+      end
     end
 
     # Track stats about an attack detected by our scanners.
@@ -80,25 +121,49 @@ module Aikido::Zen
     # @param attack [Aikido::Zen::Attack]
     # @return [void]
     def track_attack(attack)
+      add_event(Events::TrackAttack.new(attack.sink.name, being_blocked: attack.blocked?))
+    end
+
+    def handle_track_attack(sink_name, being_blocked:)
       synchronize(@stats) do |stats|
-        stats.add_attack(attack, being_blocked: attack.blocked?)
+        stats.add_attack(sink_name, being_blocked: being_blocked)
       end
     end
 
+    # Track the user reported by the developer to be behind this request.
+    #
+    # @param actor [Aikido::Zen::Actor]
+    # @return [void]
+    def track_user(actor)
+      add_event(Events::TrackUser.new(actor))
+    end
+
+    def handle_track_user(actor)
+      synchronize(@users) { |users| users.add(actor) }
+    end
+
     # Track an HTTP connections to an external host.
     #
     # @param connection [Aikido::Zen::OutboundConnection]
     # @return [void]
     def track_outbound(connection)
+      add_event(Events::TrackOutbound.new(connection))
+    end
+
+    def handle_track_outbound(connection)
       synchronize(@hosts) { |hosts| hosts.add(connection) }
     end
 
-    # Track the user reported by the developer to be behind this request.
+    # Record the visited endpoint, and if enabled, the API schema for this endpoint.
     #
-    # @param actor [Aikido::Zen::Actor]
+    # @param request [Aikido::Zen::Request]
     # @return [void]
-    def track_user(actor)
-      synchronize(@users) { |users| users.add(actor) }
+    def track_route(request)
+      add_event(Events::TrackRoute.new(request.route, request.schema))
+    end
+
+    def handle_track_route(route, schema)
+      synchronize(@routes) { |routes| routes.add(route, schema) }
     end
 
     def middleware_installed!
@@ -106,26 +171,40 @@ module Aikido::Zen
     end
 
     # @api private
-    def routes
-      @routes.get
+    #
+    # @note Visible for testing.
+    def stats
+      handle
+      @stats.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def users
+      handle
       @users.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def hosts
+      handle
       @hosts.get
     end
 
     # @api private
-    def stats
-      @stats.get
+    #
+    # @note Visible for testing.
+    def routes
+      handle
+      @routes.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def middleware_installed?
       @middleware_installed.true?
     end