Merge pull request #30 from lunks/allow-limit-to-be-a-proc

Allow limit option to be a proc.

README.md

@@ -147,6 +147,13 @@ how the parameters work.
     Rack::Attack.throttle('logins/email', :limit => 6, :period => 60.seconds) do |req|
       req.params['email'] if req.path == '/login' && req.post?
     end
+
+    # You can also set a limit using a proc instead of a number. For
+    # instance, after Rack::Auth::Basic has authenticated the user:
+    limit_based_on_proc = proc {|req| req.env["REMOTE_USER"] == "god" ? 100 : 1}
+    Rack::Attack.throttle('req/ip', :limit => limit_based_on_proc, :period => 1.second) do |req|
+      req.ip
+    end
 ```
 
 ### Tracks

lib/rack/attack/throttle.rb

@@ -22,14 +22,15 @@ module Rack
 
         key = "#{name}:#{discriminator}"
         count = cache.count(key, period)
+        current_limit = limit.respond_to?(:call) ? limit.call(req) : limit
         data = {
           :count => count,
           :period => period,
-          :limit => limit
+          :limit => current_limit
         }
         (req.env['rack.attack.throttle_data'] ||= {})[name] = data
 
-        (count > limit).tap do |throttled|
+        (count > current_limit).tap do |throttled|
           if throttled
             req.env['rack.attack.matched']    = name
             req.env['rack.attack.match_type'] = :throttle
@@ -38,7 +39,6 @@ module Rack
           end
         end
       end
-
     end
   end
 end

spec/rack_attack_throttle_spec.rb

@@ -37,7 +37,27 @@ describe 'Rack::Attack.throttle' do
       last_response.headers['Retry-After'].must_equal @period.to_s
     end
   end
-
 end
 
+describe 'Rack::Attack.throttle with limit as proc' do
+  before do
+    @period = 60 # Use a long period; failures due to cache key rotation less likely
+    Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new
+    Rack::Attack.throttle('ip/sec', :limit => lambda {|env| 1}, :period => @period) { |req| req.ip }
+  end
 
+  allow_ok_requests
+
+  describe 'a single request' do
+    before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' }
+    it 'should set the counter for one request' do
+      key = "rack::attack:#{Time.now.to_i/@period}:ip/sec:1.2.3.4"
+      Rack::Attack.cache.store.read(key).must_equal 1
+    end
+
+    it 'should populate throttle data' do
+      data = { :count => 1, :limit => 1, :period => @period }
+      last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
+    end
+  end
+end