From 1ebe1c3517363cf88f50b22764be3c011fc0c197 Mon Sep 17 00:00:00 2001
From: Paul Coates <paul.coates@livingsocial.com>
Date: Mon, 19 May 2014 11:11:01 -0700
Subject: [PATCH] Added limit and period options to track. Delegates [] to
 Throttle if they are present otherwise Check.

---
 lib/rack/attack.rb             |  4 ++--
 lib/rack/attack/check.rb       |  4 ++--
 lib/rack/attack/throttle.rb    |  5 +++--
 lib/rack/attack/track.rb       | 19 +++++++++++++++----
 spec/rack_attack_track_spec.rb | 14 ++++++++++++++
 5 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index 0ec982c..eb9043e 100644
--- a/lib/rack/attack.rb
+++ b/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
diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb
index caea75a..21451de 100644
--- a/lib/rack/attack/check.rb
+++ b/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)
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
index 2824014..6bf91c9 100644
--- a/lib/rack/attack/throttle.rb
+++ b/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
diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb
index 3e516d7..4bfc145 100644
--- a/lib/rack/attack/track.rb
+++ b/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
diff --git a/spec/rack_attack_track_spec.rb b/spec/rack_attack_track_spec.rb
index 3f0e56e..8603266 100644
--- a/spec/rack_attack_track_spec.rb
+++ b/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