From cf508e1d18ae246c9c781ec9728ae49374f05cd5 Mon Sep 17 00:00:00 2001
From: Aaron Suggs <aaron@ktheory.com>
Date: Mon, 25 Feb 2013 12:09:59 -0500
Subject: [PATCH 1/2] Support RedisStore as cache store

Add tests for different cache stores
---
 lib/rack/attack/cache.rb          | 32 ++++++++++++++--
 rack-attack.gemspec               |  2 +
 spec/rack_attack_cache_spec.rb    | 61 +++++++++++++++++++++++++++++++
 spec/rack_attack_throttle_spec.rb |  1 -
 4 files changed, 92 insertions(+), 4 deletions(-)
 create mode 100644 spec/rack_attack_cache_spec.rb

diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index 06914c7..e5d2bfb 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -2,17 +2,43 @@ module Rack
   module Attack
     class Cache
 
-      attr_accessor :store, :prefix
+      attr_accessor :prefix
+
       def initialize
-        @store = ::Rails.cache if defined?(::Rails.cache)
+        self.store = ::Rails.cache if defined?(::Rails.cache)
         @prefix = 'rack::attack'
       end
 
+      attr_reader :store
+      def store=(store)
+        # RedisStore#increment needs different behavior, so detect that
+        # (method has an arity of 2; must call #expire seperately
+        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
+          @redis_store = false
+          @store = store
+        end
+      end
+
       def count(unprefixed_key, period)
         epoch_time = Time.now.to_i
         expires_in = period - (epoch_time % period)
         key = "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
-        result = store.increment(key, 1, :expires_in => expires_in)
+        do_count(key, expires_in)
+      end
+
+      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
         # NB: Some stores return nil when incrementing uninitialized values
         if result.nil?
           store.write(key, 1, :expires_in => expires_in)
diff --git a/rack-attack.gemspec b/rack-attack.gemspec
index eb2bf97..1cd56a0 100644
--- a/rack-attack.gemspec
+++ b/rack-attack.gemspec
@@ -27,5 +27,7 @@ Gem::Specification.new do |s|
   s.add_development_dependency 'rake'
   s.add_development_dependency 'activesupport', '>= 3.0.0'
   s.add_development_dependency 'debugger', '~> 1.1.3'
+  s.add_development_dependency 'redis-activesupport'
+  s.add_development_dependency 'dalli'
 end
 
diff --git a/spec/rack_attack_cache_spec.rb b/spec/rack_attack_cache_spec.rb
new file mode 100644
index 0000000..1419688
--- /dev/null
+++ b/spec/rack_attack_cache_spec.rb
@@ -0,0 +1,61 @@
+require_relative 'spec_helper'
+
+if ENV['TEST_INTEGRATION']
+  describe Rack::Attack::Cache do
+    def delete(key)
+      if @cache.store.respond_to?(:delete)
+        @cache.store.delete(key)
+      else
+        @cache.store.del(key)
+      end
+    end
+
+    require 'active_support/cache/dalli_store'
+    require 'active_support/cache/redis_store'
+    cache_stores = [
+      ActiveSupport::Cache::MemoryStore.new,
+      ActiveSupport::Cache::DalliStore.new("localhost"),
+      ActiveSupport::Cache::RedisStore.new("localhost"),
+      Redis::Store.new
+    ]
+
+    cache_stores.each do |store|
+      describe "with #{store.class}" do
+
+        before {
+          @cache ||= Rack::Attack::Cache.new
+          @key = "rack::attack:cache-test-key"
+          @expires_in = 1
+          @cache.store = store
+          delete(@key)
+        }
+
+        after { delete(@key) }
+
+        describe "do_count once" do
+          it "should be 1" do
+            @cache.send(:do_count, @key, @expires_in).must_equal 1
+          end
+        end
+
+        describe "do_count twice" do
+          it "must be 2" do
+            @cache.send(:do_count, @key, @expires_in)
+            @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)
+            sleep @expires_in # sigh
+            @cache.send(:do_count, @key, @expires_in).must_equal 1
+          end
+        end
+      end
+
+    end
+
+  end
+else
+  puts 'Skipping cache store integration tests (set ENV["TEST_INTEGRATION"] to enable)'
+end
diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb
index 91bb922..7938d76 100644
--- a/spec/rack_attack_throttle_spec.rb
+++ b/spec/rack_attack_throttle_spec.rb
@@ -1,4 +1,3 @@
-
 require_relative 'spec_helper'
 describe 'Rack::Attack.throttle' do
   before do

From 1c01e6097ce18f486d887ebbeb9f0c4b434cd8f5 Mon Sep 17 00:00:00 2001
From: Aaron Suggs <aaron@ktheory.com>
Date: Mon, 25 Feb 2013 22:39:11 -0500
Subject: [PATCH 2/2] bump to version 2.1.0

---
 lib/rack/attack/version.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb
index cc80cca..ba37a94 100644
--- a/lib/rack/attack/version.rb
+++ b/lib/rack/attack/version.rb
@@ -1,5 +1,5 @@
 module Rack
   module Attack
-    VERSION = '2.0.0'
+    VERSION = '2.1.0'
   end
 end