From c90a0182ebc67d34756d74a172d202b081e62901 Mon Sep 17 00:00:00 2001 From: Aaron Suggs Date: Sat, 28 Jul 2012 19:51:24 -0400 Subject: [PATCH] Allow custom responses based on env --- lib/rack/attack.rb | 21 +++++++------- lib/rack/attack/check.rb | 5 +++- lib/rack/attack/throttle.rb | 9 ++++-- spec/rack_attack_spec.rb | 58 +++++++++++++++++++++++++------------ 4 files changed, 59 insertions(+), 34 deletions(-) diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index ccb69af..62f8d8b 100644 --- a/lib/rack/attack.rb +++ b/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 diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb index f7aaee3..8e97558 100644 --- a/lib/rack/attack/check.rb +++ b/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 diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index 020439d..068ff6e 100644 --- a/lib/rack/attack/throttle.rb +++ b/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 diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb index e8fc78d..ecaa7ee 100644 --- a/spec/rack_attack_spec.rb +++ b/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