Feature proposal: Request instead of Env in callbacks (#419)

feat: allow easy access to the request object in the callbacks

README.md

@@ -313,18 +313,18 @@ Note that `Rack::Attack.cache` is only used for throttling, allow2ban and fail2b
 Customize the response of blocklisted and throttled requests using an object that adheres to the [Rack app interface](http://www.rubydoc.info/github/rack/rack/file/SPEC).
 
 ```ruby
-Rack::Attack.blocklisted_response = lambda do |env|
+Rack::Attack.blocklisted_callback = lambda do |request|
   # Using 503 because it may make attacker think that they have successfully
   # DOSed the site. Rack::Attack returns 403 for blocklists by default
   [ 503, {}, ['Blocked']]
 end
 
-Rack::Attack.throttled_response = lambda do |env|
+Rack::Attack.throttled_callback = lambda do |request|
   # NB: you have access to the name and other data about the matched throttle
-  #  env['rack.attack.matched'],
-  #  env['rack.attack.match_type'],
-  #  env['rack.attack.match_data'],
-  #  env['rack.attack.match_discriminator']
+  #  request.env['rack.attack.matched'],
+  #  request.env['rack.attack.match_type'],
+  #  request.env['rack.attack.match_data'],
+  #  request.env['rack.attack.match_discriminator']
 
   # Using 503 because it may make attacker think that they have successfully
   # DOSed the site. Rack::Attack returns 429 for throttling by default
@@ -411,9 +411,10 @@ def call(env)
   if safelisted?(req)
     @app.call(env)
   elsif blocklisted?(req)
-    self.class.blocklisted_response.call(env)
+    self.class.blocklisted_callback.call(req)
   elsif throttled?(req)
     self.class.throttled_response.call(env)
+    self.class.throttled_callback.call(req)
   else
     tracked?(req)
     @app.call(env)

docs/development.md

@@ -8,6 +8,10 @@ Install dependencies by running
 
     $ bundle install
 
+Install test dependencies by running:
+
+    $ bundle exec appraisal install
+
 Then run the test suite by running
 
     $ bundle exec appraisal rake test

lib/rack/attack.rb

@@ -66,6 +66,10 @@ module Rack
         :safelist_ip,
         :throttle,
         :track,
+        :throttled_callback,
+        :throttled_callback=,
+        :blocklisted_callback,
+        :blocklisted_callback=,
         :blocklisted_response,
         :blocklisted_response=,
         :throttled_response,
@@ -105,9 +109,13 @@ module Rack
       if configuration.safelisted?(request)
         @app.call(env)
       elsif configuration.blocklisted?(request)
-        configuration.blocklisted_response.call(env)
+        # Deprecated: Keeping blocklisted_response for backwards compatibility
+        configuration.blocklisted_response ? configuration.blocklisted_response.call(env) :
+            configuration.blocklisted_callback.call(request)
       elsif configuration.throttled?(request)
-        configuration.throttled_response.call(env)
+        # Deprecated: Keeping throttled_response for backwards compatibility
+        configuration.throttled_response ? configuration.throttled_response.call(env) :
+            configuration.throttled_callback.call(request)
       else
         configuration.tracked?(request)
         @app.call(env)

lib/rack/attack/configuration.rb

@@ -5,11 +5,11 @@ require "ipaddr"
 module Rack
   class Attack
     class Configuration
-      DEFAULT_BLOCKLISTED_RESPONSE = lambda { |_env| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
+      DEFAULT_BLOCKLISTED_CALLBACK = lambda { |_req| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
 
-      DEFAULT_THROTTLED_RESPONSE = lambda do |env|
+      DEFAULT_THROTTLED_CALLBACK = lambda do |req|
         if Rack::Attack.configuration.throttled_response_retry_after_header
-          match_data = env['rack.attack.match_data']
+          match_data = req.env['rack.attack.match_data']
           now = match_data[:epoch_time]
           retry_after = match_data[:period] - (now % match_data[:period])
 
@@ -20,7 +20,23 @@ module Rack
       end
 
       attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists
-      attr_accessor :blocklisted_response, :throttled_response, :throttled_response_retry_after_header
+      attr_accessor :blocklisted_callback, :throttled_callback, :throttled_response_retry_after_header
+
+      attr_reader :blocklisted_response, :throttled_response # Keeping these for backwards compatibility
+
+      def blocklisted_response=(callback)
+        # TODO: uncomment in 7.0
+        # warn "[DEPRECATION] Rack::Attack.blocklisted_response is deprecated. "\
+        #   "Please use Rack::Attack.blocklisted_callback instead."
+        @blocklisted_response = callback
+      end
+
+      def throttled_response=(callback)
+        # TODO: uncomment in 7.0
+        # warn "[DEPRECATION] Rack::Attack.throttled_response is deprecated. "\
+        #   "Please use Rack::Attack.throttled_callback instead"
+        @throttled_response = callback
+      end
 
       def initialize
         set_defaults
@@ -99,8 +115,12 @@ module Rack
         @anonymous_safelists = []
         @throttled_response_retry_after_header = false
 
-        @blocklisted_response = DEFAULT_BLOCKLISTED_RESPONSE
-        @throttled_response = DEFAULT_THROTTLED_RESPONSE
+        @blocklisted_callback = DEFAULT_BLOCKLISTED_CALLBACK
+        @throttled_callback = DEFAULT_THROTTLED_CALLBACK
+
+        # Deprecated: Keeping these for backwards compatibility
+        @blocklisted_response = nil
+        @throttled_response = nil
       end
     end
   end

spec/acceptance/customizing_blocked_response_spec.rb

@@ -14,7 +14,7 @@ describe "Customizing block responses" do
 
     assert_equal 403, last_response.status
 
-    Rack::Attack.blocklisted_response = lambda do |_env|
+    Rack::Attack.blocklisted_callback = lambda do |_req|
       [503, {}, ["Blocked"]]
     end
 
@@ -28,9 +28,9 @@ describe "Customizing block responses" do
     matched = nil
     match_type = nil
 
-    Rack::Attack.blocklisted_response = lambda do |env|
-      matched = env['rack.attack.matched']
-      match_type = env['rack.attack.match_type']
+    Rack::Attack.blocklisted_callback = lambda do |req|
+      matched = req.env['rack.attack.matched']
+      match_type = req.env['rack.attack.match_type']
 
       [503, {}, ["Blocked"]]
     end
@@ -40,4 +40,19 @@ describe "Customizing block responses" do
     assert_equal "block 1.2.3.4", matched
     assert_equal :blocklist, match_type
   end
+
+  it "supports old style" do
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 403, last_response.status
+
+    Rack::Attack.blocklisted_response = lambda do |_env|
+      [503, {}, ["Blocked"]]
+    end
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 503, last_response.status
+    assert_equal "Blocked", last_response.body
+  end
 end

spec/acceptance/customizing_throttled_response_spec.rb

@@ -20,7 +20,7 @@ describe "Customizing throttled response" do
 
     assert_equal 429, last_response.status
 
-    Rack::Attack.throttled_response = lambda do |_env|
+    Rack::Attack.throttled_callback = lambda do |_req|
       [503, {}, ["Throttled"]]
     end
 
@@ -36,11 +36,11 @@ describe "Customizing throttled response" do
     match_data = nil
     match_discriminator = nil
 
-    Rack::Attack.throttled_response = lambda do |env|
-      matched = env['rack.attack.matched']
-      match_type = env['rack.attack.match_type']
-      match_data = env['rack.attack.match_data']
-      match_discriminator = env['rack.attack.match_discriminator']
+    Rack::Attack.throttled_callback = lambda do |req|
+      matched = req.env['rack.attack.matched']
+      match_type = req.env['rack.attack.match_type']
+      match_data = req.env['rack.attack.match_data']
+      match_discriminator = req.env['rack.attack.match_discriminator']
 
       [429, {}, ["Throttled"]]
     end
@@ -58,4 +58,23 @@ describe "Customizing throttled response" do
     get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
     assert_equal 3, match_data[:count]
   end
+
+  it "supports old style" do
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 200, last_response.status
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 429, last_response.status
+
+    Rack::Attack.throttled_response = lambda do |_req|
+      [503, {}, ["Throttled"]]
+    end
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 503, last_response.status
+    assert_equal "Throttled", last_response.body
+  end
 end

spec/rack_attack_spec.rb

@@ -64,15 +64,15 @@ describe 'Rack::Attack' do
       end
     end
 
-    describe '#blocklisted_response' do
+    describe '#blocklisted_callback' do
       it 'should exist' do
-        _(Rack::Attack.blocklisted_response).must_respond_to :call
+        _(Rack::Attack.blocklisted_callback).must_respond_to :call
       end
     end
 
-    describe '#throttled_response' do
+    describe '#throttled_callback' do
       it 'should exist' do
-        _(Rack::Attack.throttled_response).must_respond_to :call
+        _(Rack::Attack.throttled_callback).must_respond_to :call
       end
     end
   end