diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index 1e0ae9ac..672f2335 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -11,6 +11,7 @@ class Rack::Attack autoload :Track, 'rack/attack/track' autoload :StoreProxy, 'rack/attack/store_proxy' autoload :DalliProxy, 'rack/attack/store_proxy/dalli_proxy' + autoload :MemCacheProxy, 'rack/attack/store_proxy/mem_cache_proxy' autoload :RedisStoreProxy, 'rack/attack/store_proxy/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 5c476ffd..d88ad674 100644 --- a/lib/rack/attack/store_proxy.rb +++ b/lib/rack/attack/store_proxy.rb @@ -1,19 +1,25 @@ module Rack class Attack module StoreProxy - PROXIES = [DalliProxy, RedisStoreProxy] + PROXIES = [DalliProxy, MemCacheProxy, RedisStoreProxy] 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) + if (defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore)) || + (defined?(::ActiveSupport::Cache::MemCacheStore) && store.is_a?(::ActiveSupport::Cache::MemCacheStore)) + # 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) + # so use the raw Redis::Store instead. + # We also want to use the underlying Dalli client instead of ::ActiveSupport::Cache::MemCacheStore, + # and the MemCache client if using Rails 3.x + client = store.instance_variable_get(:@data) + if (defined?(::Redis::Store) && client.is_a?(Redis::Store)) || + (defined?(Dalli::Client) && client.is_a?(Dalli::Client)) || (defined?(MemCache) && client.is_a?(MemCache)) + store = store.instance_variable_get(:@data) + end end - klass = PROXIES.find { |proxy| proxy.handle?(store) } - klass ? klass.new(store) : store end diff --git a/lib/rack/attack/store_proxy/mem_cache_proxy.rb b/lib/rack/attack/store_proxy/mem_cache_proxy.rb new file mode 100644 index 00000000..098e0480 --- /dev/null +++ b/lib/rack/attack/store_proxy/mem_cache_proxy.rb @@ -0,0 +1,51 @@ +module Rack + class Attack + module StoreProxy + class MemCacheProxy < SimpleDelegator + def self.handle?(store) + defined?(::MemCache) && store.is_a?(::MemCache) + end + + def initialize(store) + super(store) + stub_with_if_missing + end + + def read(key) + # Second argument: reading raw value + get(key, true) + rescue MemCache::MemCacheError + end + + def write(key, value, options={}) + # Third argument: writing raw value + set(key, value, options.fetch(:expires_in, 0), true) + rescue MemCache::MemCacheError + end + + def increment(key, amount, options={}) + incr(key, amount) + rescue MemCache::MemCacheError + end + + def delete(key, options={}) + with do |client| + client.delete(key) + end + rescue MemCache::MemCacheError + end + + private + + def stub_with_if_missing + unless __getobj__.respond_to?(:with) + class << self + def with; yield __getobj__; end + end + end + end + + end + end + end +end diff --git a/rack-attack.gemspec b/rack-attack.gemspec index 057aa6d8..1fcdd42a 100644 --- a/rack-attack.gemspec +++ b/rack-attack.gemspec @@ -32,4 +32,5 @@ Gem::Specification.new do |s| s.add_development_dependency 'redis-activesupport' s.add_development_dependency 'dalli' s.add_development_dependency 'connection_pool' + s.add_development_dependency 'memcache-client' end diff --git a/spec/integration/rack_attack_cache_spec.rb b/spec/integration/rack_attack_cache_spec.rb index 06aa51a4..6eb27eff 100644 --- a/spec/integration/rack_attack_cache_spec.rb +++ b/spec/integration/rack_attack_cache_spec.rb @@ -17,12 +17,14 @@ def sleep_until_expired end require 'active_support/cache/dalli_store' + require 'active_support/cache/mem_cache_store' require 'active_support/cache/redis_store' require 'connection_pool' cache_stores = [ ActiveSupport::Cache::MemoryStore.new, ActiveSupport::Cache::DalliStore.new("127.0.0.1"), ActiveSupport::Cache::RedisStore.new("127.0.0.1"), + ActiveSupport::Cache::MemCacheStore.new("127.0.0.1"), Dalli::Client.new, ConnectionPool.new { Dalli::Client.new }, Redis::Store.new @@ -54,6 +56,7 @@ def sleep_until_expired @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)