Make it possible to block requests based on response.

lib/rack/attack.rb

@@ -31,6 +31,7 @@ class IncompatibleStoreError < Error; end
     autoload :Track,                'rack/attack/track'
     autoload :Fail2Ban,             'rack/attack/fail2ban'
     autoload :Allow2Ban,            'rack/attack/allow2ban'
+    autoload :Postrequest,          'rack/attack/postrequest'
 
     class << self
       attr_accessor :enabled, :notifier, :throttle_discriminator_normalizer
@@ -81,6 +82,7 @@ def reset!
         :clear_configuration,
         :safelists,
         :blocklists,
+        :postrequest,
         :throttles,
         :tracks
       )
@@ -126,7 +128,9 @@ def call(env)
         end
       else
         configuration.tracked?(request)
-        @app.call(env)
+        response = @app.call(env)
+        configuration.process_postrequests(request, response)
+        response
       end
     end
   end

lib/rack/attack/configuration.rb

@@ -19,7 +19,7 @@ class Configuration
         end
       end
 
-      attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists
+      attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists, :postrequests
       attr_accessor :blocklisted_responder, :throttled_responder, :throttled_response_retry_after_header
 
       attr_reader :blocklisted_response, :throttled_response # Keeping these for backwards compatibility
@@ -80,14 +80,19 @@ def track(name, options = {}, &block)
         @tracks[name] = Track.new(name, options, &block)
       end
 
+      def postrequest(name, &block)
+        @postrequests[name] = Postrequest.new(name, &block)
+      end
+
       def safelisted?(request)
         @anonymous_safelists.any? { |safelist| safelist.matched_by?(request) } ||
           @safelists.any? { |_name, safelist| safelist.matched_by?(request) }
       end
 
       def blocklisted?(request)
         @anonymous_blocklists.any? { |blocklist| blocklist.matched_by?(request) } ||
-          @blocklists.any? { |_name, blocklist| blocklist.matched_by?(request) }
+          @blocklists.any? { |_name, blocklist| blocklist.matched_by?(request) } ||
+          @postrequests.any? { |_name, postrequest| postrequest.matched_by?(request, nil) }
       end
 
       def throttled?(request)
@@ -102,6 +107,10 @@ def tracked?(request)
         end
       end
 
+      def process_postrequests(request, response)
+        @postrequests.each { |_name, postrequest| postrequest.matched_by?(request, response) }
+      end
+
       def clear_configuration
         set_defaults
       end
@@ -113,6 +122,7 @@ def set_defaults
         @blocklists = {}
         @throttles = {}
         @tracks = {}
+        @postrequests = {}
         @anonymous_blocklists = []
         @anonymous_safelists = []
         @throttled_response_retry_after_header = false

lib/rack/attack/postrequest.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Rack
+  class Attack
+    class Postrequest < Check
+      def initialize(name = nil, &block)
+        super
+        @type = :postrequest
+      end
+
+      def matched_by?(request, response)
+        block.call(request, response).tap do |match|
+          if match
+            request.env["rack.attack.matched"] = name
+            request.env["rack.attack.match_type"] = type
+            Rack::Attack.instrument(request)
+          end
+        end
+      end
+    end
+  end
+end

spec/acceptance/postrequest_spec.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require_relative "../spec_helper"
+require "timecop"
+
+describe "postrequest" do
+  before do
+    Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new
+
+    Rack::Attack.postrequest("fail2ban for 404") do |request, response|
+      Rack::Attack::Fail2Ban.filter(request.ip, maxretry: 2, findtime: 30, bantime: 60) do
+        if response.nil?
+          false
+        else
+          response[0] == 404
+        end
+      end
+    end
+  end
+
+  it "returns OK for many requests with 200 status" do
+    get "/"
+    assert_equal 200, last_response.status
+
+    get "/"
+    assert_equal 200, last_response.status
+  end
+
+
+  it "returns OK for few requests with 404 status" do
+    get "/not_found"
+    assert_equal 404, last_response.status
+
+    get "/not_found"
+    assert_equal 404, last_response.status
+  end
+
+  it "forbids all access after reaching maxretry limit" do
+    get "/not_found"
+    assert_equal 404, last_response.status
+
+    get "/not_found"
+    assert_equal 404, last_response.status
+
+    get "/not_found"
+    assert_equal 403, last_response.status
+
+    get "/"
+    assert_equal 403, last_response.status
+  end
+
+end

spec/spec_helper.rb

@@ -46,7 +46,13 @@ def app
       use Rack::Attack
       use Rack::Lint
 
-      run lambda { |_env| [200, {}, ['Hello World']] }
+      run lambda { |env|
+        if env['PATH_INFO'] == '/not_found'
+          [404, {}, ['Not Found']]
+        else
+          [200, {}, ['Hello World']]
+        end
+      }
     end.to_app
   end