Support for period ranges

lib/rack/attack/cache.rb

@@ -16,8 +16,17 @@ def store=(store)
 
       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}"
+        if period.is_a? Range
+          expires_in = period.end - epoch_time
+          # Manually create string representation to force cast both ends to int
+          # (instead of using Range#to_s)
+          period_key = "#{period.begin.to_i}..#{period.end.to_i}"
+        else
+          expires_in = period - (epoch_time % period)
+          period_key = (epoch_time/period).to_i
+        end
+
+        key = "#{prefix}:#{period_key}:#{unprefixed_key}"
         do_count(key, expires_in)
       end
 

lib/rack/attack/throttle.rb

@@ -9,7 +9,7 @@ def initialize(name, options, block)
           raise ArgumentError.new("Must pass #{opt.inspect} option") unless options[opt]
         end
         @limit  = options[:limit]
-        @period = options[:period].to_i
+        @period = options[:period]
         @type   = options.fetch(:type, :throttle)
       end
 
@@ -22,11 +22,12 @@ def [](req)
         return false unless discriminator
 
         key = "#{name}:#{discriminator}"
-        count = cache.count(key, period)
+        current_period = period.respond_to?(:call) ? period.call(req) : period
         current_limit = limit.respond_to?(:call) ? limit.call(req) : limit
+        count = cache.count(key, current_period)
         data = {
           :count => count,
-          :period => period,
+          :period => current_period,
           :limit => current_limit
         }
         (req.env['rack.attack.throttle_data'] ||= {})[name] = data

spec/rack_attack_cache_spec.rb

@@ -0,0 +1,27 @@
+require_relative 'spec_helper'
+
+describe Rack::Attack::Cache do
+  describe "#count" do
+    before do
+      @cache = Rack::Attack::Cache.new
+      @store = MiniTest::Mock.new
+      @cache.store = @store
+    end
+    describe "with a integer period" do
+      it "should set key with using current time over period" do
+        Time.stub(:now, 30) do
+          @store.expect(:increment, true, ["rack::attack:1:cache-test-key", 1, {expires_in: 20}])
+          @cache.count("cache-test-key", 25)
+        end
+      end
+    end
+    describe "with a range period" do
+      it "should set key with using current time over period" do
+        Time.stub(:now, 6) do
+          @store.expect(:increment, true, ["rack::attack:5..10:cache-test-key", 1, {expires_in: 4}])
+          @cache.count("cache-test-key", (5..10))
+        end
+      end
+    end
+  end
+end