Merge pull request #67 from chiliburger/track_only

Limit and period options for track

README.md

@@ -188,11 +188,16 @@ end
 ### Tracks
 
 ```ruby
-# Track requests from a special user agent
+# Track requests from a special user agent.
 Rack::Attack.track("special_agent") do |req|
   req.user_agent == "SpecialAgent"
 end
 
+# Supports optional limit and period, triggers the notification only when the limit is reached.
+Rack::Attack.track("special_agent", :limit 6, :period => 60.seconds) do |req|
+  req.user_agent == "SpecialAgent"
+end
+
 # Track it using ActiveSupport::Notification
 ActiveSupport::Notifications.subscribe("rack.attack") do |name, start, finish, request_id, req|
   if req.env['rack.attack.matched'] == "special_agent" && req.env['rack.attack.match_type'] == :track

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 :filter
+
+      def initialize(name, options = {}, block)
+        options[:type] = :track
+
+        if options[:limit] && options[:period]
+          @filter = Throttle.new(name, options, block)
+        else
+          @filter = Check.new(name, options, block)
+        end
       end
+
+      def_delegator :@filter, :[]
     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 assign the track filter to a Check instance" do
+      tracker = Rack::Attack.track("homepage") { |req| req.path == "/"}
+      tracker.filter.class.must_equal Rack::Attack::Check
+    end
+  end
+
+  describe "with limit and period options" do
+    it "should assign the track filter to a Throttle instance" do
+      tracker = Rack::Attack.track("homepage", :limit => 10, :period => 10) { |req| req.path == "/"}
+      tracker.filter.class.must_equal Rack::Attack::Throttle
+    end
+  end
 end