From 6c259ea9be8fa639e8a119941e67188a61416897 Mon Sep 17 00:00:00 2001
From: madlep <madlep@madlep.com>
Date: Thu, 6 Jun 2013 15:50:35 +1000
Subject: [PATCH] delegate Redis custom logic to StoreProxy

this removes ugly `if redis blah` code from cache
---
 lib/rack/attack.rb             |  1 +
 lib/rack/attack/cache.rb       | 19 ++-----------
 lib/rack/attack/store_proxy.rb | 51 ++++++++++++++++++++++++++++++++++
 spec/rack_attack_cache_spec.rb |  1 +
 4 files changed, 56 insertions(+), 16 deletions(-)
 create mode 100644 lib/rack/attack/store_proxy.rb

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index d402775..b741478 100644
--- a/lib/rack/attack.rb
+++ b/lib/rack/attack.rb
@@ -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
 
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index 2388283..6b29ebb 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -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)
diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb
new file mode 100644
index 0000000..0aa9e62
--- /dev/null
+++ b/lib/rack/attack/store_proxy.rb
@@ -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
diff --git a/spec/rack_attack_cache_spec.rb b/spec/rack_attack_cache_spec.rb
index 26287bc..a3998ab 100644
--- a/spec/rack_attack_cache_spec.rb
+++ b/spec/rack_attack_cache_spec.rb
@@ -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 {