Added limit and period options to track. Delegates [] to Throttle if they are present otherwise Check.

lib/rack/attack.rb

@@ -31,8 +31,8 @@ class Rack::Attack
       self.throttles[name] = Throttle.new(name, options, block)
     end
 
-    def track(name, &block)
-      self.tracks[name] = Track.new(name, block)
+    def track(name, options = {}, &block)
+      self.tracks[name] = Track.new(name, options, block)
     end
 
     def whitelists; @whitelists ||= {}; end

lib/rack/attack/check.rb

@@ -2,9 +2,9 @@ module Rack
   class Attack
     class Check
       attr_reader :name, :block, :type
-      def initialize(name, block)
+      def initialize(name, options = {}, block)
         @name, @block = name, block
-        @type = nil
+        @type = options.fetch(:type, nil)
       end
 
       def [](req)

lib/rack/attack/throttle.rb

@@ -2,7 +2,7 @@ module Rack
   class Attack
     class Throttle
       MANDATORY_OPTIONS = [:limit, :period]
-      attr_reader :name, :limit, :period, :block
+      attr_reader :name, :limit, :period, :block, :type
       def initialize(name, options, block)
         @name, @block = name, block
         MANDATORY_OPTIONS.each do |opt|
@@ -10,6 +10,7 @@ module Rack
         end
         @limit  = options[:limit]
         @period = options[:period].to_i
+        @type   = options.fetch(:type, :throttle)
       end
 
       def cache
@@ -34,7 +35,7 @@ module Rack
           if throttled
             req.env['rack.attack.matched']             = name
             req.env['rack.attack.match_discriminator'] = discriminator
-            req.env['rack.attack.match_type']          = :throttle
+            req.env['rack.attack.match_type']          = type
             req.env['rack.attack.match_data']          = data
             Rack::Attack.instrument(req)
           end

lib/rack/attack/track.rb

@@ -1,10 +1,21 @@
 module Rack
   class Attack
-    class Track < Check
-      def initialize(name, block)
-        super
-        @type = :track
+    class Track
+      extend Forwardable
+
+      attr_reader :checker
+
+      def initialize(name, options = {}, block)
+        options[:type] = :track
+
+        if options[:limit] && options[:period]
+          @checker = Throttle.new(name, options, block)
+        else
+          @checker = Check.new(name, options, block)
+        end
       end
+
+      def_delegator :@checker, :[], :[]
     end
   end
 end

spec/rack_attack_track_spec.rb

@@ -41,4 +41,18 @@ describe 'Rack::Attack.track' do
       Counter.check.must_equal 2
     end
   end
+
+  describe "without limit and period options" do
+    it "should delegate [] to check" do
+      tracker = Rack::Attack.track("homepage") { |req| req.path == "/"}
+      tracker.checker.class.must_equal Rack::Attack::Check
+    end
+  end
+
+  describe "with limit and period options" do
+    it "should delegate [] method to throttle" do
+      tracker = Rack::Attack.track("homepage", :limit => 10, :period => 10) { |req| req.path == "/"}
+      tracker.checker.class.must_equal Rack::Attack::Throttle
+    end
+  end
 end