Merge commit '8fcd6c855915c802adeeb1784c503fc74115f5a3' into 6-stable

README.md

@@ -342,6 +342,11 @@ end
 While Rack::Attack's primary focus is minimizing harm from abusive clients, it
 can also be used to return rate limit data that's helpful for well-behaved clients.
 
+If you want to return to user how many seconds to wait until he can start sending requests again, this can be done through enabling `Retry-After` header:
+```ruby
+Rack::Attack.throttled_response_retry_after_header = true
+```
+
 Here's an example response that includes conventional `RateLimit-*` headers:
 
 ```ruby

lib/rack/attack.rb

@@ -2,9 +2,10 @@
 
 require 'rack'
 require 'forwardable'
+require 'rack/attack/cache'
+require 'rack/attack/configuration'
 require 'rack/attack/path_normalizer'
 require 'rack/attack/request'
-require "ipaddr"
 
 require 'rack/attack/railtie' if defined?(::Rails)
 
@@ -14,7 +15,6 @@ module Rack
     class MisconfiguredStoreError < Error; end
     class MissingStoreError < Error; end
 
-    autoload :Cache,                'rack/attack/cache'
     autoload :Check,                'rack/attack/check'
     autoload :Throttle,             'rack/attack/throttle'
     autoload :Safelist,             'rack/attack/safelist'
@@ -31,82 +31,8 @@ module Rack
     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 :configuration
 
       def instrument(request)
         if notifier
@@ -122,34 +48,44 @@ module Rack
         @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
+        @configuration.clear_configuration
       end
+
+      extend Forwardable
+      def_delegators(
+        :@configuration,
+        :safelist,
+        :blocklist,
+        :blocklist_ip,
+        :safelist_ip,
+        :throttle,
+        :track,
+        :blocklisted_response,
+        :blocklisted_response=,
+        :throttled_response,
+        :throttled_response=,
+        :throttled_response_retry_after_header,
+        :throttled_response_retry_after_header=,
+        :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
+    @configuration = Configuration.new
+
+    attr_reader :configuration
 
     def initialize(app)
       @app = app
+      @configuration = self.class.configuration
     end
 
     def call(env)
@@ -159,19 +95,16 @@ module Rack
       env['PATH_INFO'] = PathNormalizer.normalize_path(env['PATH_INFO'])
       request = Rack::Attack::Request.new(env)
 
-      if safelisted?(request)
+      if configuration.safelisted?(request)
         @app.call(env)
-      elsif blocklisted?(request)
-        self.class.blocklisted_response.call(env)
-      elsif throttled?(request)
-        self.class.throttled_response.call(env)
+      elsif configuration.blocklisted?(request)
+        configuration.blocklisted_response.call(env)
+      elsif configuration.throttled?(request)
+        configuration.throttled_response.call(env)
       else
-        tracked?(request)
+        configuration.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,107 @@
+# frozen_string_literal: true
+
+require "ipaddr"
+
+module Rack
+  class Attack
+    class Configuration
+      DEFAULT_BLOCKLISTED_RESPONSE = lambda { |_env| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
+
+      DEFAULT_THROTTLED_RESPONSE = lambda do |env|
+        if Rack::Attack.configuration.throttled_response_retry_after_header
+          match_data = env['rack.attack.match_data']
+          now = match_data[:epoch_time]
+          retry_after = match_data[:period] - (now % match_data[:period])
+
+          [429, { 'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s }, ["Retry later\n"]]
+        else
+          [429, { 'Content-Type' => 'text/plain' }, ["Retry later\n"]]
+        end
+      end
+
+      attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists
+      attr_accessor :blocklisted_response, :throttled_response, :throttled_response_retry_after_header
+
+      def initialize
+        set_defaults
+      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
+        set_defaults
+      end
+
+      private
+
+      def set_defaults
+        @safelists = {}
+        @blocklists = {}
+        @throttles = {}
+        @tracks = {}
+        @anonymous_blocklists = []
+        @anonymous_safelists = []
+        @throttled_response_retry_after_header = false
+
+        @blocklisted_response = DEFAULT_BLOCKLISTED_RESPONSE
+        @throttled_response = DEFAULT_THROTTLED_RESPONSE
+      end
+    end
+  end
+end

lib/rack/attack/store_proxy/redis_cache_store_proxy.rb

@@ -15,33 +15,17 @@ module Rack
           #
           # So in order to workaround this we use RedisCacheStore#write (which sets expiration) to initialize
           # the counter. After that we continue using the original RedisCacheStore#increment.
-          rescuing do
-            if options[:expires_in] && !read(name)
-              write(name, amount, options)
+          if options[:expires_in] && !read(name)
+            write(name, amount, options)
 
-              amount
-            else
-              super
-            end
+            amount
+          else
+            super
           end
         end
 
-        def read(*_args)
-          rescuing { super }
-        end
-
         def write(name, value, options = {})
-          rescuing do
-            super(name, value, options.merge!(raw: true))
-          end
-        end
-
-        private
-
-        def rescuing
-          yield
-        rescue Redis::BaseError
-          nil
+          super(name, value, options.merge!(raw: true))
         end
       end
     end

lib/rack/attack/store_proxy/redis_proxy.rb

@@ -47,7 +47,7 @@ module Rack
 
         def rescuing
           yield
-        rescue Redis::BaseError
+        rescue Redis::BaseConnectionError
           nil
         end
       end

lib/rack/attack/throttle.rb

@@ -23,34 +23,50 @@ module Rack
 
       def matched_by?(request)
         discriminator = block.call(request)
+
         return false unless discriminator
 
-        current_period = period.respond_to?(:call) ? period.call(request) : period
-        current_limit  = limit.respond_to?(:call) ? limit.call(request) : limit
-        key            = "#{name}:#{discriminator}"
-        count          = cache.count(key, current_period)
-        epoch_time     = cache.last_epoch_time
+        current_period  = period_for(request)
+        current_limit   = limit_for(request)
+        count           = cache.count("#{name}:#{discriminator}", current_period)
 
         data = {
           discriminator: discriminator,
           count: count,
           period: current_period,
           limit: current_limit,
-          epoch_time: epoch_time
+          epoch_time: cache.last_epoch_time
         }
 
-        (request.env['rack.attack.throttle_data'] ||= {})[name] = data
-
         (count > current_limit).tap do |throttled|
+          annotate_request_with_throttle_data(request, data)
           if throttled
-            request.env['rack.attack.matched']             = name
-            request.env['rack.attack.match_discriminator'] = discriminator
-            request.env['rack.attack.match_type']          = type
-            request.env['rack.attack.match_data']          = data
+            annotate_request_with_matched_data(request, data)
             Rack::Attack.instrument(request)
           end
         end
       end
+
+      private
+
+      def period_for(request)
+        period.respond_to?(:call) ? period.call(request) : period
+      end
+
+      def limit_for(request)
+        limit.respond_to?(:call) ? limit.call(request) : limit
+      end
+
+      def annotate_request_with_throttle_data(request, data)
+        (request.env['rack.attack.throttle_data'] ||= {})[name] = data
+      end
+
+      def annotate_request_with_matched_data(request, data)
+        request.env['rack.attack.matched']             = name
+        request.env['rack.attack.match_discriminator'] = data[:discriminator]
+        request.env['rack.attack.match_type']          = type
+        request.env['rack.attack.match_data']          = data
+      end
     end
   end
 end

spec/acceptance/throttling_spec.rb

@@ -20,7 +20,7 @@ describe "#throttle" do
     get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
 
     assert_equal 429, last_response.status
-    assert_equal "60", last_response.headers["Retry-After"]
+    assert_nil last_response.headers["Retry-After"]
     assert_equal "Retry later\n", last_response.body
 
     get "/", {}, "REMOTE_ADDR" => "5.6.7.8"
@@ -34,6 +34,24 @@ describe "#throttle" do
     end
   end
 
+  it "returns correct Retry-After header if enabled" do
+    Rack::Attack.throttled_response_retry_after_header = true
+
+    Rack::Attack.throttle("by ip", limit: 1, period: 60) do |request|
+      request.ip
+    end
+
+    Timecop.freeze(Time.at(0)) do
+      get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+      assert_equal 200, last_response.status
+    end
+
+    Timecop.freeze(Time.at(25)) do
+      get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+      assert_equal "35", last_response.headers["Retry-After"]
+    end
+  end
+
   it "supports limit to be dynamic" do
     # Could be used to have different rate limits for authorized
     # vs general requests

spec/integration/offline_spec.rb

@@ -13,7 +13,11 @@ OfflineExamples = Minitest::SharedExamples.new do
   end
 
   it 'should count' do
-    @cache.send(:do_count, 'rack::attack::cache-test-key', 1)
+    @cache.count('cache-test-key', 1)
+  end
+
+  it 'should delete' do
+    @cache.delete('cache-test-key')
   end
 end
 
@@ -29,6 +33,18 @@ if defined?(::ActiveSupport::Cache::RedisStore)
   end
 end
 
+if defined?(Redis) && defined?(ActiveSupport::Cache::RedisCacheStore) && Redis::VERSION >= '4'
+  describe 'when Redis is offline' do
+    include OfflineExamples
+
+    before do
+      @cache = Rack::Attack::Cache.new
+      # Use presumably unused port for Redis client
+      @cache.store = ActiveSupport::Cache::RedisCacheStore.new(host: '127.0.0.1', port: 3333)
+    end
+  end
+end
+
 if defined?(::Dalli)
   describe 'when Memcached is offline' do
     include OfflineExamples
@@ -46,6 +62,23 @@ if defined?(::Dalli)
   end
 end
 
+if defined?(::Dalli) && defined?(::ActiveSupport::Cache::MemCacheStore)
+  describe 'when Memcached is offline' do
+    include OfflineExamples
+
+    before do
+      Dalli.logger.level = Logger::FATAL
+
+      @cache = Rack::Attack::Cache.new
+      @cache.store = ActiveSupport::Cache::MemCacheStore.new('127.0.0.1:22122')
+    end
+
+    after do
+      Dalli.logger.level = Logger::INFO
+    end
+  end
+end
+
 if defined?(Redis)
   describe 'when Redis is offline' do
     include OfflineExamples

spec/rack_attack_throttle_spec.rb

@@ -57,10 +57,6 @@ describe 'Rack::Attack.throttle' do
 
       _(last_request.env['rack.attack.match_discriminator']).must_equal('1.2.3.4')
     end
-
-    it 'should set a Retry-After header' do
-      _(last_response.headers['Retry-After']).must_equal @period.to_s
-    end
   end
 end
 

spec/spec_helper.rb

@@ -30,16 +30,11 @@ class MiniTest::Spec
 
   before do
     Rails.cache = nil
-    @_original_throttled_response = Rack::Attack.throttled_response
-    @_original_blocklisted_response = Rack::Attack.blocklisted_response
   end
 
   after do
     Rack::Attack.clear_configuration
     Rack::Attack.instance_variable_set(:@cache, nil)
-
-    Rack::Attack.throttled_response = @_original_throttled_response
-    Rack::Attack.blocklisted_response = @_original_blocklisted_response
   end
 
   def app