diff --git a/README.md b/README.md index f18f9f3..866cf35 100644 --- a/README.md +++ b/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 diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index add91c7..de5f745 100644 --- a/lib/rack/attack.rb +++ b/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 diff --git a/lib/rack/attack/configuration.rb b/lib/rack/attack/configuration.rb new file mode 100644 index 0000000..f2d0189 --- /dev/null +++ b/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 diff --git a/lib/rack/attack/store_proxy/redis_cache_store_proxy.rb b/lib/rack/attack/store_proxy/redis_cache_store_proxy.rb index f4081be..78807f8 100644 --- a/lib/rack/attack/store_proxy/redis_cache_store_proxy.rb +++ b/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 diff --git a/lib/rack/attack/store_proxy/redis_proxy.rb b/lib/rack/attack/store_proxy/redis_proxy.rb index 1bcb4b4..d4e6f3a 100644 --- a/lib/rack/attack/store_proxy/redis_proxy.rb +++ b/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 diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index c222325..3b80d9e 100644 --- a/lib/rack/attack/throttle.rb +++ b/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 diff --git a/spec/acceptance/throttling_spec.rb b/spec/acceptance/throttling_spec.rb index bf3a79b..0db89dd 100644 --- a/spec/acceptance/throttling_spec.rb +++ b/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 diff --git a/spec/integration/offline_spec.rb b/spec/integration/offline_spec.rb index f9bacab..d496063 100644 --- a/spec/integration/offline_spec.rb +++ b/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 diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb index 8346611..dcb1e41 100644 --- a/spec/rack_attack_throttle_spec.rb +++ b/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 diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index 0389bd0..10f856b 100644 --- a/spec/spec_helper.rb +++ b/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