Allow custom responses based on env

lib/rack/attack.rb

@@ -8,6 +8,7 @@ module Rack::Attack
   class << self
 
     attr_reader :cache, :notifier
+    attr_accessor :blacklisted_response, :throttled_response
 
     def whitelist(name, &block)
       (@whitelists ||= {})[name] = Whitelist.new(name, block)
@@ -28,6 +29,12 @@ module Rack::Attack
     def new(app)
       @cache ||= Cache.new
       @notifier = ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
+      @blacklisted_response = lambda {|env| [503, {}, ['Blocked']] }
+      @throttled_response   = lambda {|env|
+        retry_after = env['rack.attack.throttled'][:period] rescue nil
+        [503, {'Retry-After' => retry_after}, ['Retry later']]
+      }
+
       @app = app
       self
     end
@@ -41,9 +48,9 @@ module Rack::Attack
       end
 
       if blacklisted?(req)
-        blacklisted_response
+        blacklisted_response[env]
       elsif throttled?(req)
-        throttled_response
+        throttled_response[env]
       else
         @app.call(env)
       end
@@ -71,15 +78,7 @@ module Rack::Attack
       notifier.instrument('rack.attack', payload) if notifier
     end
 
-    def blacklisted_response
-      [503, {}, ['Blocked']]
-    end
-
-    def throttled_response
-      [503, {}, ['Throttled']]
-    end
-
-    def clear!
+   def clear!
       @whitelists, @blacklists, @throttles = {}, {}, {}
     end
 

lib/rack/attack/check.rb

@@ -9,7 +9,10 @@ module Rack
 
       def [](req)
         block[req].tap {|match|
-          Rack::Attack.instrument(:type => type, :name => name, :request => req) if match
+          if match
+            Rack::Attack.instrument(:type => type, :name => name, :request => req)
+            req.env["rack.attack.#{type}"] = name
+          end
         }
       end
 

lib/rack/attack/throttle.rb

@@ -21,9 +21,12 @@ module Rack
 
         key = "#{name}:#{discriminator}"
         count = cache.count(key, period)
-        throttled = count > limit
-        Rack::Attack.instrument(:type => :throttle, :name => name, :request => req, :count => count, :throttled => throttled)
-        throttled
+        (count > limit).tap do |throttled|
+          Rack::Attack.instrument(:type => :throttle, :name => name, :request => req, :count => count, :throttled => throttled)
+          if throttled
+            req.env['rack.attack.throttled'] = {:name => name, :count => count, :period => period, :limit => limit}
+          end
+        end
       end
 
     end

spec/rack_attack_spec.rb

@@ -30,12 +30,18 @@ describe 'Rack::Attack' do
 
     it('has a blacklist') { Rack::Attack.blacklists.key?("ip #{@bad_ip}") }
 
-    it "should blacklist bad requests" do
-      get '/', {}, 'REMOTE_ADDR' => @bad_ip
-      last_response.status.must_equal 503
-    end
+    describe "a bad request" do
+      before { get '/', {}, 'REMOTE_ADDR' => @bad_ip }
+      it "should return a blacklist response" do
+        get '/', {}, 'REMOTE_ADDR' => @bad_ip
+        last_response.status.must_equal 503
+      end
+      it "should tag the env" do
+        last_request.env['rack.attack.blacklist'].must_equal "ip #{@bad_ip}"
+      end
 
-    allow_ok_requests
+      allow_ok_requests
+    end
 
     describe "and with a whitelist" do
       before do
@@ -44,9 +50,15 @@ describe 'Rack::Attack' do
       end
 
       it('has a whitelist'){ Rack::Attack.whitelists.key?("good ua") }
-      it "should allow whitelists before blacklists" do
-        get '/', {}, 'REMOTE_ADDR' => @bad_ip, 'HTTP_USER_AGENT' => @good_ua
-        last_response.status.must_equal 200
+      describe "with a request match both whitelist & blacklist" do
+        before { get '/', {}, 'REMOTE_ADDR' => @bad_ip, 'HTTP_USER_AGENT' => @good_ua }
+        it "should allow whitelists before blacklists" do
+          get '/', {}, 'REMOTE_ADDR' => @bad_ip, 'HTTP_USER_AGENT' => @good_ua
+          last_response.status.must_equal 200
+        end
+        it "should tag the env" do
+          last_request.env['rack.attack.whitelist'].must_equal 'good ua'
+        end
       end
     end
   end
@@ -60,18 +72,26 @@ describe 'Rack::Attack' do
     it('should have a throttle'){ Rack::Attack.throttles.key?('ip/sec') }
     allow_ok_requests
 
-    it 'should set the counter for one request' do
-      get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
-      Rack::Attack.cache.store.read('rack::attack:ip/sec:1.2.3.4').must_equal 1
-    end
-
-    it 'should block 2 requests' do
-      2.times do
-        get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
+    describe 'a single request' do
+      before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' }
+      it 'should set the counter for one request' do
+        Rack::Attack.cache.store.read('rack::attack:ip/sec:1.2.3.4').must_equal 1
       end
-      last_response.status.must_equal 503
     end
+    describe "with 2 requests" do
+      before do
+        2.times { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' }
+      end
+      it 'should block the last request' do
+        last_response.status.must_equal 503
+      end
+      it 'should tag the env' do
+        last_request.env['rack.attack.throttled'].must_equal({:name => 'ip/sec', :count => 2, :limit => 1, :period => 1})
+      end
+      it 'should set a Retry-After header' do
+        last_response.headers['Retry-After'].must_equal 1
+      end
+    end
+
   end
-
-
 end