aikido-zen 0.1.0.alpha4-arm64-darwin → 0.1.1-arm64-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/internals.rb

@@ -31,6 +31,8 @@ module Aikido::Zen
       attach_function :detect_sql_injection_native, :detect_sql_injection,
         [:string, :string, :int], :int
     rescue LoadError, FFI::NotFoundError => err
+      # :nocov:
+
       # Emit an $stderr warning at startup.
       warn "Zen could not load its binary extension #{libzen_name}: #{err}"
 
@@ -38,6 +40,8 @@ module Aikido::Zen
         attempt = format("%p for SQL injection", query)
         raise InternalsError.new(attempt, "loading", Internals.libzen_name)
       end
+
+      # :nocov:
     else
       # Analyzes the SQL query to detect if the provided user input is being
       # passed as-is without escaping.

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/scanners/ssrf_scanner.rb

@@ -112,7 +112,8 @@ module Aikido::Zen
         is_port_relevant = input_uri.port != input_uri.default_port
         return false if is_port_relevant && input_uri.port != conn_uri.port
 
-        conn_uri.hostname == input_uri.hostname
+        conn_uri.hostname == input_uri.hostname &&
+          conn_uri.port == input_uri.port
       end
 
       def private_ip?(hostname)
@@ -128,8 +129,11 @@ module Aikido::Zen
       # * The input itself, if it already looks like a URI.
       # * The input prefixed with http://
       # * The input prefixed with https://
+      # * The input prefixed with the scheme of the request's URI, to consider
+      #   things like an FTP request (to "ftp://localhost") with a plain host
+      #   as a user-input ("localhost").
       #
-      # @return [Set<URI>]
+      # @return [Array<URI>] a list of unique URIs based on the above criteria.
       def uris_from_input
         input = @input.to_s
 
@@ -138,10 +142,12 @@ module Aikido::Zen
         # valid hostname. We should do the same for the input.
         input = format("[%s]", input) if unescaped_ipv6?(input)
 
-        [input, "http://#{input}", "https://#{input}"]
-          .map { |candidate| as_uri(candidate) }
-          .compact
-          .uniq
+        [
+          input,
+          "http://#{input}",
+          "https://#{input}",
+          "#{@request_uri.scheme}://#{input}"
+        ].map { |candidate| as_uri(candidate) }.compact.uniq
       end
 
       def as_uri(string)

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/http.rb

@@ -66,7 +66,7 @@ module Aikido::Zen
 
           response
         ensure
-          context["ssrf.request"] = prev_request
+          context["ssrf.request"] = prev_request if context
         end
       end
     end

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

@@ -7,6 +7,17 @@ module Aikido::Zen
     module PG
       SINK = Sinks.add("pg", scanners: [Scanners::SQLInjectionScanner])
 
+      # For some reason, the ActiveRecord pg adapter does not wrap exceptions in
+      # StatementInvalid, which leads to inconsistent handling. This guarantees
+      # that all Zen errors are wrapped in a StatementInvalid, so documentation
+      # can be consistent.
+      WRAP_EXCEPTIONS = if defined?(ActiveRecord::StatementInvalid)
+        <<~RUBY
+          rescue Aikido::Zen::SQLInjectionError
+            raise ActiveRecord::StatementInvalid
+        RUBY
+      end
+
       module Extensions
         %i[
           send_query exec sync_exec async_exec
@@ -16,12 +27,7 @@ module Aikido::Zen
             def #{method}(query, *)
               SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
               super
-            rescue Aikido::Zen::SQLInjectionError
-              # The pg adapter does not wrap exceptions in StatementInvalid, which
-              # leads to inconsistent handling. This guarantees that all Aikido
-              # errors are wrapped in a StatementInvalid, so documentation can be
-              # consistent.
-              raise ActiveRecord::StatementInvalid
+            #{WRAP_EXCEPTIONS}
             end
           RUBY
         end
@@ -33,12 +39,7 @@ module Aikido::Zen
             def #{method}(_, query, *)
               SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
               super
-            rescue Aikido::Zen::SQLInjectionError
-              # The pg adapter does not wrap exceptions in StatementInvalid, which
-              # leads to inconsistent handling. This guarantees that all Aikido
-              # errors are wrapped in a StatementInvalid, so documentation can be
-              # consistent.
-              raise ActiveRecord::StatementInvalid
+            #{WRAP_EXCEPTIONS}
             end
           RUBY
         end

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

@@ -66,7 +66,7 @@ module Aikido::Zen
             operation: "request"
           )
         ensure
-          context["ssrf.request"] = nil
+          context["ssrf.request"] = nil if context
         end
 
         true

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.1"
 
     # The version of libzen_internals that we build against.
-    LIBZEN_VERSION = "0.1.26"
+    LIBZEN_VERSION = "0.1.31"
   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