diff --git a/README.md b/README.md index b5381e7..8fe6d36 100644 --- a/README.md +++ b/README.md @@ -254,13 +254,13 @@ Here's an example response that includes conventional `X-RateLimit-*` headers: ```ruby Rack::Attack.throttled_response = lambda do |env| - now = Time.now match_data = env['rack.attack.match_data'] + now = match_data[:epoch_time] headers = { 'X-RateLimit-Limit' => match_data[:limit].to_s, 'X-RateLimit-Remaining' => '0', - 'X-RateLimit-Reset' => (now + (match_data[:period] - now.to_i % match_data[:period])).to_s + 'X-RateLimit-Reset' => (now + (match_data[:period] - now % match_data[:period])).to_s } [ 429, headers, ["Throttled\n"]] @@ -271,7 +271,7 @@ end For responses that did not exceed a throttle limit, Rack::Attack annotates the env with match data: ```ruby -request.env['rack.attack.throttle_data'][name] # => { :count => n, :period => p, :limit => l } +request.env['rack.attack.throttle_data'][name] # => { :count => n, :period => p, :limit => l, :epoch_time => t } ``` ## Logging & Instrumentation diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb index c2dd06c..73f0ca2 100644 --- a/lib/rack/attack/cache.rb +++ b/lib/rack/attack/cache.rb @@ -3,6 +3,7 @@ module Rack class Cache attr_accessor :prefix + attr_reader :last_epoch_time def initialize self.store = ::Rails.cache if defined?(::Rails.cache) @@ -39,10 +40,10 @@ module Rack private def key_and_expiry(unprefixed_key, period) - epoch_time = Time.now.to_i + @last_epoch_time = Time.now.to_i # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA - expires_in = (period - (epoch_time % period) + 1).to_i - ["#{prefix}:#{(epoch_time / period).to_i}:#{unprefixed_key}", expires_in] + expires_in = (period - (@last_epoch_time % period) + 1).to_i + ["#{prefix}:#{(@last_epoch_time / period).to_i}:#{unprefixed_key}", expires_in] end def do_count(key, expires_in) diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index 19a6519..4dc999d 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -26,12 +26,15 @@ module Rack current_limit = limit.respond_to?(:call) ? limit.call(req) : limit key = "#{name}:#{discriminator}" count = cache.count(key, current_period) + epoch_time = cache.last_epoch_time data = { :count => count, :period => current_period, - :limit => current_limit + :limit => current_limit, + :epoch_time => epoch_time } + (req.env['rack.attack.throttle_data'] ||= {})[name] = data (count > current_limit).tap do |throttled| diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb index 28ab26c..5befbe9 100644 --- a/spec/rack_attack_throttle_spec.rb +++ b/spec/rack_attack_throttle_spec.rb @@ -20,7 +20,7 @@ describe 'Rack::Attack.throttle' do end it 'should populate throttle data' do - data = { :count => 1, :limit => 1, :period => @period } + data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i } last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data end end @@ -37,7 +37,7 @@ describe 'Rack::Attack.throttle' 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 => @period}) + last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i}) last_request.env['rack.attack.match_discriminator'].must_equal('1.2.3.4') end @@ -65,7 +65,7 @@ describe 'Rack::Attack.throttle with limit as proc' do end it 'should populate throttle data' do - data = { :count => 1, :limit => 1, :period => @period } + data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i } last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data end end @@ -89,7 +89,7 @@ describe 'Rack::Attack.throttle with period as proc' do end it 'should populate throttle data' do - data = { :count => 1, :limit => 1, :period => @period } + data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i } last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data end end