From c2f5ad2fa89c67cca5618e24db03c85ed7d3e992 Mon Sep 17 00:00:00 2001 From: Rozhnov Alexandr Date: Sat, 6 Jan 2024 19:23:59 +0100 Subject: [PATCH] Make it possible to block requests based on response. --- lib/rack/attack.rb | 6 +++- lib/rack/attack/configuration.rb | 14 ++++++-- lib/rack/attack/postrequest.rb | 22 ++++++++++++ spec/acceptance/postrequest_spec.rb | 52 +++++++++++++++++++++++++++++ spec/spec_helper.rb | 8 ++++- 5 files changed, 98 insertions(+), 4 deletions(-) create mode 100644 lib/rack/attack/postrequest.rb create mode 100644 spec/acceptance/postrequest_spec.rb diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index c9094b21..4c046a6e 100644 --- a/lib/rack/attack.rb +++ b/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 diff --git a/lib/rack/attack/configuration.rb b/lib/rack/attack/configuration.rb index a4bdc987..4fc68ebd 100644 --- a/lib/rack/attack/configuration.rb +++ b/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,6 +80,10 @@ 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) } @@ -87,7 +91,8 @@ def safelisted?(request) 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 diff --git a/lib/rack/attack/postrequest.rb b/lib/rack/attack/postrequest.rb new file mode 100644 index 00000000..221a0c51 --- /dev/null +++ b/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 diff --git a/spec/acceptance/postrequest_spec.rb b/spec/acceptance/postrequest_spec.rb new file mode 100644 index 00000000..0b4c029f --- /dev/null +++ b/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 diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index f529e6a1..58fc1514 100644 --- a/spec/spec_helper.rb +++ b/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