mirror of
https://github.com/samsonjs/rack-attack.git
synced 2026-03-25 09:25:49 +00:00
delegate Redis custom logic to StoreProxy
this removes ugly `if redis blah` code from cache
This commit is contained in:
parent
22fc386bad
commit
6c259ea9be
4 changed files with 56 additions and 16 deletions
|
|
@ -6,6 +6,7 @@ module Rack::Attack
|
|||
autoload :Whitelist, 'rack/attack/whitelist'
|
||||
autoload :Blacklist, 'rack/attack/blacklist'
|
||||
autoload :Track, 'rack/attack/track'
|
||||
autoload :StoreProxy,'rack/attack/store_proxy'
|
||||
|
||||
class << self
|
||||
|
||||
|
|
|
|||
|
|
@ -11,15 +11,7 @@ module Rack
|
|||
|
||||
attr_reader :store
|
||||
def store=(store)
|
||||
# RedisStore#increment needs different behavior, so detect that
|
||||
# (method has an arity of 2; must call #expire separately
|
||||
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
|
||||
@store = store
|
||||
end
|
||||
@store = StoreProxy.build(store)
|
||||
end
|
||||
|
||||
def count(unprefixed_key, period)
|
||||
|
|
@ -39,13 +31,8 @@ module Rack
|
|||
|
||||
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
|
||||
result = store.increment(key, 1, :expires_in => expires_in)
|
||||
|
||||
# NB: Some stores return nil when incrementing uninitialized values
|
||||
if result.nil?
|
||||
store.write(key, 1, :expires_in => expires_in)
|
||||
|
|
|
|||
51
lib/rack/attack/store_proxy.rb
Normal file
51
lib/rack/attack/store_proxy.rb
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
require 'delegate'
|
||||
|
||||
module Rack
|
||||
module Attack
|
||||
class StoreProxy
|
||||
def self.build(store)
|
||||
# RedisStore#increment needs different behavior, so detect that
|
||||
# (method has an arity of 2; must call #expire separately
|
||||
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)
|
||||
end
|
||||
|
||||
if defined?(::Redis::Store) && store.is_a?(::Redis::Store)
|
||||
RedisStoreProxy.new(store)
|
||||
else
|
||||
store
|
||||
end
|
||||
end
|
||||
|
||||
class RedisStoreProxy < SimpleDelegator
|
||||
def initialize(store)
|
||||
super(store)
|
||||
end
|
||||
|
||||
def read(key)
|
||||
self.get(key)
|
||||
end
|
||||
|
||||
def write(key, value, options={})
|
||||
if (expires_in = options[:expires_in])
|
||||
self.setex(key, expires_in, value)
|
||||
else
|
||||
self.set(key, value)
|
||||
end
|
||||
end
|
||||
|
||||
def increment(key, amount, options={})
|
||||
count = nil
|
||||
self.pipelined do
|
||||
count = self.incrby(key, amount)
|
||||
self.expire(key, options[:expires_in]) if options[:expires_in]
|
||||
end
|
||||
count.value if count
|
||||
end
|
||||
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
@ -20,6 +20,7 @@ if ENV['TEST_INTEGRATION']
|
|||
]
|
||||
|
||||
cache_stores.each do |store|
|
||||
store = Rack::Attack::StoreProxy.build(store)
|
||||
describe "with #{store.class}" do
|
||||
|
||||
before {
|
||||
|
|
|
|||
Loading…
Reference in a new issue