From 9284a08cc3377330b8e79eec31da5a57aa81f3c4 Mon Sep 17 00:00:00 2001
From: Aaron Suggs <aaron@ktheory.com>
Date: Thu, 26 Jul 2012 17:29:09 -0400
Subject: [PATCH] Whitelists support

---
 lib/rack/attack.rb         | 86 +++++++++++++++++++++++++++++++-------
 lib/rack/attack/cache.rb   | 21 ++++++++++
 lib/rack/attack/version.rb |  2 +-
 rack-attack.gemspec        |  1 +
 spec/rack_attack_spec.rb   | 47 ++++++++++++++++-----
 5 files changed, 130 insertions(+), 27 deletions(-)
 create mode 100644 lib/rack/attack/cache.rb

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index 39737d1..c72c62a 100644
--- a/lib/rack/attack.rb
+++ b/lib/rack/attack.rb
@@ -1,29 +1,83 @@
 require 'rack'
-module Rack
-  class Attack
-    class << self
+module Rack::Attack
+  require 'rack/attack/cache'
 
-      attr_reader :blocks, :throttles, :whitelists
+  class << self
 
-      def block(name, &block)
-        (@blocks ||= {})[name] = block
-      end
-
-      def throttle
-      end
-
-      def whitelist
-      end
+    attr_reader :cache, :notifier
 
+    def whitelist(name, &block)
+      (@whitelists ||= {})[name] = block
     end
 
-    def initialize(app)
+    def block(name, &block)
+      (@blocks ||= {})[name] = block
+    end
+
+    def throttle
+    end
+
+    def whitelists; @whitelists ||= {}; end
+    def blocks;     @blocks     ||= {}; end
+    def throttles;  @throttles  ||= {}; end
+
+    def new(app)
+      @cache = Cache.new
+      @notifier = ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
       @app = app
+      self
     end
 
+
     def call(env)
-      puts 'Rack attack!'
-      @app.call(env)
+      req = Rack::Request.new(env)
+
+      if whitelisted?(req)
+        return @app.call(env)
+      end
+
+      if blocked?(req)
+        blocked_response
+      elsif throttled?(req)
+      else
+        @app.call(env)
+      end
+    end
+
+    def whitelisted?(req)
+      whitelists.any? do |name, block|
+        block[req].tap{ |match|
+          instrument(:type => :whitelist, :name => name, :request => req) if match
+        }
+      end
+    end
+
+    def blocked?(req)
+      blocks.any? do |name, block|
+        block[req].tap { |match|
+          instrument(:type => :block, :name => name, :request => req) if match
+        }
+      end
+    end
+
+    def throttled?(req)
+      false
+    end
+
+    def instrument(payload)
+      notifier.instrument('rack.attack', payload) if notifier
+    end
+
+    def blocked_response
+      [503, {}, ['Blocked']]
+    end
+
+    def throttled_response
+      [503, {}, ['Throttled']]
+    end
+
+    def clear!
+      @whitelists, @blocks, @throttles = {}, {}, {}
     end
 
   end
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
new file mode 100644
index 0000000..d7d7c91
--- /dev/null
+++ b/lib/rack/attack/cache.rb
@@ -0,0 +1,21 @@
+module Rack
+  module Attack
+    class Cache
+
+      attr_accessor :store
+      def initialize
+        @store = ::Rails.cache if defined?(::Rails.cache)
+      end
+
+      def count(key, expires_in)
+        result = store.increment(1, :expires_in => expires_in)
+        # NB: Some stores return nil when incrementing uninitialized values
+        if result.nil?
+          store.write(key, 1, :expires_in => expires_in)
+        end
+        result || 1
+      end
+
+    end
+  end
+end
diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb
index d1560dd..4e63ea4 100644
--- a/lib/rack/attack/version.rb
+++ b/lib/rack/attack/version.rb
@@ -1,5 +1,5 @@
 module Rack
-  class Attack
+  module Attack
     VERSION = '0.0.1'
   end
 end
diff --git a/rack-attack.gemspec b/rack-attack.gemspec
index 9b6bcf0..49efd7d 100644
--- a/rack-attack.gemspec
+++ b/rack-attack.gemspec
@@ -24,6 +24,7 @@ Gem::Specification.new do |s|
   s.add_dependency 'rack'
   s.add_development_dependency 'minitest'
   s.add_development_dependency 'rack-test'
+  s.add_development_dependency 'activesupport', '>= 3.0.0'
   s.add_development_dependency 'debugger', '~> 1.1.3'
 end
 
diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb
index 03aafe3..8a986a6 100644
--- a/spec/rack_attack_spec.rb
+++ b/spec/rack_attack_spec.rb
@@ -3,10 +3,6 @@ require_relative 'spec_helper'
 describe 'Rack::Attack' do
   include Rack::Test::Methods
 
-  before do
-    Rack::Attack.block("ip 1.2.3.4") {|req| req.ip == '1.2.3.4' }
-  end
-
   def app
     Rack::Builder.new {
       use Rack::Attack
@@ -14,14 +10,45 @@ describe 'Rack::Attack' do
     }.to_app
   end
 
-  it 'has a block' do
-    Rack::Attack.blocks.class.must_equal Hash
+  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
 
-  it "says hello" do
-    get '/'
-    last_response.status.must_equal 200
-    last_response.body.must_equal 'Hello World'
+  after { Rack::Attack.clear! }
+
+  allow_ok_requests
+
+  describe 'with a block' do
+    before do
+      @bad_ip = '1.2.3.4'
+      Rack::Attack.block("ip #{@bad_ip}") {|req| req.ip == @bad_ip }
+    end
+
+    it('has a block') { Rack::Attack.blocks.key?("ip #{@bad_ip}") }
+
+    it "should block bad requests" do
+      get '/', {}, 'REMOTE_ADDR' => @bad_ip
+      last_response.status.must_equal 503
+    end
+
+    allow_ok_requests
+
+    describe "and with a whitelist" do
+      before do
+        @good_ua = 'GoodUA'
+        Rack::Attack.whitelist("good ua") {|req| req.user_agent == @good_ua }
+      end
+
+      it('has a whitelist'){ Rack::Attack.whitelists.key?("good ua") }
+      it "should allow whitelists before blocks" do
+        get '/', {}, 'REMOTE_ADDR' => @bad_ip, 'HTTP_USER_AGENT' => @good_ua
+        last_response.status.must_equal 200
+      end
+    end
   end