diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index d98c08d..7a1b60a 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -5,6 +5,7 @@ module Rack::Attack autoload :Throttle, 'rack/attack/throttle' autoload :Whitelist, 'rack/attack/whitelist' autoload :Blacklist, 'rack/attack/blacklist' + autoload :Track, 'rack/attack/track' class << self @@ -22,9 +23,14 @@ module Rack::Attack self.throttles[name] = Throttle.new(name, options, block) end + def track(name, &block) + self.tracks[name] = Track.new(name, block) + end + def whitelists; @whitelists ||= {}; end def blacklists; @blacklists ||= {}; end def throttles; @throttles ||= {}; end + def tracks; @tracks ||= {}; end def new(app) @app = app @@ -52,6 +58,7 @@ module Rack::Attack elsif throttled?(req) throttled_response[env] else + tracked?(req) @app.call(env) end end @@ -74,6 +81,12 @@ module Rack::Attack end end + def tracked?(req) + tracks.each do |name, tracker| + tracker[req] + end + end + def instrument(req) notifier.instrument('rack.attack', req) if notifier end diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb new file mode 100644 index 0000000..c9d9152 --- /dev/null +++ b/lib/rack/attack/track.rb @@ -0,0 +1,10 @@ +module Rack + module Attack + class Track < Check + def initialize(name, block) + super + @type = :track + end + end + end +end diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb index adab8d5..ed437a9 100644 --- a/spec/rack_attack_spec.rb +++ b/spec/rack_attack_spec.rb @@ -1,28 +1,9 @@ require_relative 'spec_helper' describe 'Rack::Attack' do - include Rack::Test::Methods - - def app - Rack::Builder.new { - use Rack::Attack - run lambda {|env| [200, {}, ['Hello World']]} - }.to_app - end - - def self.allow_ok_requests - it "must allow ok requests" do - get '/', {}, 'REMOTE_ADDR' => '127.0.0.1' - last_response.status.must_equal 200 - last_response.body.must_equal 'Hello World' - end - end - - after { Rack::Attack.clear! } - allow_ok_requests - describe 'with a blacklist' do + describe 'blacklist' do before do @bad_ip = '1.2.3.4' Rack::Attack.blacklist("ip #{@bad_ip}") {|req| req.ip == @bad_ip } @@ -44,7 +25,7 @@ describe 'Rack::Attack' do allow_ok_requests end - describe "and with a whitelist" do + describe "and whitelist" do before do @good_ua = 'GoodUA' Rack::Attack.whitelist("good ua") {|req| req.user_agent == @good_ua } @@ -65,44 +46,4 @@ describe 'Rack::Attack' do end end - describe 'with a throttle' do - before do - @period = 60 # Use a long period; failures due to cache key rotation less likely - Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new - Rack::Attack.throttle('ip/sec', :limit => 1, :period => @period) { |req| req.ip } - end - - it('should have a throttle'){ Rack::Attack.throttles.key?('ip/sec') } - allow_ok_requests - - describe 'a single request' do - before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' } - it 'should set the counter for one request' do - key = "rack::attack:#{Time.now.to_i/@period}:ip/sec:1.2.3.4" - Rack::Attack.cache.store.read(key).must_equal 1 - end - - it 'should populate throttle data' do - data = { :count => 1, :limit => 1, :period => @period } - last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data - end - 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.matched'].must_equal 'ip/sec' - last_request.env['rack.attack.match_type'].must_equal :throttle - last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period}) - end - it 'should set a Retry-After header' do - last_response.headers['Retry-After'].must_equal @period.to_s - end - end - - end end diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb new file mode 100644 index 0000000..91bb922 --- /dev/null +++ b/spec/rack_attack_throttle_spec.rb @@ -0,0 +1,44 @@ + +require_relative 'spec_helper' +describe 'Rack::Attack.throttle' do + before do + @period = 60 # Use a long period; failures due to cache key rotation less likely + Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new + Rack::Attack.throttle('ip/sec', :limit => 1, :period => @period) { |req| req.ip } + end + + it('should have a throttle'){ Rack::Attack.throttles.key?('ip/sec') } + allow_ok_requests + + describe 'a single request' do + before { get '/', {}, 'REMOTE_ADDR' => '1.2.3.4' } + it 'should set the counter for one request' do + key = "rack::attack:#{Time.now.to_i/@period}:ip/sec:1.2.3.4" + Rack::Attack.cache.store.read(key).must_equal 1 + end + + it 'should populate throttle data' do + data = { :count => 1, :limit => 1, :period => @period } + last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data + end + 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.matched'].must_equal 'ip/sec' + last_request.env['rack.attack.match_type'].must_equal :throttle + last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period}) + end + it 'should set a Retry-After header' do + last_response.headers['Retry-After'].must_equal @period.to_s + end + end + +end + + diff --git a/spec/rack_attack_track_spec.rb b/spec/rack_attack_track_spec.rb new file mode 100644 index 0000000..3f0e56e --- /dev/null +++ b/spec/rack_attack_track_spec.rb @@ -0,0 +1,44 @@ +require_relative 'spec_helper' + +describe 'Rack::Attack.track' do + class Counter + def self.incr + @counter += 1 + end + + def self.reset + @counter = 0 + end + + def self.check + @counter + end + end + + before do + Rack::Attack.track("everything"){ |req| true } + end + allow_ok_requests + it "should tag the env" do + get '/' + last_request.env['rack.attack.matched'].must_equal 'everything' + last_request.env['rack.attack.match_type'].must_equal :track + end + + describe "with a notification subscriber and two tracks" do + before do + Counter.reset + # A second track + Rack::Attack.track("homepage"){ |req| req.path == "/"} + + ActiveSupport::Notifications.subscribe("rack.attack") do |*args| + Counter.incr + end + get "/" + end + + it "should notify twice" do + Counter.check.must_equal 2 + end + end +end diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index b36fc3c..706f5d5 100644 --- a/spec/spec_helper.rb +++ b/spec/spec_helper.rb @@ -7,3 +7,25 @@ require 'debugger' require 'active_support' require "rack/attack" + +class Minitest::Spec + + include Rack::Test::Methods + + after { Rack::Attack.clear! } + + def app + Rack::Builder.new { + use Rack::Attack + run lambda {|env| [200, {}, ['Hello World']]} + }.to_app + end + + def self.allow_ok_requests + it "must allow ok requests" do + get '/', {}, 'REMOTE_ADDR' => '127.0.0.1' + last_response.status.must_equal 200 + last_response.body.must_equal 'Hello World' + end + end +end