diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb index 747812d..2bedd1f 100644 --- a/lib/rack/attack/cache.rb +++ b/lib/rack/attack/cache.rb @@ -8,12 +8,12 @@ module Rack @prefix = 'rack::attack' end - def count(unprefixed_key, expires_in) - key = "#{prefix}:#{unprefixed_key}" - result = store.increment(key, 1, :expires_in => expires_in) + def count(unprefixed_key, period) + key = "#{prefix}:#{Time.now.to_i/period}:#{unprefixed_key}" + result = store.increment(key, 1) # NB: Some stores return nil when incrementing uninitialized values if result.nil? - store.write(key, 1, :expires_in => expires_in) + store.write(key, 1) end result || 1 end diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index 29c2002..014f1e4 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -21,11 +21,18 @@ module Rack key = "#{name}:#{discriminator}" count = cache.count(key, period) + data = { + :count => count, + :period => period, + :limit => limit + } + (req.env['rack.attack.throttle_data'] ||= {})[name] = data + (count > limit).tap do |throttled| if throttled req.env['rack.attack.matched'] = name req.env['rack.attack.match_type'] = :throttle - req.env['rack.attack.match_data'] = {:count => count, :period => period, :limit => limit} + req.env['rack.attack.match_data'] = data Rack::Attack.instrument(req) end end diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb index 0709ab6..91ccab8 100644 --- a/lib/rack/attack/version.rb +++ b/lib/rack/attack/version.rb @@ -1,5 +1,5 @@ module Rack module Attack - VERSION = '1.2.0' + VERSION = '1.3.0' end end diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb index 8a56f98..adab8d5 100644 --- a/spec/rack_attack_spec.rb +++ b/spec/rack_attack_spec.rb @@ -67,8 +67,9 @@ describe 'Rack::Attack' do describe 'with a throttle' 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 => 1, :period => 1) { |req| req.ip } + Rack::Attack.throttle('ip/sec', :limit => 1, :period => @period) { |req| req.ip } end it('should have a throttle'){ Rack::Attack.throttles.key?('ip/sec') } @@ -77,7 +78,13 @@ describe 'Rack::Attack' do describe 'a single request' do before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' } it 'should set the counter for one request' do - Rack::Attack.cache.store.read('rack::attack:ip/sec:1.2.3.4').must_equal 1 + 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 describe "with 2 requests" do @@ -90,10 +97,10 @@ describe 'Rack::Attack' do it 'should tag the env' do last_request.env['rack.attack.matched'].must_equal 'ip/sec' last_request.env['rack.attack.match_type'].must_equal :throttle - last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => 1}) + last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period}) end it 'should set a Retry-After header' do - last_response.headers['Retry-After'].must_equal '1' + last_response.headers['Retry-After'].must_equal @period.to_s end end