Merge pull request #4 from kickstarter/track

Add Rack::Attack.track.

README.md

@@ -2,7 +2,7 @@
 *A DSL for blocking & throttling abusive clients*
 
 Rack::Attack is a rack middleware to protect your web app from bad clients.
-It allows *whitelisting*, *blacklisting*, and *throttling* based on arbitrary properties of the request.
+It allows *whitelisting*, *blacklisting*, *throttling*, and *tracking* based on arbitrary properties of the request.
 
 Throttle state is stored in a configurable cache (e.g. `Rails.cache`), presumably backed by memcached.
 
@@ -32,17 +32,31 @@ Note that `Rack::Attack.cache` is only used for throttling; not blacklisting & w
 
 ## How it works
 
-The Rack::Attack middleware compares each request against *whitelists*, *blacklists*, and *throttles* that you define. There are none by default.
+The Rack::Attack middleware compares each request against *whitelists*, *blacklists*, *throttles*, and *tracks* that you define. There are none by default.
 
- * If the request matches any whitelist, it is allowed. Blacklists and throttles are not checked.
+ * If the request matches any **whitelist**, it is allowed. Blacklists and throttles are not checked.
- * If the request matches any blacklist, it is blocked. Throttles are not checked.
+ * If the request matches any **blacklist**, it is blocked. Throttles are not checked.
- * If the request matches any throttle, a counter is incremented in the Rack::Attack.cache. If the throttle limit is exceeded, the request is blocked and further throttles are not checked.
+ * If the request matches any **throttle**, a counter is incremented in the Rack::Attack.cache. If the throttle limit is exceeded, the request is blocked and further throttles are not checked.
+ * If the request was not whitelisted, blacklisted, or throttled; all **tracks** are checked.
+
+## About Tracks
+
+`Rack::Attack.track` doesn't affect request processing. Tracks are an easy way to log and measure requests matching arbitrary attributes.
 
 ## Usage
 
-Define blacklists, throttles, and whitelists as blocks that return truthy values if matched, falsy otherwise.
+Define whitelists, blacklists, throttles, and tracks as blocks that return truthy values if matched, falsy otherwise.
 A [Rack::Request](http://rack.rubyforge.org/doc/classes/Rack/Request.html) object is passed to the block (named 'req' in the examples).
 
+### Whitelists
+
+    # Always allow requests from localhost
+    # (blacklist & throttles are skipped)
+    Rack::Attack.whitelist('allow from localhost') do |req|
+      # Requests are allowed if the return value is truthy
+      '127.0.0.1' == req.ip
+    end
+
 ### Blacklists
 
     # Block requests from 1.2.3.4
@@ -71,18 +85,25 @@ A [Rack::Request](http://rack.rubyforge.org/doc/classes/Rack/Request.html) objec
 
     # Throttle login attempts for a given email parameter to 6 reqs/minute
     Rack::Attack.throttle('logins/email', :limit => 6, :period => 60.seconds) do |req|
-      request.path == '/login' && req.post? && req.params['email']
+      req.path == '/login' && req.post? && req.params['email']
     end
 
-### Whitelists
+### Tracks
 
-    # Always allow requests from localhost
+    # Track requests from a special user agent
-    # (blacklist & throttles are skipped)
+    Rack::Attack.track("special_agent") do |req|
-    Rack::Attack.whitelist('allow from localhost') do |req|
+      req.user_agent == "SpecialAgent"
-      # Requests are allowed if the return value is truthy
-      '127.0.0.1' == req.ip
     end
 
+    # Track it using ActiveSupport::Notification
+    ActiveSupport::Notifications.subscribe("rack.attack") do |name, start, finish, request_id, req|
+      if req.env['rack.attack.matched'] == "special_agent" && req.env['rack.attack.match_type'] == :track
+        Rails.logger.info "special_agent: #{req.path}"
+        STATSD.increment("special_agent")
+      end
+    end
+
+
 ## Responses
 
 Customize the response of blacklisted and throttled requests using an object that adheres to the [Rack app interface](http://rack.rubyforge.org/doc/SPEC.html).

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_value do |tracker|
+        tracker[req]
+      end
+    end
+
     def instrument(req)
       notifier.instrument('rack.attack', req) if notifier
     end

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

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

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
+
+

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

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