Add an reader for the epoch_time variable in the cache so that it can also be returned in the data from the throttle.

This is allows access to the same time that the cache uses for the count. This can be important for clients that want to provide rate limit information for well-behaved clients
This commit is contained in:
Kyle d'Oliveira 2018-02-15 12:50:42 -08:00
parent 39c04b311f
commit 9dbece5272
4 changed files with 15 additions and 11 deletions

View file

@ -254,13 +254,13 @@ Here's an example response that includes conventional `X-RateLimit-*` headers:
```ruby ```ruby
Rack::Attack.throttled_response = lambda do |env| Rack::Attack.throttled_response = lambda do |env|
now = Time.now
match_data = env['rack.attack.match_data'] match_data = env['rack.attack.match_data']
now = match_data[:epoch_time]
headers = { headers = {
'X-RateLimit-Limit' => match_data[:limit].to_s, 'X-RateLimit-Limit' => match_data[:limit].to_s,
'X-RateLimit-Remaining' => '0', '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"]] [ 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: For responses that did not exceed a throttle limit, Rack::Attack annotates the env with match data:
```ruby ```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 ## Logging & Instrumentation

View file

@ -3,6 +3,7 @@ module Rack
class Cache class Cache
attr_accessor :prefix attr_accessor :prefix
attr_reader :last_epoch_time
def initialize def initialize
self.store = ::Rails.cache if defined?(::Rails.cache) self.store = ::Rails.cache if defined?(::Rails.cache)
@ -39,10 +40,10 @@ module Rack
private private
def key_and_expiry(unprefixed_key, period) 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 # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
expires_in = (period - (epoch_time % period) + 1).to_i expires_in = (period - (@last_epoch_time % period) + 1).to_i
["#{prefix}:#{(epoch_time / period).to_i}:#{unprefixed_key}", expires_in] ["#{prefix}:#{(@last_epoch_time / period).to_i}:#{unprefixed_key}", expires_in]
end end
def do_count(key, expires_in) def do_count(key, expires_in)

View file

@ -26,12 +26,15 @@ module Rack
current_limit = limit.respond_to?(:call) ? limit.call(req) : limit current_limit = limit.respond_to?(:call) ? limit.call(req) : limit
key = "#{name}:#{discriminator}" key = "#{name}:#{discriminator}"
count = cache.count(key, current_period) count = cache.count(key, current_period)
epoch_time = cache.last_epoch_time
data = { data = {
:count => count, :count => count,
:period => current_period, :period => current_period,
:limit => current_limit :limit => current_limit,
:epoch_time => epoch_time
} }
(req.env['rack.attack.throttle_data'] ||= {})[name] = data (req.env['rack.attack.throttle_data'] ||= {})[name] = data
(count > current_limit).tap do |throttled| (count > current_limit).tap do |throttled|

View file

@ -20,7 +20,7 @@ describe 'Rack::Attack.throttle' do
end end
it 'should populate throttle data' do 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 last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
end end
end end
@ -37,7 +37,7 @@ describe 'Rack::Attack.throttle' do
it 'should tag the env' do it 'should tag the env' do
last_request.env['rack.attack.matched'].must_equal 'ip/sec' 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_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') last_request.env['rack.attack.match_discriminator'].must_equal('1.2.3.4')
end end
@ -65,7 +65,7 @@ describe 'Rack::Attack.throttle with limit as proc' do
end end
it 'should populate throttle data' do 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 last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
end end
end end
@ -89,7 +89,7 @@ describe 'Rack::Attack.throttle with period as proc' do
end end
it 'should populate throttle data' do 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 last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
end end
end end