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

data/lib/aikido/zen/config.rb

@@ -8,16 +8,10 @@ require_relative "context"
 
 module Aikido::Zen
   class Config
-    # @api private
-    # @return [Boolean] whether Aikido should protect.
-    def protect?
-      !api_token.nil? || blocking_mode? || debugging?
-    end
-
     # @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.
+    #   via the AIKIDO_DISABLE environment variable.
     attr_accessor :disabled
     alias_method :disabled?, :disabled
 
@@ -46,9 +40,9 @@ module Aikido::Zen
     #   settings changes. Defaults to evey 60 seconds.
     attr_accessor :polling_interval
 
-    # @return [Integer] the amount in seconds to wait before sending an initial
-    #   heartbeat event when the server reports no stats have been sent yet.
-    attr_accessor :initial_heartbeat_delay
+    # @return [Array<Integer>] the delays in seconds to wait before sending
+    #   each initial heartbeat event.
+    attr_accessor :initial_heartbeat_delays
 
     # @return [#call] Callable that can be passed an Object and returns a String
     #   of JSON. Defaults to the standard library's JSON.dump method.
@@ -61,15 +55,18 @@ module Aikido::Zen
     # @return [Logger]
     attr_reader :logger
 
-    # @return [string] Path of the socket where the detached agent will listen.
+    # @return [String] Path of the socket where the detached agent will listen.
     # By default, is stored under the root application path with file name
     # `aikido-detached-agent.sock`
-    attr_reader :detached_agent_socket_path
+    attr_accessor :detached_agent_socket_path
 
     # @return [Boolean] is the agent in debugging mode?
     attr_accessor :debugging
     alias_method :debugging?, :debugging
 
+    # @return [String] environment specific HTTP header providing the client IP.
+    attr_accessor :client_ip_header
+
     # @return [Integer] maximum number of timing measurements to keep in memory
     #   before compressing them.
     attr_accessor :max_performance_samples
@@ -146,25 +143,56 @@ module Aikido::Zen
     #   the server returns a 429 response.
     attr_accessor :server_rate_limit_deadline
 
+    # @return [Boolean] whether Aikido Zen should scan for stored SSSRF attacks.
+    #   Defaults to true. Can be set through AIKIDO_FEATURE_STORED_SSRF
+    #   environment variable.
+    attr_accessor :stored_ssrf
+    alias_method :stored_ssrf?, :stored_ssrf
+
     # @return [Array<String>] when checking for stored SSRF attacks, we want to
     #   allow known hosts that should be able to resolve to the IMDS service.
     attr_accessor :imds_allowed_hosts
 
+    # @return [Boolean] whether Aikido Zen should harden methods where possible.
+    #   Defaults to true. Can be set through AIKIDO_HARDEN environment variable.
+    attr_accessor :harden
+    alias_method :harden?, :harden
+
+    # @return [Integer] how many suspicious requests are allowed before an
+    #   attack wave detected event is reported.
+    #   Defaults to 15 requests.
+    attr_accessor :attack_wave_threshold
+
+    # @return [Integer] the minimum time in milliseconds between requests for
+    #   requests to be part of an attack wave.
+    #   Defaults to 1 minute in milliseconds.
+    attr_accessor :attack_wave_min_time_between_requests
+
+    # @return [Integer] the minimum time in milliseconds between reporting
+    #   attack wave events.
+    #   Defaults to 20 minutes in milliseconds.
+    attr_accessor :attack_wave_min_time_between_events
+
+    # @return [Integer] the maximum number of entries in the LRU cache.
+    #   Defaults to 10,000 entries.
+    attr_accessor :attack_wave_max_cache_entries
+
     def initialize
-      self.disabled = read_boolean_from_env(ENV.fetch("AIKIDO_DISABLED", false))
+      self.disabled = read_boolean_from_env(ENV.fetch("AIKIDO_DISABLE", false)) || read_boolean_from_env(ENV.fetch("AIKIDO_DISABLED", false))
       self.blocking_mode = read_boolean_from_env(ENV.fetch("AIKIDO_BLOCK", false))
       self.api_timeouts = 10
       self.api_endpoint = ENV.fetch("AIKIDO_ENDPOINT", DEFAULT_AIKIDO_ENDPOINT)
       self.realtime_endpoint = ENV.fetch("AIKIDO_REALTIME_ENDPOINT", DEFAULT_RUNTIME_BASE_URL)
       self.api_token = ENV.fetch("AIKIDO_TOKEN", nil)
-      self.polling_interval = 60
-      self.initial_heartbeat_delay = 60
+      self.polling_interval = 60 # 1 min
+      self.initial_heartbeat_delays = [30, 60 * 2] # 30 sec, 2 min
       self.json_encoder = DEFAULT_JSON_ENCODER
       self.json_decoder = DEFAULT_JSON_DECODER
       self.debugging = read_boolean_from_env(ENV.fetch("AIKIDO_DEBUG", false))
       self.logger = Logger.new($stdout, progname: "aikido", level: debugging ? Logger::DEBUG : Logger::INFO)
-      self.max_performance_samples = 5000
       self.detached_agent_socket_path = ENV.fetch("AIKIDO_DETACHED_AGENT_SOCKET_PATH", DEFAULT_DETACHED_AGENT_SOCKET_PATH)
+      self.client_ip_header = ENV.fetch("AIKIDO_CLIENT_IP_HEADER", nil)
+      self.max_performance_samples = 5000
       self.max_compressed_stats = 100
       self.max_outbound_connections = 200
       self.max_users_tracked = 1000
@@ -172,14 +200,20 @@ module Aikido::Zen
       self.blocked_responder = DEFAULT_BLOCKED_RESPONDER
       self.rate_limited_responder = DEFAULT_RATE_LIMITED_RESPONDER
       self.rate_limiting_discriminator = DEFAULT_RATE_LIMITING_DISCRIMINATOR
-      self.server_rate_limit_deadline = 1800 # 30 min
-      self.client_rate_limit_period = 3600 # 1 hour
+      self.server_rate_limit_deadline = 30 * 60 # 30 min
+      self.client_rate_limit_period = 60 * 60 # 1 hour
       self.client_rate_limit_max_events = 100
       self.collect_api_schema = read_boolean_from_env(ENV.fetch("AIKIDO_FEATURE_COLLECT_API_SCHEMA", true))
       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.stored_ssrf = read_boolean_from_env(ENV.fetch("AIKIDO_FEATURE_STORED_SSRF", true))
       self.imds_allowed_hosts = ["metadata.google.internal", "metadata.goog"]
+      self.harden = read_boolean_from_env(ENV.fetch("AIKIDO_HARDEN", true))
+      self.attack_wave_threshold = 15
+      self.attack_wave_min_time_between_requests = 60 * 1000 # 1 min (ms)
+      self.attack_wave_min_time_between_events = 20 * 60 * 1000 # 20 min (ms)
+      self.attack_wave_max_cache_entries = 10_000
     end
 
     # Set the base URL for API requests.
@@ -222,9 +256,8 @@ module Aikido::Zen
       @api_timeouts.update(value)
     end
 
-    def detached_agent_socket_path=(path)
-      @detached_agent_socket_path = path
-      @detached_agent_socket_path = "drbunix:" + @detached_agent_socket_path unless @detached_agent_socket_path.start_with?("drbunix:")
+    def detached_agent_socket_uri
+      "drbunix:" + @detached_agent_socket_path
     end
 
     private

data/lib/aikido/zen/context/rack_request.rb

@@ -6,6 +6,9 @@ require_relative "../request/heuristic_router"
 module Aikido::Zen
   # @!visibility private
   Context::RACK_REQUEST_BUILDER = ->(env) do
+    # Normalize PATH_INFO so routes are correctly recognized in middleware.
+    env["PATH_INFO"] = Helpers.normalize_path(env["PATH_INFO"])
+
     delegate = Rack::Request.new(env)
     router = Aikido::Zen::Request::HeuristicRouter.new
     request = Aikido::Zen::Request.new(delegate, framework: "rack", router: router)

data/lib/aikido/zen/context/rails_request.rb

@@ -12,6 +12,9 @@ module Aikido::Zen
 
   # @!visibility private
   Context::RAILS_REQUEST_BUILDER = ->(env) do
+    # Normalize PATH_INFO so routes are correctly recognized in middleware.
+    env["PATH_INFO"] = Helpers.normalize_path(env["PATH_INFO"])
+
     # Duplicate the Rack environment to prevent unexpected modifications from
     # breaking Rails routing.
     delegate = ActionDispatch::Request.new(env.dup)

data/lib/aikido/zen/context.rb

@@ -81,8 +81,8 @@ module Aikido::Zen
     def protection_disabled?
       return false if request.nil?
 
-      !@settings.endpoints[request.route].protected? ||
-        @settings.skip_protection_for_ips.include?(request.ip)
+      !@settings.endpoints.match(request.route).all?(&:protected?) ||
+        @settings.allowed_ips.include?(request.ip)
     end
 
     # @!visibility private
@@ -92,18 +92,51 @@ module Aikido::Zen
 
     private
 
+    # @!visibility private
     def extract_payloads_from(data, source_type, prefix = nil)
       if data.respond_to?(:to_hash)
-        data.to_hash.flat_map { |name, val|
-          extract_payloads_from(val, source_type, [prefix, name].compact.join("."))
-        }
+        data.to_hash.flat_map do |key, value|
+          extract_payloads_from(value, source_type, [prefix, key].compact.join("."))
+        end
       elsif data.respond_to?(:to_ary)
-        data.to_ary.flat_map.with_index { |val, idx|
-          extract_payloads_from(val, source_type, [prefix, idx].compact.join("."))
-        }
+        array = data.to_ary
+        return array if array.empty?
+
+        payloads = array.flat_map.with_index do |value, index|
+          extract_payloads_from(value, source_type, [prefix, index].compact.join("."))
+        end
+
+        unless Aikido::Zen.config.harden?
+          # Special case for File.join given a possibly nested array of strings,
+          # as might occur when a query parameter is an array.
+          begin
+            string = File.join__internal_for_aikido_zen(*array)
+            if unsafe_path?(string)
+              payloads << Payload.new(string, source_type, [prefix, "__File.join__"].compact.join("."))
+            end
+          rescue
+            # Could not create special payload for File.join.
+          end
+        end
+
+        payloads
       else
-        Payload.new(data, source_type, prefix.to_s)
+        [Payload.new(data, source_type, prefix.to_s)]
+      end
+    end
+
+    def unsafe_path?(filepath)
+      normalized_filepath = Pathname.new(filepath).cleanpath.to_s.downcase
+
+      Scanners::PathTraversal::DANGEROUS_PATH_PARTS.each do |dangerous_path_part|
+        return true if normalized_filepath.include?(dangerous_path_part)
+      end
+
+      Scanners::PathTraversal::DANGEROUS_PATH_STARTS.each do |dangerous_path_start|
+        return true if normalized_filepath.start_with?(dangerous_path_start)
       end
+
+      false
     end
   end
 end

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

@@ -14,53 +14,39 @@ module Aikido::Zen::DetachedAgent
   # parent process. We want to have the freshest data.
   #
   # It's possible to use `extend Forwardable` here for one-line forward calls to the
-  # @detached_agent_front object. Unfortunately, the methods to be called are
+  # @front_object object. Unfortunately, the methods to be called are
   # created at runtime by `DRbObject`, which leads to an ugly warning about
   # private methods after the delegator is bound.
   class Agent
     attr_reader :worker
 
     def initialize(
+      config: Aikido::Zen.config,
+      worker: Aikido::Zen::Worker.new(config: config),
       heartbeat_interval: 10,
       polling_interval: 10,
-      config: Aikido::Zen.config,
-      collector: Aikido::Zen.collector,
-      worker: Aikido::Zen::Worker.new(config: config)
+      collector: Aikido::Zen.collector
     )
       @config = config
+      @worker = worker
       @heartbeat_interval = heartbeat_interval
       @polling_interval = polling_interval
-      @worker = worker
+
       @collector = collector
-      @detached_agent_front = DRbObject.new_with_uri(config.detached_agent_socket_path)
+
+      @front_object = DRbObject.new_with_uri(config.detached_agent_socket_uri)
+
       @has_forked = false
       schedule_tasks
     end
 
-    def send_heartbeat(at: Time.now.utc)
-      return unless @collector.stats.any?
-
-      heartbeat = @collector.flush(at: at)
-      @detached_agent_front.send_heartbeat_to_parent_process(heartbeat.as_json)
-    end
-
-    private def schedule_tasks
-      # For heartbeats is correct to send them from parent or child process. Otherwise, we'll lose
-      # stats made by the parent process.
-      @worker.every(@heartbeat_interval, run_now: false) { send_heartbeat }
-
-      # Runtime_settings fetch must happens only in the child processes, otherwise, due to
-      # we are updating the global runtime_settings, we could have an infinite recursion.
-      if @has_forked
-        @worker.every(@polling_interval) do
-          Aikido::Zen.runtime_settings = @detached_agent_front.updated_settings
-          @config.logger.debug "Updated runtime settings after polling from child process #{Process.pid}"
-        end
-      end
+    def send_collector_events
+      events_data = @collector.flush_events.map(&:as_json)
+      @front_object.send_collector_events(events_data)
     end
 
     def calculate_rate_limits(request)
-      @detached_agent_front.calculate_rate_limits(request.route, request.ip, request.actor.to_json)
+      @front_object.calculate_rate_limits(request.route.as_json, request.ip, request.actor.as_json)
     end
 
     # Every time a fork occurs (a new child process is created), we need to start
@@ -74,5 +60,20 @@ module Aikido::Zen::DetachedAgent
       @worker.restart
       schedule_tasks
     end
+
+    private
+
+    def schedule_tasks
+      @worker.every(@heartbeat_interval, run_now: false) { send_collector_events }
+
+      # Runtime_settings fetch must happens only in the child processes, otherwise, due to
+      # we are updating the global runtime_settings, we could have an infinite recursion.
+      if @has_forked
+        @worker.every(@polling_interval) do
+          Aikido::Zen.runtime_settings = @front_object.updated_settings
+          @config.logger.debug "Updated runtime settings after polling from child process #{Process.pid}"
+        end
+      end
+    end
   end
 end

data/lib/aikido/zen/detached_agent/front_object.rb

@@ -7,20 +7,23 @@ module Aikido::Zen::DetachedAgent
   class FrontObject
     def initialize(
       config: Aikido::Zen.config,
-      collector: Aikido::Zen.collector,
       runtime_settings: Aikido::Zen.runtime_settings,
+      collector: Aikido::Zen.collector,
       rate_limiter: Aikido::Zen::RateLimiter.new
     )
       @config = config
+      @runtime_settings = runtime_settings
       @collector = collector
       @rate_limiter = rate_limiter
-      @runtime_settings = runtime_settings
     end
 
     RequestKind = Struct.new(:route, :schema, :ip, :actor)
 
-    def send_heartbeat_to_parent_process(heartbeat)
-      @collector.push_heartbeat(heartbeat)
+    def send_collector_events(events_data)
+      events_data.each do |event_data|
+        event = Aikido::Zen::Collector::Event.from_json(event_data)
+        @collector.add_event(event)
+      end
     end
 
     # Method called by child processes to get an up-to-date version of the
@@ -29,8 +32,9 @@ module Aikido::Zen::DetachedAgent
       @runtime_settings
     end
 
-    def calculate_rate_limits(route, ip, actor_hash)
-      actor = Aikido::Zen::Actor(actor_hash) if actor_hash
+    def calculate_rate_limits(route_data, ip, actor_data)
+      actor = Aikido::Zen::Actor.from_json(actor_data) if actor_data
+      route = Aikido::Zen::Route.from_json(route_data)
       @rate_limiter.calculate_rate_limits(RequestKind.new(route, nil, ip, actor))
     end
   end

data/lib/aikido/zen/detached_agent/server.rb

@@ -1,41 +1,78 @@
 # frozen_string_literal: true
 
+require "fileutils"
+
 module Aikido::Zen::DetachedAgent
   class Server
+    # Initialize and start a detached agent server instance.
+    #
+    # @return [Aikido::Zen::DetachedAgent::Server]
+    def self.start(**opts)
+      new(**opts).tap(&:start!)
+    end
+
     def initialize(config: Aikido::Zen.config)
-      @detached_agent_front = FrontObject.new
-      @drb_server = DRb.start_service(config.detached_agent_socket_path, @detached_agent_front)
+      @started_at = nil
 
-      # We don't want to see drb logs unless in debug mode
-      @drb_server.verbose = config.logger.debug?
-    end
+      @config = config
 
-    def alive?
-      @drb_server.alive?
+      @socket_path = config.detached_agent_socket_path
+      @socket_uri = config.detached_agent_socket_uri
     end
 
-    def stop!
-      @drb_server.stop_service
-      DRb.stop_service
+    def started?
+      !!@started_at
     end
 
-    class << self
-      def start!
-        Aikido::Zen.config.logger.debug("Starting DRb Server...")
-        max_attempts = 10
-        @server = new
-
-        attempts = 0
-        until @server.alive?
-          Aikido::Zen.config.logger.info("DRb Server still not alive. #{max_attempts - attempts} attempts remaining")
-          sleep 0.1
-          attempts += 1
-          raise Aikido::Zen::DetachedAgentError.new("Impossible to start the dRB server (socket=#{Aikido::Zen.config.detached_agent_socket_path})") \
-            if attempts == max_attempts
-        end
-
-        @server
+    def start!
+      @config.logger.info("Starting DRb Server...")
+
+      # Try to ensure that the DRb service can start if the DRb service did
+      # not stop cleanly.
+      begin
+        # Check whether the Unix domain socket is in use by another process.
+        UNIXSocket.new(@socket_path).close
+      rescue Errno::ECONNREFUSED
+        @config.logger.debug("Removing residual Unix domain socket...")
+
+        # Remove the residual Unix domain socket.
+        FileUtils.rm_f(@socket_path)
+      rescue
+        # empty
+      end
+
+      @front = FrontObject.new
+
+      # If the Unix domain socket is in use by another process and/or the
+      # residual Unix domain socket could not be removed DRb will raise an
+      # appropriate error.
+      @drb_server = DRb.start_service(@socket_uri, @front)
+
+      # Only show DRb output in debug mode.
+      @drb_server.verbose = @config.logger.debug?
+
+      # Ensure that the DRb server is alive.
+      max_attempts = 10
+      attempts = 0
+      until @drb_server.alive?
+        @config.logger.info("DRb Server still not alive. #{max_attempts - attempts} attempts remaining")
+        sleep 0.1
+        attempts += 1
+        raise Aikido::Zen::DetachedAgentError.new("Impossible to start the dRB server (socket=#{Aikido::Zen.config.detached_agent_socket_path})") \
+          if attempts == max_attempts
       end
+
+      @started_at = Time.now.utc
+
+      at_exit { stop! if started? }
+    end
+
+    def stop!
+      @config.logger.info("Stopping DRb Server...")
+      @started_at = nil
+
+      @drb_server.stop_service if @drb_server.alive?
+      DRb.stop_service
     end
   end
 end

data/lib/aikido/zen/event.rb

@@ -41,8 +41,10 @@ module Aikido::Zen
 
       def as_json
         super.update(
-          attack: @attack.as_json,
-          request: @attack.context.request.as_json
+          {
+            attack: @attack.as_json,
+            request: @attack.context&.request&.as_json
+          }.compact
         )
       end
     end
@@ -67,5 +69,48 @@ module Aikido::Zen
         )
       end
     end
+
+    class AttackWave < Event
+      # @param [Aikido::Zen::Context] a context
+      # @return [Aikido::Zen::Events::AttackWave] an attack wave event
+      def self.from_context(context)
+        request = Aikido::Zen::AttackWave::Request.new(
+          ip_address: context.request.client_ip,
+          user_agent: context.request.user_agent,
+          source: context.request.framework
+        )
+
+        attack = Aikido::Zen::AttackWave::Attack.new(
+          metadata: {}, # not used yet
+          user: context.request.actor
+        )
+
+        new(request: request, attack: attack)
+      end
+
+      # @return [Aikido::Zen::AttackWave::Request]
+      attr_reader :request
+
+      # @return [Aikido::Zen::AttackWave::Attack]
+      attr_reader :attack
+
+      # @param [Aikido::Zen::AttackWave::Request] the attack wave request
+      # @param [Aikido::Zen::AttackWave::Attack] the attack wave attack
+      # @param opts [Hash<Symbol, Object>] any other options to pass to
+      #   the superclass initializer.
+      # @return [Aikido::Zen::Events::AttackWave] an attack wave event
+      def initialize(request:, attack:, **opts)
+        super(type: "detected_attack_wave", **opts)
+        @request = request
+        @attack = attack
+      end
+
+      def as_json
+        super.update(
+          request: @request.as_json,
+          attack: @attack.as_json
+        )
+      end
+    end
   end
 end

data/lib/aikido/zen/helpers.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Aikido
+  module Zen
+    # @api private
+    module Helpers
+      # Normalizes a path by:
+      #
+      #   1. Collapsing consecutive forward slashes into a single forward slash.
+      #   2. Removing forward trailing slash, unless the normalized path is "/".
+      #
+      # @param path [String, nil] the path to normalize.
+      # @return [String, nil] the normalized path.
+      def self.normalize_path(path)
+        return path unless path
+
+        normalized_path = path.dup
+        normalized_path.squeeze!("/")
+        normalized_path.chomp!("/") unless normalized_path == "/"
+        normalized_path
+      end
+    end
+  end
+end

data/lib/aikido/zen/internals.rb

@@ -60,11 +60,11 @@ module Aikido::Zen
       # @!method self.detect_sql_injection_native(query, input, dialect)
       # @param (see .detect_sql_injection)
       # @returns [Integer] 0 if no injection detected, 1 if an injection was
-      #   detected, or 2 if there was an internal error.
+      #   detected, 2 if there was an internal error, or 3 if SQL tokenization failed.
       # @raise [Aikido::Zen::InternalsError] if there's a problem loading or
       #   calling libzen.
       attach_function :detect_sql_injection_native, :detect_sql_injection,
-        [:string, :string, :int], :int
+        [:pointer, :size_t, :pointer, :size_t, :int], :int
     rescue LoadError, FFI::NotFoundError => err # rubocop:disable Lint/ShadowedException
       # :nocov:
 
@@ -90,14 +90,34 @@ module Aikido::Zen
       # @raise [Aikido::Zen::InternalsError] if there's a problem loading or
       #   calling libzen.
       def self.detect_sql_injection(query, input, dialect)
-        case detect_sql_injection_native(query, input, dialect)
+        query_bytes = encode_safely(query)
+        input_bytes = encode_safely(input)
+
+        query_ptr = FFI::MemoryPointer.new(:uint8, query_bytes.bytesize)
+        input_ptr = FFI::MemoryPointer.new(:uint8, input_bytes.bytesize)
+
+        query_ptr.put_bytes(0, query_bytes)
+        input_ptr.put_bytes(0, input_bytes)
+
+        case detect_sql_injection_native(query_ptr, query_bytes.bytesize, input_ptr, input_bytes.bytesize, dialect)
         when 0 then false
         when 1 then true
         when 2
           attempt = format("%s query %p with input %p", dialect, query, input)
           raise InternalsError.new(attempt, "calling detect_sql_injection in", libzen_name)
+        when 3
+          # SQL tokenization failed - return false (no injection detected)
+          false
         end
       end
     end
+
+    class << self
+      private
+
+      def encode_safely(string)
+        string.encode("UTF-8", invalid: :replace, undef: :replace)
+      end
+    end
   end
 end

data/lib/aikido/zen/middleware/{check_allowed_addresses.rb → allowed_address_checker.rb}

@@ -3,7 +3,7 @@
 module Aikido::Zen
   module Middleware
     # Middleware that rejects requests from IPs blocked in the Aikido dashboard.
-    class CheckAllowedAddresses
+    class AllowedAddressChecker
       def initialize(app, config: Aikido::Zen.config, settings: Aikido::Zen.runtime_settings)
         @app = app
         @config = config

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

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module Aikido
+  module Zen
+    module Middleware
+      class AttackWaveProtector
+        def initialize(app, zen: Aikido::Zen, settings: Aikido::Zen.runtime_settings)
+          @app = app
+          @zen = zen
+          @settings = settings
+        end
+
+        def call(env)
+          response = @app.call(env)
+
+          context = @zen.current_context
+          protect(context)
+
+          response
+        end
+
+        # @api private
+        # Visible for testing.
+        def attack_wave?(context)
+          request = context.request
+          return false if request.nil?
+
+          # Bypass attack wave protection for allowed IPs
+          return false if @settings.allowed_ips.include?(request.ip)
+
+          @zen.attack_wave_detector.attack_wave?(context)
+        end
+
+        # @api private
+        # Visible for testing.
+        def protect(context)
+          if attack_wave?(context)
+            attack_wave = Aikido::Zen::Events::AttackWave.from_context(context)
+            @zen.track_attack_wave(attack_wave)
+            @zen.agent.report(attack_wave)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/aikido/zen/middleware/{set_context.rb → context_setter.rb}

@@ -9,7 +9,7 @@ module Aikido::Zen
   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.
-    class SetContext
+    class ContextSetter
       def initialize(app)
         @app = app
       end