diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index d7c7252..c77245e 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -22,6 +22,7 @@ class Rack::Attack autoload :RedisProxy, 'rack/attack/store_proxy/redis_proxy' autoload :RedisStoreProxy, 'rack/attack/store_proxy/redis_store_proxy' autoload :RedisCacheStoreProxy, 'rack/attack/store_proxy/redis_cache_store_proxy' + autoload :ActiveSupportRedisStoreProxy, 'rack/attack/store_proxy/active_support_redis_store_proxy' autoload :Fail2Ban, 'rack/attack/fail2ban' autoload :Allow2Ban, 'rack/attack/allow2ban' diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb index bf9ec90..55c63e2 100644 --- a/lib/rack/attack/store_proxy.rb +++ b/lib/rack/attack/store_proxy.rb @@ -3,22 +3,18 @@ module Rack class Attack module StoreProxy - PROXIES = [DalliProxy, MemCacheStoreProxy, RedisStoreProxy, RedisProxy, RedisCacheStoreProxy].freeze + PROXIES = [ + DalliProxy, + MemCacheStoreProxy, + RedisStoreProxy, + RedisProxy, + RedisCacheStoreProxy, + ActiveSupportRedisStoreProxy + ].freeze def self.build(store) - client = unwrap_active_support_stores(store) - klass = PROXIES.find { |proxy| proxy.handle?(client) } - klass ? klass.new(client) : client - end - - def self.unwrap_active_support_stores(store) - # ActiveSupport::Cache::RedisStore doesn't expose any way to set an expiry, - # so use the raw Redis::Store instead. - if store.class.name == 'ActiveSupport::Cache::RedisStore' - store.instance_variable_get(:@data) - else - store - end + klass = PROXIES.find { |proxy| proxy.handle?(store) } + klass ? klass.new(store) : store end end end diff --git a/lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb b/lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb new file mode 100644 index 0000000..a18d081 --- /dev/null +++ b/lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb @@ -0,0 +1,37 @@ +# frozen_string_literal: true + +require 'delegate' + +module Rack + class Attack + module StoreProxy + class ActiveSupportRedisStoreProxy < SimpleDelegator + def self.handle?(store) + defined?(::Redis) && defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore) + end + + def increment(name, amount = 1, options = {}) + # #increment ignores options[:expires_in]. + # + # So in order to workaround this we use #write (which sets expiration) to initialize + # the counter. After that we continue using the original #increment. + if options[:expires_in] && !read(name) + write(name, amount, options) + + amount + else + super + end + end + + def read(name, options = {}) + super(name, options.merge!(raw: true)) + end + + def write(name, value, options = {}) + super(name, value, options.merge!(raw: true)) + end + end + end + end +end diff --git a/spec/acceptance/stores/active_support_redis_store_spec.rb b/spec/acceptance/stores/active_support_redis_store_spec.rb index 7e7ce54..75e4d68 100644 --- a/spec/acceptance/stores/active_support_redis_store_spec.rb +++ b/spec/acceptance/stores/active_support_redis_store_spec.rb @@ -12,7 +12,7 @@ if defined?(::ActiveSupport::Cache::RedisStore) end after do - Rack::Attack.cache.store.flushdb + Rack::Attack.cache.store.clear end it_works_for_cache_backed_features(fetch_from_store: ->(key) { Rack::Attack.cache.store.read(key) })