rack-attack 0.0.1 → 0.0.3

data/README.md

@@ -1 +1,9 @@
 # Rack::Attack - middleware for throttling & blocking abusive clients
+
+## Processing order
+ * If any whitelist matches, the request is allowed
+ * If any blacklist matches, the request is blocked (unless a whitelist matched)
+ * If any throttle matches, the request is throttled (unless a whitelist or blacklist matched)
+
+[![Travis CI](https://secure.travis-ci.org/ktheory/rack-attack.png)](http://travis-ci.org/ktheory/rack-attack)
+

data/Rakefile

@@ -0,0 +1,9 @@
+require "rubygems"
+require "bundler/setup"
+require 'rake/testtask'
+
+Rake::TestTask.new do |t|
+  t.pattern = "spec/*_spec.rb"
+end
+
+task :default => :test

data/lib/rack/attack.rb

@@ -1,29 +1,86 @@
 require 'rack'
-module Rack
-  class Attack
-    class << self
+module Rack::Attack
+  require 'rack/attack/cache'
+  require 'rack/attack/throttle'
+  require 'rack/attack/whitelist'
+  require 'rack/attack/blacklist'
 
-      attr_reader :blocks, :throttles, :whitelists
+  class << self
 
-      def block(name, &block)
-        (@blocks ||= {})[name] = block
+    attr_reader :cache, :notifier
+
+    def whitelist(name, &block)
+      (@whitelists ||= {})[name] = Whitelist.new(name, block)
+    end
+
+    def blacklist(name, &block)
+      (@blacklists ||= {})[name] = Blacklist.new(name, block)
+    end
+
+    def throttle(name, options, &block)
+      (@throttles ||= {})[name] = Throttle.new(name, options, block)
+    end
+
+    def whitelists; @whitelists ||= {}; end
+    def blacklists; @blacklists ||= {}; 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)
+      req = Rack::Request.new(env)
+
+      if whitelisted?(req)
+        return @app.call(env)
       end
 
-      def throttle
+      if blacklisted?(req)
+        blacklisted_response
+      elsif throttled?(req)
+        throttled_response
+      else
+        @app.call(env)
       end
+    end
 
-      def whitelist
+    def whitelisted?(req)
+      whitelists.any? do |name, whitelist|
+        whitelist[req]
       end
+    end
 
+    def blacklisted?(req)
+      blacklists.any? do |name, blacklist|
+        blacklist[req]
+      end
     end
 
-    def initialize(app)
-      @app = app
+    def throttled?(req)
+      throttles.any? do |name, throttle|
+        throttle[req]
+      end
     end
 
-    def call(env)
-      puts 'Rack attack!'
-      @app.call(env)
+    def instrument(payload)
+      notifier.instrument('rack.attack', payload) if notifier
+    end
+
+    def blacklisted_response
+      [503, {}, ['Blocked']]
+    end
+
+    def throttled_response
+      [503, {}, ['Throttled']]
+    end
+
+    def clear!
+      @whitelists, @blacklists, @throttles = {}, {}, {}
     end
 
   end

data/lib/rack/attack/blacklist.rb

@@ -0,0 +1,13 @@
+require_relative 'check'
+module Rack
+  module Attack
+    class Blacklist < Check
+      def initialize(name, block)
+        super
+        @type = :blacklist
+      end
+
+    end
+  end
+end
+

data/lib/rack/attack/cache.rb

@@ -0,0 +1,23 @@
+module Rack
+  module Attack
+    class Cache
+
+      attr_accessor :store, :prefix
+      def initialize
+        @store = ::Rails.cache if defined?(::Rails.cache)
+        @prefix = 'rack::attack'
+      end
+
+      def count(unprefixed_key, expires_in)
+        key = "#{prefix}:#{unprefixed_key}"
+        result = store.increment(key, 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

data/lib/rack/attack/check.rb

@@ -0,0 +1,19 @@
+module Rack
+  module Attack
+    class Check
+      attr_reader :name, :block, :type
+      def initialize(name, block)
+        @name, @block = name, block
+        @type = nil
+      end
+
+      def [](req)
+        block[req].tap {|match|
+          Rack::Attack.instrument(:type => type, :name => name, :request => req) if match
+        }
+      end
+
+    end
+  end
+end
+

data/lib/rack/attack/throttle.rb

@@ -0,0 +1,31 @@
+module Rack
+  module Attack
+    class Throttle
+      attr_reader :name, :limit, :period, :block
+      def initialize(name, options, block)
+        @name, @block = name, block
+        [:limit, :period].each do |opt|
+          raise ArgumentError.new("Must pass #{opt.inspect} option") unless options[opt]
+        end
+        @limit  = options[:limit]
+        @period = options[:period]
+      end
+
+      def cache
+        Rack::Attack.cache
+      end
+
+      def [](req)
+        discriminator = block[req]
+        return false unless discriminator
+
+        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
+      end
+
+    end
+  end
+end

data/lib/rack/attack/version.rb

@@ -1,5 +1,5 @@
 module Rack
-  class Attack
-    VERSION = '0.0.1'
+  module Attack
+    VERSION = '0.0.3'
   end
 end

data/lib/rack/attack/whitelist.rb

@@ -0,0 +1,12 @@
+require_relative 'check'
+module Rack
+  module Attack
+    class Whitelist < Check
+      def initialize(name, block)
+        super
+        @type = :whitelist
+      end
+
+    end
+  end
+end

data/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,67 @@ 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
+
+  after { Rack::Attack.clear! }
+
+  allow_ok_requests
+
+  describe 'with a blacklist' do
+    before do
+      @bad_ip = '1.2.3.4'
+      Rack::Attack.blacklist("ip #{@bad_ip}") {|req| req.ip == @bad_ip }
+    end
+
+    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
+
+    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 blacklists" do
+        get '/', {}, 'REMOTE_ADDR' => @bad_ip, 'HTTP_USER_AGENT' => @good_ua
+        last_response.status.must_equal 200
+      end
+    end
   end
 
-  it "says hello" do
-    get '/'
-    last_response.status.must_equal 200
-    last_response.body.must_equal 'Hello World'
+  describe 'with a throttle' do
+    before do
+      Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new
+      Rack::Attack.throttle('ip/sec', :limit => 1, :period => 1) { |req| req.ip }
+    end
+
+    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'
+      end
+      last_response.status.must_equal 503
+    end
   end
 
 

data/spec/spec_helper.rb

@@ -3,5 +3,7 @@ require "bundler/setup"
 
 require "minitest/autorun"
 require "rack/test"
+require 'debugger'
+require 'active_support'
 
 require "rack/attack"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: rack-attack
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.0.3
   prerelease: 
 platform: ruby
 authors:
@@ -9,7 +9,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2012-07-25 00:00:00.000000000 Z
+date: 2012-07-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rack
@@ -59,6 +59,38 @@ dependencies:
     - - ! '>='
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: activesupport
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: 3.0.0
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: 3.0.0
 - !ruby/object:Gem::Dependency
   name: debugger
   requirement: !ruby/object:Gem::Requirement
@@ -81,8 +113,14 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- lib/rack/attack/blacklist.rb
+- lib/rack/attack/cache.rb
+- lib/rack/attack/check.rb
+- lib/rack/attack/throttle.rb
 - lib/rack/attack/version.rb
+- lib/rack/attack/whitelist.rb
 - lib/rack/attack.rb
+- Rakefile
 - LICENSE
 - README.md
 - spec/rack_attack_spec.rb
@@ -108,7 +146,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 1.8.23
+rubygems_version: 1.8.24
 signing_key: 
 specification_version: 3
 summary: Block & throttle abusive requests