Merge pull request #153 from elhu/master

feature: support for ActiveSupport::MemCacheStore

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'

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
+          # 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
 

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

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

spec/integration/rack_attack_cache_spec.rb

@@ -17,12 +17,14 @@ describe Rack::Attack::Cache do
   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 @@ describe Rack::Attack::Cache do
           @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)