Add ability to use multiple instances of middleware

Author: fatkodima

lib/rack/attack.rb

@@ -14,6 +14,7 @@ class Error < StandardError; end
     class MisconfiguredStoreError < Error; end
     class MissingStoreError < Error; end
 
+    autoload :Configuration,        'rack/attack/configuration'
     autoload :Cache,                'rack/attack/cache'
     autoload :Check,                'rack/attack/check'
     autoload :Throttle,             'rack/attack/throttle'
@@ -31,82 +32,8 @@ class MissingStoreError < Error; end
     autoload :Allow2Ban,            'rack/attack/allow2ban'
 
     class << self
-      attr_accessor :enabled, :notifier, :blocklisted_response, :throttled_response,
-                    :anonymous_blocklists, :anonymous_safelists
-
-      def safelist(name = nil, &block)
-        safelist = Safelist.new(name, &block)
-
-        if name
-          safelists[name] = safelist
-        else
-          anonymous_safelists << safelist
-        end
-      end
-
-      def blocklist(name = nil, &block)
-        blocklist = Blocklist.new(name, &block)
-
-        if name
-          blocklists[name] = blocklist
-        else
-          anonymous_blocklists << blocklist
-        end
-      end
-
-      def blocklist_ip(ip_address)
-        anonymous_blocklists << Blocklist.new { |request| IPAddr.new(ip_address).include?(IPAddr.new(request.ip)) }
-      end
-
-      def safelist_ip(ip_address)
-        anonymous_safelists << Safelist.new { |request| IPAddr.new(ip_address).include?(IPAddr.new(request.ip)) }
-      end
-
-      def throttle(name, options, &block)
-        throttles[name] = Throttle.new(name, options, &block)
-      end
-
-      def track(name, options = {}, &block)
-        tracks[name] = Track.new(name, options, &block)
-      end
-
-      def safelists
-        @safelists ||= {}
-      end
-
-      def blocklists
-        @blocklists ||= {}
-      end
-
-      def throttles
-        @throttles ||= {}
-      end
-
-      def tracks
-        @tracks ||= {}
-      end
-
-      def safelisted?(request)
-        anonymous_safelists.any? { |safelist| safelist.matched_by?(request) } ||
-          safelists.any? { |_name, safelist| safelist.matched_by?(request) }
-      end
-
-      def blocklisted?(request)
-        anonymous_blocklists.any? { |blocklist| blocklist.matched_by?(request) } ||
-          blocklists.any? { |_name, blocklist| blocklist.matched_by?(request) }
-      end
-
-      def throttled?(request)
-        throttles.any? do |_name, throttle|
-          throttle.matched_by?(request)
-        end
-      end
-
-      def tracked?(request)
-        tracks.each_value do |track|
-          track.matched_by?(request)
-        end
-      end
+      attr_accessor :enabled, :notifier
+      attr_reader :config
 
       def instrument(request)
         if notifier
@@ -122,34 +49,31 @@ def cache
         @cache ||= Cache.new
       end
 
-      def clear_configuration
-        @safelists = {}
-        @blocklists = {}
-        @throttles = {}
-        @tracks = {}
-        self.anonymous_blocklists = []
-        self.anonymous_safelists = []
-      end
-
       def clear!
         warn "[DEPRECATION] Rack::Attack.clear! is deprecated. Please use Rack::Attack.clear_configuration instead"
-        clear_configuration
+        @config.clear_configuration
       end
+
+      extend Forwardable
+      def_delegators :@config, :safelist, :blocklist, :blocklist_ip, :safelist_ip, :throttle, :track,
+                     :blocklisted_response, :blocklisted_response=, :throttled_response, :throttled_response=,
+                     :clear_configuration, :safelists, :blocklists, :throttles, :tracks
     end
 
     # Set defaults
     @enabled = true
-    @anonymous_blocklists = []
-    @anonymous_safelists = []
     @notifier = ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
-    @blocklisted_response = lambda { |_env| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
-    @throttled_response   = lambda do |env|
-      retry_after = (env['rack.attack.match_data'] || {})[:period]
-      [429, { 'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s }, ["Retry later\n"]]
-    end
+    @config = Configuration.new
 
-    def initialize(app)
+    def initialize(app, &config_block)
       @app = app
+      @config =
+        if block_given?
+          config = Configuration.new
+          config.instance_exec(&config_block)
+        else
+          self.class.config
+        end
     end
 
     def call(env)
@@ -158,19 +82,16 @@ def call(env)
       env['PATH_INFO'] = PathNormalizer.normalize_path(env['PATH_INFO'])
       request = Rack::Attack::Request.new(env)
 
-      if safelisted?(request)
+      if @config.safelisted?(request)
         @app.call(env)
-      elsif blocklisted?(request)
-        self.class.blocklisted_response.call(env)
-      elsif throttled?(request)
-        self.class.throttled_response.call(env)
+      elsif @config.blocklisted?(request)
+        @config.blocklisted_response.call(env)
+      elsif @config.throttled?(request)
+        @config.throttled_response.call(env)
       else
-        tracked?(request)
+        @config.tracked?(request)
         @app.call(env)
       end
     end
-
-    extend Forwardable
-    def_delegators self, :safelisted?, :blocklisted?, :throttled?, :tracked?
   end
 end

lib/rack/attack/configuration.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module Rack
+  class Attack
+    class Configuration
+      attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists
+      attr_accessor :blocklisted_response, :throttled_response
+
+      def initialize
+        @safelists = {}
+        @blocklists = {}
+        @throttles = {}
+        @tracks = {}
+        @anonymous_blocklists = []
+        @anonymous_safelists = []
+
+        @blocklisted_response = lambda { |_env| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
+        @throttled_response   = lambda do |env|
+          retry_after = (env['rack.attack.match_data'] || {})[:period]
+          [429, { 'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s }, ["Retry later\n"]]
+        end
+      end
+
+      def safelist(name = nil, &block)
+        safelist = Safelist.new(name, &block)
+
+        if name
+          @safelists[name] = safelist
+        else
+          @anonymous_safelists << safelist
+        end
+      end
+
+      def blocklist(name = nil, &block)
+        blocklist = Blocklist.new(name, &block)
+
+        if name
+          @blocklists[name] = blocklist
+        else
+          @anonymous_blocklists << blocklist
+        end
+      end
+
+      def blocklist_ip(ip_address)
+        @anonymous_blocklists << Blocklist.new { |request| IPAddr.new(ip_address).include?(IPAddr.new(request.ip)) }
+      end
+
+      def safelist_ip(ip_address)
+        @anonymous_safelists << Safelist.new { |request| IPAddr.new(ip_address).include?(IPAddr.new(request.ip)) }
+      end
+
+      def throttle(name, options, &block)
+        @throttles[name] = Throttle.new(name, options, &block)
+      end
+
+      def track(name, options = {}, &block)
+        @tracks[name] = Track.new(name, options, &block)
+      end
+
+      def safelisted?(request)
+        @anonymous_safelists.any? { |safelist| safelist.matched_by?(request) } ||
+          @safelists.any? { |_name, safelist| safelist.matched_by?(request) }
+      end
+
+      def blocklisted?(request)
+        @anonymous_blocklists.any? { |blocklist| blocklist.matched_by?(request) } ||
+          @blocklists.any? { |_name, blocklist| blocklist.matched_by?(request) }
+      end
+
+      def throttled?(request)
+        @throttles.any? do |_name, throttle|
+          throttle.matched_by?(request)
+        end
+      end
+
+      def tracked?(request)
+        @tracks.each_value do |track|
+          track.matched_by?(request)
+        end
+      end
+
+      def clear_configuration
+        @safelists = {}
+        @blocklists = {}
+        @throttles = {}
+        @tracks = {}
+        @anonymous_blocklists = []
+        @anonymous_safelists = []
+      end
+    end
+  end
+end