aikido-zen 1.0.1.beta.2-arm64-linux-musl

data/lib/aikido/zen/background_worker.rb

@@ -0,0 +1,52 @@
+module Aikido::Zen
+  # Generic background worker class backed by queue. Meant to be used by any
+  # background process that needs to do heavy tasks.
+  class BackgroundWorker
+    # @param block [block] A block that receives 1 message directly from the queue
+    def initialize(&block)
+      @queue = Queue.new
+      @block = block
+    end
+
+    # starts the background thread, blocking the thread until a new messages arrives
+    # or the queue is stopped.
+    def start
+      @thread = Thread.new do
+        while running? || actions?
+          action = wait_for_action
+          @block.call(action) unless action.nil?
+        end
+      end
+    end
+
+    def restart
+      stop
+      @queue = Queue.new # re-open the queue
+      start
+    end
+
+    # Drain the queue to do not lose any messages
+    def stop
+      @queue.close # stop accepting messages
+      @thread.join # wait for the queue to be drained
+    end
+
+    def enqueue(scan)
+      @queue.push(scan)
+    end
+
+    private
+
+    def actions?
+      !@queue.empty?
+    end
+
+    def running?
+      !@queue.closed?
+    end
+
+    def wait_for_action
+      @queue.pop(false)
+    end
+  end
+end

data/lib/aikido/zen/capped_collections.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "forwardable"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Provides a FIFO set with a maximum size. Adding an element after the
+  # capacity has been reached kicks the oldest element in the set out,
+  # while maintaining the uniqueness property of a set (relying on #eql?
+  # and #hash).
+  class CappedSet
+    include Enumerable
+    extend Forwardable
+
+    def_delegators :@data, :size, :empty?
+
+    # @return [Integer]
+    attr_reader :capacity
+
+    def initialize(capacity)
+      @data = CappedMap.new(capacity)
+    end
+
+    def <<(element)
+      @data[element] = nil
+      self
+    end
+    alias_method :add, :<<
+    alias_method :push, :<<
+
+    def each(&b)
+      @data.each_key(&b)
+    end
+
+    def as_json
+      map(&:as_json)
+    end
+  end
+
+  # @api private
+  #
+  # Provides a FIFO hash-like structure with a maximum size. Adding a new key
+  # after the capacity has been reached kicks the first element pair added out.
+  class CappedMap
+    include Enumerable
+    extend Forwardable
+
+    def_delegators :@data,
+      :[], :fetch, :delete, :key?,
+      :each, :each_key, :each_value,
+      :size, :empty?, :to_hash
+
+    # @return [Integer]
+    attr_reader :capacity
+
+    def initialize(capacity)
+      raise ArgumentError, "cannot set capacity lower than 1: #{capacity}" if capacity < 1
+      @capacity = capacity
+      @data = {}
+    end
+
+    def []=(key, value)
+      @data[key] = value
+      @data.delete(@data.each_key.first) if @data.size > @capacity
+    end
+  end
+end

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,66 @@
+# frozen_string_literal: true
+
+require_relative "../request/schema/empty_schema"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Keeps track of the visited routes.
+  class Collector::Routes
+    attr_reader :visits
+
+    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/collector/sink_stats.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require_relative "../capped_collections"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Tracks data specific to a single Sink.
+  class Collector::SinkStats
+    # @return [Integer] number of total calls to Sink#scan.
+    attr_accessor :scans
+
+    # @return [Integer] number of scans where our scanners raised an
+    #   error that was handled.
+    attr_accessor :errors
+
+    # @return [Integer] number of scans where an attack was detected.
+    attr_accessor :attacks
+
+    # @return [Integer] number of scans where an attack was detected
+    #   _and_ blocked by the Zen.
+    attr_accessor :blocked_attacks
+
+    # @return [Set<Float>] keeps the duration of individual scans. If
+    #   this grows to match Config#max_performance_samples, the set is
+    #   cleared and the data is aggregated into #compressed_timings.
+    attr_accessor :timings
+
+    # @return [Array<CompressedTiming>] list of aggregated stats.
+    attr_accessor :compressed_timings
+
+    def initialize(name, config)
+      @name = name
+      @config = config
+
+      @scans = 0
+      @errors = 0
+
+      @attacks = 0
+      @blocked_attacks = 0
+
+      @timings = Set.new
+      @compressed_timings = CappedSet.new(@config.max_compressed_stats)
+    end
+
+    def add_timing(duration)
+      compress_timings if @timings.size >= @config.max_performance_samples
+      @timings << duration
+    end
+
+    def compress_timings(at: Time.now.utc)
+      return if @timings.empty?
+
+      list = @timings.sort
+      @timings.clear
+
+      mean = list.sum / list.size
+      percentiles = percentiles(list, 50, 75, 90, 95, 99)
+
+      @compressed_timings << CompressedTiming.new(mean, percentiles, at)
+    end
+
+    def as_json
+      {
+        total: @scans,
+        interceptorThrewError: @errors,
+        withoutContext: 0,
+        attacksDetected: {
+          total: @attacks,
+          blocked: @blocked_attacks
+        },
+        compressedTimings: @compressed_timings.as_json
+      }
+    end
+
+    private def percentiles(sorted, *scores)
+      return {} if sorted.empty? || scores.empty?
+
+      scores.map { |p|
+        index = (sorted.size * (p / 100.0)).floor
+        [p, sorted.at(index)]
+      }.to_h
+    end
+
+    CompressedTiming = Struct.new(:mean, :percentiles, :compressed_at) do
+      def as_json
+        {
+          averageInMs: mean * 1000,
+          percentiles: percentiles.transform_values { |t| t * 1000 },
+          compressedAt: compressed_at.to_i * 1000
+        }
+      end
+    end
+  end
+end

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/collector/users.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require_relative "../capped_collections"
+
+module Aikido::Zen
+  # @api private
+  #
+  # Keeps track of the users that were seen by the app.
+  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
+      else
+        self[actor.id] = actor
+      end
+    end
+
+    def each(&b)
+      each_value(&b)
+    end
+
+    def as_json
+      map(&:as_json)
+    end
+  end
+end

data/lib/aikido/zen/collector.rb

@@ -0,0 +1,144 @@
+# 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))
+      @heartbeats = Queue.new
+      @middleware_installed = Concurrent::AtomicBoolean.new
+    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, 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.
+    # @return [void]
+    def start(at: Time.now.utc)
+      synchronize(@stats) { |stats| stats.start(at) }
+    end
+
+    # Track stats about the requests
+    #
+    # @param request [Aikido::Zen::Request]
+    # @return [void]
+    def 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 }
+    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
+
+    def middleware_installed!
+      @middleware_installed.make_true
+    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
+
+    # @api private
+    def middleware_installed?
+      @middleware_installed.true?
+    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"