diff --git a/README.md b/README.md index d6d53e1..5cdac58 100644 --- a/README.md +++ b/README.md @@ -181,10 +181,11 @@ Rack::Attack.throttle('logins/email', :limit => 6, :period => 60.seconds) do |re req.params['email'] if req.path == '/login' && req.post? end -# You can also set a limit using a proc instead of a number. For -# instance, after Rack::Auth::Basic has authenticated the user: -limit_based_on_proc = proc {|req| req.env["REMOTE_USER"] == "admin" ? 100 : 1} -Rack::Attack.throttle('req/ip', :limit => limit_based_on_proc, :period => 1.second) do |req| +# You can also set a limit and period using a proc. For instance, after +# Rack::Auth::Basic has authenticated the user: +limit_proc = proc {|req| req.env["REMOTE_USER"] == "admin" ? 100 : 1} +period_proc = proc {|req| req.env["REMOTE_USER"] == "admin" ? 1.second : 1.minute} +Rack::Attack.throttle('req/ip', :limit => limit_proc, :period => period_proc) do |req| req.ip end ``` diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index 6bf91c9..5af953c 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -9,7 +9,7 @@ module Rack raise ArgumentError.new("Must pass #{opt.inspect} option") unless options[opt] end @limit = options[:limit] - @period = options[:period].to_i + @period = options[:period].respond_to?(:call) ? options[:period] : options[:period].to_i @type = options.fetch(:type, :throttle) end @@ -21,12 +21,14 @@ module Rack discriminator = block[req] return false unless discriminator - key = "#{name}:#{discriminator}" - count = cache.count(key, period) - current_limit = limit.respond_to?(:call) ? limit.call(req) : limit + current_period = period.respond_to?(:call) ? period.call(req) : period + current_limit = limit.respond_to?(:call) ? limit.call(req) : limit + key = "#{name}:#{discriminator}" + count = cache.count(key, current_period) + data = { :count => count, - :period => period, + :period => current_period, :limit => current_limit } (req.env['rack.attack.throttle_data'] ||= {})[name] = data diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb index 2acf55f..0c65872 100644 --- a/spec/rack_attack_throttle_spec.rb +++ b/spec/rack_attack_throttle_spec.rb @@ -44,7 +44,30 @@ describe 'Rack::Attack.throttle with limit as proc' do before do @period = 60 # Use a long period; failures due to cache key rotation less likely Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new - Rack::Attack.throttle('ip/sec', :limit => lambda {|req| 1}, :period => @period) { |req| req.ip } + Rack::Attack.throttle('ip/sec', :limit => lambda { |req| 1 }, :period => @period) { |req| req.ip } + end + + allow_ok_requests + + describe 'a single request' do + before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' } + it 'should set the counter for one request' do + key = "rack::attack:#{Time.now.to_i/@period}:ip/sec:1.2.3.4" + Rack::Attack.cache.store.read(key).must_equal 1 + end + + it 'should populate throttle data' do + data = { :count => 1, :limit => 1, :period => @period } + last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data + end + end +end + +describe 'Rack::Attack.throttle with period as proc' do + before do + @period = 60 # Use a long period; failures due to cache key rotation less likely + Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new + Rack::Attack.throttle('ip/sec', :limit => lambda { |req| 1 }, :period => lambda { |req| @period }) { |req| req.ip } end allow_ok_requests