From ab6d7b239d7b8239527419953e726687f39f378f Mon Sep 17 00:00:00 2001
From: Pedro Nascimento <pnascimento@gmail.com>
Date: Mon, 19 Aug 2013 16:22:25 -0300
Subject: [PATCH] Allow limit option to be a proc.

This allows you to do stuff like:
`req.env["USER"] == "god" ? 1000 : 1`
---
 README.md                         |  7 +++++++
 lib/rack/attack/throttle.rb       |  6 +++---
 spec/rack_attack_throttle_spec.rb | 22 +++++++++++++++++++++-
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 93792b0..3ccb950 100644
--- a/README.md
+++ b/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
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
index 29689e2..fa82266 100644
--- a/lib/rack/attack/throttle.rb
+++ b/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
diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb
index 13e2ec4..5b40ace 100644
--- a/spec/rack_attack_throttle_spec.rb
+++ b/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