aikido-zen 0.1.0.alpha4-x86_64-darwin → 0.1.0-x86_64-darwin

data/lib/aikido/zen/config.rb

@@ -8,6 +8,13 @@ require_relative "context"
 
 module Aikido::Zen
   class Config
+    # @return [Boolean] whether Aikido should be turned completely off (no
+    #   intercepting calls to protect the app, no agent process running, no
+    #   middleware installed). Defaults to false (so, enabled). Can be set
+    #   via the AIKIDO_DISABLED environment variable.
+    attr_accessor :disabled
+    alias_method :disabled?, :disabled
+
     # @return [Boolean] whether Aikido should only report infractions or block
     #   the request by raising an Exception. Defaults to whether AIKIDO_BLOCKING
     #   is set to a non-empty value in your environment, or +false+ otherwise.
@@ -45,7 +52,7 @@ module Aikido::Zen
     #   into an Object. Defaults to the standard library's JSON.parse method.
     attr_accessor :json_decoder
 
-    # @returns [Logger]
+    # @return [Logger]
     attr_accessor :logger
 
     # @return [Integer] maximum number of timing measurements to keep in memory
@@ -83,11 +90,9 @@ module Aikido::Zen
     #   differentiate different clients. By default this uses the request IP.
     attr_accessor :rate_limiting_discriminator
 
-    # @return [Boolean] whether Zen should infer the schema from request bodies
-    #   sent to the app. Defaults to +false+ or the value of the environment
-    #   variable AIKIDO_FEATURE_COLLECT_API_SCHEMA.
-    attr_accessor :api_schema_collection_enabled
-    alias_method :api_schema_collection_enabled?, :api_schema_collection_enabled
+    # @return [Integer] max number of requests we sample per endpoint when
+    #   computing the schema.
+    attr_accessor :api_schema_max_samples
 
     # @api private
     # @return [Integer] max number of levels deep we want to read a nested
@@ -125,7 +130,8 @@ module Aikido::Zen
     attr_accessor :imds_allowed_hosts
 
     def initialize
-      self.blocking_mode = !!ENV.fetch("AIKIDO_BLOCKING", false)
+      self.disabled = read_boolean_from_env(ENV.fetch("AIKIDO_DISABLED", false))
+      self.blocking_mode = read_boolean_from_env(ENV.fetch("AIKIDO_BLOCKING", false))
       self.api_timeouts = 10
       self.api_base_url = ENV.fetch("AIKIDO_BASE_URL", DEFAULT_API_BASE_URL)
       self.runtime_api_base_url = ENV.fetch("AIKIDO_RUNTIME_URL", DEFAULT_RUNTIME_BASE_URL)
@@ -146,11 +152,9 @@ module Aikido::Zen
       self.server_rate_limit_deadline = 1800 # 30 min
       self.client_rate_limit_period = 3600 # 1 hour
       self.client_rate_limit_max_events = 100
-
-      self.api_schema_collection_enabled = read_boolean_from_env(ENV.fetch("AIKIDO_FEATURE_COLLECT_API_SCHEMA", false))
+      self.api_schema_max_samples = Integer(ENV.fetch("AIKIDO_MAX_API_DISCOVERY_SAMPLES", 10))
       self.api_schema_collection_max_depth = 20
       self.api_schema_collection_max_properties = 20
-
       self.imds_allowed_hosts = ["metadata.google.internal", "metadata.goog"]
     end
 
@@ -224,6 +228,8 @@ module Aikido::Zen
     end
 
     # @!visibility private
-    DEFAULT_RATE_LIMITING_DISCRIMINATOR = ->(request) { request.ip }
+    DEFAULT_RATE_LIMITING_DISCRIMINATOR = ->(request) {
+      request.actor ? "actor:#{request.actor.id}" : request.ip
+    }
   end
 end

data/lib/aikido/zen/context.rb

@@ -7,13 +7,20 @@ require_relative "payload"
 
 module Aikido::Zen
   class Context
+    # Build a Context object for the current HTTP request based on the currently
+    # configured request builder.
+    #
+    # @param env [Hash] the Rack env hash.
+    # @param config [Aikido::Zen::Config]
+    # @return [Aikido::Zen::Context]
     def self.from_rack_env(env, config = Aikido::Zen.config)
       config.request_builder.call(env)
     end
 
+    # @return [Aikido::Zen::Request]
     attr_reader :request
 
-    # @param [Rack::Request] a Request object that implements the
+    # @param request [Rack::Request] a Request object that implements the
     #   Rack::Request API, to which we will delegate behavior.
     # @param settings [Aikido::Zen::RuntimeSettings]
     #

data/lib/aikido/zen/errors.rb

@@ -3,11 +3,13 @@
 require "forwardable"
 
 module Aikido
+  # @!visibility private
   # Support rescuing Aikido::Error without forcing a single base class to all
   # errors (so things that should be e.g. a TypeError, can have the correct
   # superclass).
-  module Error; end
+  module Error; end # :nodoc:
 
+  # @!visibility private
   # Generic error for problems with the Agent.
   class ZenError < RuntimeError
     include Error

data/lib/aikido/zen/event.rb

@@ -48,17 +48,20 @@ module Aikido::Zen
     end
 
     class Heartbeat < Event
-      def initialize(stats:, **opts)
+      def initialize(stats:, users:, hosts:, routes:, **opts)
         super(type: "heartbeat", **opts)
         @stats = stats
+        @users = users
+        @hosts = hosts
+        @routes = routes
       end
 
       def as_json
         super.update(
           stats: @stats.as_json,
-          routes: @stats.routes.as_json,
-          hostnames: @stats.outbound_connections.as_json,
-          users: @stats.users.as_json
+          users: @users.as_json,
+          routes: @routes.as_json,
+          hostnames: @hosts.as_json
         )
       end
     end

data/lib/aikido/zen/middleware/set_context.rb

@@ -3,6 +3,9 @@
 require_relative "../context"
 
 module Aikido::Zen
+  # @!visibility private
+  ENV_KEY = "aikido.context"
+
   module Middleware
     # Rack middleware that keeps the current context in a Thread/Fiber-local
     # variable so that other parts of the agent/firewall can access it.
@@ -14,7 +17,7 @@ module Aikido::Zen
       def call(env)
         context = Context.from_rack_env(env)
 
-        Aikido::Zen.current_context = context
+        Aikido::Zen.current_context = env[ENV_KEY] = context
         Aikido::Zen.track_request(context.request)
 
         @app.call(env)

data/lib/aikido/zen/rails_engine.rb

@@ -5,44 +5,53 @@ require "action_dispatch"
 module Aikido::Zen
   class RailsEngine < ::Rails::Engine
     config.before_configuration do
-      # Access library configuration at `Rails.application.config.aikido_zen`.
-      config.aikido_zen = Aikido::Zen.config
+      # Access library configuration at `Rails.application.config.zen`.
+      config.zen = Aikido::Zen.config
     end
 
     initializer "aikido.add_middleware" do |app|
+      next if config.zen.disabled?
+
       app.middleware.use Aikido::Zen::Middleware::SetContext
       app.middleware.use Aikido::Zen::Middleware::CheckAllowedAddresses
-      app.middleware.use Aikido::Zen::Middleware::Throttler
-
-      # Due to how Rails sets up its middleware chain, the routing is evaluated
-      # (and the Request object constructed) in the app that terminates the
-      # chain, so no amount of middleware will be able to access it.
-      #
-      # This way, we overwrite the Request object as early as we can in the
-      # request handling, so that by the time we start evaluating inputs, we
-      # have assigned the request correctly.
+
       ActiveSupport.on_load(:action_controller) do
+        # Due to how Rails sets up its middleware chain, the routing is evaluated
+        # (and the Request object constructed) in the app that terminates the
+        # chain, so no amount of middleware will be able to access it.
+        #
+        # This way, we overwrite the Request object as early as we can in the
+        # request handling, so that by the time we start evaluating inputs, we
+        # have assigned the request correctly.
         before_action { Aikido::Zen.current_context.update_request(request) }
       end
     end
 
     initializer "aikido.configuration" do |app|
-      app.config.aikido_zen.logger = ::Rails.logger.tagged("aikido")
-      app.config.aikido_zen.request_builder = Aikido::Zen::Context::RAILS_REQUEST_BUILDER
+      # Allow the logger to be configured before checking if disabled? so we can
+      # let the user know that the agent is disabled.
+      app.config.zen.logger = ::Rails.logger.tagged("aikido")
+
+      next if config.zen.disabled?
+
+      app.config.zen.request_builder = Aikido::Zen::Context::RAILS_REQUEST_BUILDER
 
       # Plug Rails' JSON encoder/decoder, but only if the user hasn't changed
       # them for something else.
-      if app.config.aikido_zen.json_encoder == Aikido::Zen::Config::DEFAULT_JSON_ENCODER
-        app.config.aikido_zen.json_encoder = ActiveSupport::JSON.method(:encode)
+      if app.config.zen.json_encoder == Aikido::Zen::Config::DEFAULT_JSON_ENCODER
+        app.config.zen.json_encoder = ActiveSupport::JSON.method(:encode)
       end
 
-      if app.config.aikido_zen.json_decoder == Aikido::Zen::Config::DEFAULT_JSON_DECODER
-        app.config.aikido_zen.json_decoder = ActiveSupport::JSON.method(:decode)
+      if app.config.zen.json_decoder == Aikido::Zen::Config::DEFAULT_JSON_DECODER
+        app.config.zen.json_decoder = ActiveSupport::JSON.method(:decode)
       end
     end
 
     config.after_initialize do
-      Aikido::Zen.initialize!
+      if config.zen.disabled?
+        config.zen.logger.warn("Zen has been disabled and will not run.")
+        next
+      end
 
       # Make sure this is run at the end of the initialization process, so
       # that any gems required after aikido-zen are detected and patched

data/lib/aikido/zen/request/schema/builder.rb

@@ -16,8 +16,6 @@ module Aikido::Zen
       end
 
       def schema
-        return unless @config.api_schema_collection_enabled?
-
         Request::Schema.new(
           content_type: body_data_type,
           body_schema: body_schema,

data/lib/aikido/zen/request.rb

@@ -11,6 +11,12 @@ module Aikido::Zen
     # @return [Aikido::Zen::Router]
     attr_reader :router
 
+    # The current user, if set by the host app.
+    #
+    # @return [Aikido::Zen::Actor, nil]
+    # @see Aikido::Zen.track_user
+    attr_accessor :actor
+
     def initialize(delegate, framework:, router:)
       super(delegate)
       @framework = framework

data/lib/aikido/zen/runtime_settings.rb

@@ -11,16 +11,11 @@ module Aikido::Zen
   #
   # You can subscribe to changes with +#add_observer(object, func_name)+, which
   # will call the function passing the settings as an argument.
-  class RuntimeSettings < Concurrent::MutableStruct.new(
-    :updated_at, :heartbeat_interval, :endpoints, :blocked_user_ids, :skip_protection_for_ips, :received_any_stats
-  )
-    include Concurrent::Concern::Observable
-
+  RuntimeSettings = Struct.new(:updated_at, :heartbeat_interval, :endpoints, :blocked_user_ids, :skip_protection_for_ips, :received_any_stats) do
     def initialize(*)
-      self.observers = Concurrent::Collection::CopyOnWriteObserverSet.new
       super
-      self.endpoints ||= Endpoints.new
-      self.skip_protection_for_ips ||= IPSet.new
+      self.endpoints ||= RuntimeSettings::Endpoints.new
+      self.skip_protection_for_ips ||= RuntimeSettings::IPSet.new
     end
 
     # @!attribute [rw] updated_at
@@ -56,12 +51,12 @@ module Aikido::Zen
 
       self.updated_at = Time.at(data["configUpdatedAt"].to_i / 1000)
       self.heartbeat_interval = (data["heartbeatIntervalInMS"].to_i / 1000)
-      self.endpoints = Endpoints.from_json(data["endpoints"])
+      self.endpoints = RuntimeSettings::Endpoints.from_json(data["endpoints"])
       self.blocked_user_ids = data["blockedUserIds"]
-      self.skip_protection_for_ips = IPSet.from_json(data["allowedIPAddresses"])
+      self.skip_protection_for_ips = RuntimeSettings::IPSet.from_json(data["allowedIPAddresses"])
       self.received_any_stats = data["receivedAnyStats"]
 
-      observers.notify_observers(self) if updated_at != last_updated_at
+      Aikido::Zen.agent.updated_settings! if updated_at != last_updated_at
     end
   end
 end

data/lib/aikido/zen/sinks/action_controller.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  module Sinks
+    module ActionController
+      # Implements the "middleware" for rate limiting in Rails apps, where we
+      # need to check at the end of the `before_action` chain, rather than in
+      # an actual Rack middleware, to allow for calls to Zen.track_user being
+      # made from before_actions in the host app, thus allowing rate-limiting
+      # by user ID rather than solely by IP.
+      class Throttler
+        def initialize(
+          config: Aikido::Zen.config,
+          settings: Aikido::Zen.runtime_settings,
+          rate_limiter: Aikido::Zen::RateLimiter.new
+        )
+          @config = config
+          @settings = settings
+          @rate_limiter = rate_limiter
+        end
+
+        def throttle(controller)
+          context = controller.request.env[Aikido::Zen::ENV_KEY]
+          request = context.request
+
+          if should_throttle?(request)
+            status, headers, body = @config.rate_limited_responder.call(request)
+            controller.headers.update(headers)
+            controller.render plain: Array(body).join, status: status
+
+            return true
+          end
+
+          false
+        end
+
+        private def should_throttle?(request)
+          return false if @settings.skip_protection_for_ips.include?(request.ip)
+
+          @rate_limiter.throttle?(request)
+        end
+      end
+
+      def self.throttler
+        @throttler ||= Aikido::Zen::Sinks::ActionController::Throttler.new
+      end
+
+      module Extensions
+        def run_callbacks(kind, *)
+          return super unless kind == :process_action
+
+          super do
+            rate_limiter = Aikido::Zen::Sinks::ActionController.throttler
+            throttled = rate_limiter.throttle(self)
+
+            yield if block_given? && !throttled
+          end
+        end
+      end
+    end
+  end
+end
+
+::AbstractController::Callbacks.prepend(Aikido::Zen::Sinks::ActionController::Extensions)

data/lib/aikido/zen/sinks.rb

@@ -4,6 +4,7 @@ require_relative "sink"
 
 require_relative "sinks/socket"
 
+require_relative "sinks/action_controller" if defined?(::ActionController)
 require_relative "sinks/resolv" if defined?(::Resolv)
 require_relative "sinks/net_http" if defined?(::Net::HTTP)
 require_relative "sinks/http" if defined?(::HTTP)

data/lib/aikido/zen/version.rb

@@ -2,9 +2,9 @@
 
 module Aikido
   module Zen
-    VERSION = "0.1.0.alpha4"
+    VERSION = "0.1.0"
 
     # The version of libzen_internals that we build against.
-    LIBZEN_VERSION = "0.1.26"
+    LIBZEN_VERSION = "0.1.30"
   end
 end

data/lib/aikido/zen/worker.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+module Aikido::Zen
+  # @api private
+  #
+  # The worker manages the background thread in which Zen communicates with the
+  # Aikido server.
+  class Worker
+    # @return [Concurrent::ExecutorService]
+    attr_reader :executor
+
+    # @!visibility private
+    attr_reader :timers, :deferrals
+
+    def initialize(config: Aikido::Zen.config)
+      @config = config
+      @timers = []
+      @deferrals = []
+      @executor = Concurrent::SingleThreadExecutor.new
+    end
+
+    # Queue a block to be run asynchronously in the background thread.
+    #
+    # @return [void]
+    def perform(&block)
+      executor.post do
+        yield
+      rescue Exception => err # rubocop:disable Lint/RescueException
+        @config.logger.error "Error in background worker: #{err.inspect}"
+      end
+    end
+
+    # Queue a block to be run asynchronously after a delay.
+    #
+    # @param interval [Integer] amount of seconds to wait.
+    # @return [void]
+    def delay(interval, &task)
+      Concurrent::ScheduledTask
+        .execute(interval, executor: executor) { perform(&task) }
+        .tap { |deferral| @deferrals << deferral }
+    end
+
+    # Queue a block to run repeatedly on a timer on the background thread. The
+    # timer will consider how long the block takes to run to schedule the next
+    # run. For example, if you schedule a block to run every 10 seconds, and the
+    # block itself takes 2 seconds, the second iteration will be run 8 seconds
+    # after the first one.
+    #
+    # If the block takes longer than the given interval, the second iteration
+    # will be run immediately.
+    #
+    # @param interval [Integer] amount of seconds to wait between runs.
+    # @param run_now [Boolean] whether to run the block immediately, or wait for
+    #   +interval+ seconds before the first run. Defaults to +true+.
+    # @return [void]
+    def every(interval, run_now: true, &task)
+      Concurrent::TimerTask
+        .execute(
+          run_now: run_now,
+          executor: executor,
+          interval_type: :fixed_rate,
+          execution_interval: interval
+        ) {
+          perform(&task)
+        }
+        .tap { |timer| @timers << timer }
+    end
+
+    # Safely clean up and kill the thread, giving time to kill any ongoing tasks
+    # on the queue.
+    #
+    # @return [void]
+    def shutdown
+      @deferrals.each { |task| task.cancel if task.pending? }
+      @timers.each { |task| task.shutdown }
+      @executor.shutdown
+      @executor.wait_for_termination(30)
+    end
+  end
+end

data/lib/aikido/zen.rb

@@ -4,7 +4,9 @@ require_relative "zen/version"
 require_relative "zen/errors"
 require_relative "zen/actor"
 require_relative "zen/config"
+require_relative "zen/collector"
 require_relative "zen/system_info"
+require_relative "zen/worker"
 require_relative "zen/agent"
 require_relative "zen/api_client"
 require_relative "zen/context"
@@ -24,12 +26,24 @@ module Aikido
       @config ||= Config.new
     end
 
+    # @return [Aikido::Zen::RuntimeSettings] the firewall configuration sourced
+    #   from your Aikido dashboard. This is periodically polled for updates.
+    def self.runtime_settings
+      @runtime_settings ||= RuntimeSettings.new
+    end
+
     # Gets information about the current system configuration, which is sent to
     # the server along with any events.
     def self.system_info
       @system_info ||= SystemInfo.new
     end
 
+    # Manages runtime metrics extracted from your app, which are uploaded to the
+    # Aikido servers if configured to do so.
+    def self.collector
+      @collector ||= Collector.new
+    end
+
     # Gets the current context object that holds all information about the
     # current request.
     #
@@ -47,24 +61,13 @@ module Aikido
       Thread.current[:_aikido_current_context_] = context
     end
 
-    # Track statistics about the result of a Sink's scan, and report it as an
-    # Attack if one is detected.
-    #
-    # @param scan [Aikido::Zen::Scan]
-    # @return [void]
-    # @raise [Aikido::Zen::UnderAttackError] if the scan detected an Attack
-    #   and blocking_mode is enabled.
-    def self.track_scan(scan)
-      agent.stats.add_scan(scan)
-      agent.handle_attack(scan.attack) if scan.attack?
-    end
-
     # Track statistics about an HTTP request the app is handling.
     #
-    # @param context [Aikido::Zen::Request]
+    # @param request [Aikido::Zen::Request]
     # @return [void]
     def self.track_request(request)
-      agent.stats.add_request(request)
+      autostart
+      collector.track_request(request)
     end
 
     # Tracks a network connection made to an external service.
@@ -72,7 +75,21 @@ module Aikido
     # @param connection [Aikido::Zen::OutboundConnection]
     # @return [void]
     def self.track_outbound(connection)
-      agent.stats.add_outbound(connection)
+      autostart
+      collector.track_outbound(connection)
+    end
+
+    # Track statistics about the result of a Sink's scan, and report it as
+    # an Attack if one is detected.
+    #
+    # @param scan [Aikido::Zen::Scan]
+    # @return [void]
+    # @raise [Aikido::Zen::UnderAttackError] if the scan detected an Attack
+    #   and blocking_mode is enabled.
+    def self.track_scan(scan)
+      autostart
+      collector.track_scan(scan)
+      agent.handle_attack(scan.attack) if scan.attack?
     end
 
     # Track the user making the current request.
@@ -80,43 +97,22 @@ module Aikido
     # @param (see Aikido::Zen.Actor)
     # @return [void]
     def self.track_user(user)
-      actor = Aikido::Zen::Actor(user)
+      return if config.disabled?
 
-      if actor
-        agent.stats.add_user(actor)
+      if (actor = Aikido::Zen::Actor(user))
+        autostart
+        collector.track_user(actor)
+        current_context.request.actor = actor if current_context
       else
-        id_attr, name_attr = config.user_attribute_mappings.values_at(:id, :name)
-        config.logger.warn(format(<<~LOG, obj: user, id: id_attr, name: name_attr))
-          Incompatible object sent to Aikido::Zen.track_user: %<obj>p
-
-          The object must satisfy one of the following:
+        config.logger.warn(format(<<~LOG, obj: user))
+          Incompatible object sent to track_user: %<obj>p
 
-          * Implement #to_aikido_actor
-          * Implement #to_model and have %<id>p and %<name>p attributes
-          * Be a Hash with :id (or "id") and, optionally, :name (or "name") keys
+          The object must either implement #to_aikido_actor, or be a Hash with
+          an :id (or "id") and, optionally, a :name (or "name") key.
         LOG
       end
     end
 
-    # Starts the background threads that keep the agent running.
-    #
-    # @return [void]
-    def self.initialize!
-      @agent ||= Agent.new
-      @agent.start!
-    end
-
-    # Stop any background threads.
-    def self.stop!
-      @agent&.stop!
-    end
-
-    # @return [Aikido::Zen::RuntimeSettings] the firewall configuration sourced
-    #   from your Aikido dashboard. This is periodically polled for updates.
-    def self.runtime_settings
-      @runtime_settings ||= RuntimeSettings.new
-    end
-
     # Load all sinks matching libraries loaded into memory. This method should
     # be called after all other dependencies have been loaded into memory (i.e.
     # at the end of the initialization process).
@@ -128,11 +124,20 @@ module Aikido
       require_relative "zen/sinks"
     end
 
-    private_class_method def self.agent
-      # We shouldn't start collecting data before we even initialize the agent,
-      # but might as well make sure we have a @agent going to report to.
-      @agent or initialize!
-      @agent
+    # @!visibility private
+    # Stop any background threads.
+    def self.stop!
+      agent&.stop!
+    end
+
+    # @!visibility private
+    # Starts the background agent if it has not been started yet.
+    def self.agent
+      @agent ||= Agent.start
+    end
+
+    class << self
+      alias_method :autostart, :agent
     end
   end
 end