mirror of
https://github.com/samsonjs/rack-attack.git
synced 2026-03-25 09:25:49 +00:00
Support RedisStore as cache store
Add tests for different cache stores
This commit is contained in:
parent
234a2cff6c
commit
cf508e1d18
4 changed files with 92 additions and 4 deletions
|
|
@ -2,17 +2,43 @@ module Rack
|
|||
module Attack
|
||||
class Cache
|
||||
|
||||
attr_accessor :store, :prefix
|
||||
attr_accessor :prefix
|
||||
|
||||
def initialize
|
||||
@store = ::Rails.cache if defined?(::Rails.cache)
|
||||
self.store = ::Rails.cache if defined?(::Rails.cache)
|
||||
@prefix = 'rack::attack'
|
||||
end
|
||||
|
||||
attr_reader :store
|
||||
def store=(store)
|
||||
# RedisStore#increment needs different behavior, so detect that
|
||||
# (method has an arity of 2; must call #expire seperately
|
||||
if defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore)
|
||||
# ActiveSupport::Cache::RedisStore doesn't expose any way to set an expiry,
|
||||
# so use the raw Redis::Store instead
|
||||
@store = store.instance_variable_get(:@data)
|
||||
else
|
||||
@redis_store = false
|
||||
@store = store
|
||||
end
|
||||
end
|
||||
|
||||
def count(unprefixed_key, period)
|
||||
epoch_time = Time.now.to_i
|
||||
expires_in = period - (epoch_time % period)
|
||||
key = "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
|
||||
result = store.increment(key, 1, :expires_in => expires_in)
|
||||
do_count(key, expires_in)
|
||||
end
|
||||
|
||||
private
|
||||
def do_count(key, expires_in)
|
||||
# Workaround Redis::Store's interface
|
||||
if defined?(::Redis::Store) && store.is_a?(::Redis::Store)
|
||||
result = store.incr(key)
|
||||
store.expire(key, expires_in)
|
||||
else
|
||||
result = store.increment(key, 1, :expires_in => expires_in)
|
||||
end
|
||||
# NB: Some stores return nil when incrementing uninitialized values
|
||||
if result.nil?
|
||||
store.write(key, 1, :expires_in => expires_in)
|
||||
|
|
|
|||
|
|
@ -27,5 +27,7 @@ Gem::Specification.new do |s|
|
|||
s.add_development_dependency 'rake'
|
||||
s.add_development_dependency 'activesupport', '>= 3.0.0'
|
||||
s.add_development_dependency 'debugger', '~> 1.1.3'
|
||||
s.add_development_dependency 'redis-activesupport'
|
||||
s.add_development_dependency 'dalli'
|
||||
end
|
||||
|
||||
|
|
|
|||
61
spec/rack_attack_cache_spec.rb
Normal file
61
spec/rack_attack_cache_spec.rb
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
require_relative 'spec_helper'
|
||||
|
||||
if ENV['TEST_INTEGRATION']
|
||||
describe Rack::Attack::Cache do
|
||||
def delete(key)
|
||||
if @cache.store.respond_to?(:delete)
|
||||
@cache.store.delete(key)
|
||||
else
|
||||
@cache.store.del(key)
|
||||
end
|
||||
end
|
||||
|
||||
require 'active_support/cache/dalli_store'
|
||||
require 'active_support/cache/redis_store'
|
||||
cache_stores = [
|
||||
ActiveSupport::Cache::MemoryStore.new,
|
||||
ActiveSupport::Cache::DalliStore.new("localhost"),
|
||||
ActiveSupport::Cache::RedisStore.new("localhost"),
|
||||
Redis::Store.new
|
||||
]
|
||||
|
||||
cache_stores.each do |store|
|
||||
describe "with #{store.class}" do
|
||||
|
||||
before {
|
||||
@cache ||= Rack::Attack::Cache.new
|
||||
@key = "rack::attack:cache-test-key"
|
||||
@expires_in = 1
|
||||
@cache.store = store
|
||||
delete(@key)
|
||||
}
|
||||
|
||||
after { delete(@key) }
|
||||
|
||||
describe "do_count once" do
|
||||
it "should be 1" do
|
||||
@cache.send(:do_count, @key, @expires_in).must_equal 1
|
||||
end
|
||||
end
|
||||
|
||||
describe "do_count twice" do
|
||||
it "must be 2" do
|
||||
@cache.send(:do_count, @key, @expires_in)
|
||||
@cache.send(:do_count, @key, @expires_in).must_equal 2
|
||||
end
|
||||
end
|
||||
describe "do_count after expires_in" do
|
||||
it "must be 1" do
|
||||
@cache.send(:do_count, @key, @expires_in)
|
||||
sleep @expires_in # sigh
|
||||
@cache.send(:do_count, @key, @expires_in).must_equal 1
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
end
|
||||
|
||||
end
|
||||
else
|
||||
puts 'Skipping cache store integration tests (set ENV["TEST_INTEGRATION"] to enable)'
|
||||
end
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
|
||||
require_relative 'spec_helper'
|
||||
describe 'Rack::Attack.throttle' do
|
||||
before do
|
||||
|
|
|
|||
Loading…
Reference in a new issue