From ccdc1f993a93f0dfec5a7f85fa07d8f7a3052dad Mon Sep 17 00:00:00 2001 From: Aaron Suggs Date: Mon, 30 Jul 2012 15:44:22 -0400 Subject: [PATCH] Change instrumentation API for simpler notifications --- README.md | 10 +++++++--- lib/rack/attack.rb | 6 +++--- lib/rack/attack/check.rb | 5 +++-- lib/rack/attack/throttle.rb | 6 ++++-- lib/rack/attack/version.rb | 2 +- spec/rack_attack_spec.rb | 10 +++++++--- 6 files changed, 25 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 46b7f00..be690a9 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,11 @@ Note that `req` is a [Rack::Request](http://rack.rubyforge.org/doc/classes/Rack/ Customize the response of throttled requests using an object that adheres to the [Rack app interface](http://rack.rubyforge.org/doc/SPEC.html). Rack:Attack.throttled_response = lambda do |env| - env['rack.attack.throttled'] # name and other data about the matched throttle + # name and other data about the matched throttle + env['rack.attack.matched'] + env['rack.attack.match_type'] + env['rack.attack.match_data'] + [ 503, {}, ['Throttled']] end @@ -97,9 +101,9 @@ Similarly for blacklisted responses: Rack::Attack uses the [ActiveSupport::Notifications](http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html) API if available. -You can subscribe to 'rack.attack.{blacklist,throttle,whitelist}' events and log it, graph it, etc: +You can subscribe to 'rack.attack' events and log it, graph it, etc: - ActiveSupport::Notifications.subscribe('rack.attack.blacklist') do |name, start, finish, request_id, req| + ActiveSupport::Notifications.subscribe('rack.attack') do |name, start, finish, request_id, req| puts req.inspect end diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index 5b22baf..e51550b 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -34,7 +34,7 @@ module Rack::Attack @notifier ||= ActiveSupport::Notifications if defined?(ActiveSupport::Notifications) @blacklisted_response ||= lambda {|env| [503, {}, ['Blocked']] } @throttled_response ||= lambda {|env| - retry_after = env['rack.attack.matched'][:period] rescue nil + retry_after = env['rack.attack.match_data'][:period] rescue nil [503, {'Retry-After' => retry_after}, ['Retry later']] } @@ -75,8 +75,8 @@ module Rack::Attack end end - def instrument(type, payload) - notifier.instrument("rack.attack.#{type}", payload) if notifier + def instrument(req) + notifier.instrument('rack.attack', req) if notifier end def clear! diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb index 032b322..79f3cd9 100644 --- a/lib/rack/attack/check.rb +++ b/lib/rack/attack/check.rb @@ -10,8 +10,9 @@ module Rack def [](req) block[req].tap {|match| if match - req.env["rack.attack.matched"] = {type => name} - Rack::Attack.instrument(type, req) + req.env["rack.attack.matched"] = name + req.env["rack.attack.match_type"] = type + Rack::Attack.instrument(req) end } end diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index b8f61ba..29c2002 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -23,8 +23,10 @@ module Rack count = cache.count(key, period) (count > limit).tap do |throttled| if throttled - req.env['rack.attack.matched'] = {:throttle => name, :count => count, :period => period, :limit => limit} - Rack::Attack.instrument(:throttle, req) + req.env['rack.attack.matched'] = name + req.env['rack.attack.match_type'] = :throttle + req.env['rack.attack.match_data'] = {:count => count, :period => period, :limit => limit} + Rack::Attack.instrument(req) end end end diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb index bec5ffc..55e2720 100644 --- a/lib/rack/attack/version.rb +++ b/lib/rack/attack/version.rb @@ -1,5 +1,5 @@ module Rack module Attack - VERSION = '0.1.0' + VERSION = '0.2.0' end end diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb index 18cd583..9a68bfc 100644 --- a/spec/rack_attack_spec.rb +++ b/spec/rack_attack_spec.rb @@ -37,7 +37,8 @@ describe 'Rack::Attack' do last_response.status.must_equal 503 end it "should tag the env" do - last_request.env['rack.attack.matched'].must_equal({:blacklist => "ip #{@bad_ip}"}) + last_request.env['rack.attack.matched'].must_equal "ip #{@bad_ip}" + last_request.env['rack.attack.match_type'].must_equal :blacklist end allow_ok_requests @@ -57,7 +58,8 @@ describe 'Rack::Attack' do last_response.status.must_equal 200 end it "should tag the env" do - last_request.env['rack.attack.matched'].must_equal({:whitelist => 'good ua'}) + last_request.env['rack.attack.matched'].must_equal 'good ua' + last_request.env['rack.attack.match_type'].must_equal :whitelist end end end @@ -86,7 +88,9 @@ describe 'Rack::Attack' do last_response.status.must_equal 503 end it 'should tag the env' do - last_request.env['rack.attack.matched'].must_equal({:throttle => 'ip/sec', :count => 2, :limit => 1, :period => 1}) + 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 => 1}) end it 'should set a Retry-After header' do last_response.headers['Retry-After'].must_equal 1