Add throttle support

examples/rack_attack.rb

@@ -0,0 +1,32 @@
+# NB: `req` is a Rack::Request object (basically an env hash with friendly accessor methods)
+
+# Throttle 10 requests/ip/second
+# NB: return value of block is key name for counter
+#     falsy values bypass throttling
+Rack::Attack.throttle("req/ip", :limit => 10, :period => 1) { |req| req.ip }
+
+# Throttle attempts to a particular path. 2 POSTs to /login per second per IP
+Rack::Attack.throttle "logins/ip", :limit => 2, :period => 1 do |req|
+  req.ip if req.post? && req.path_info =~ /^login/
+end
+
+# Throttle login attempts per email, 10/minute/email
+Rack::Attack.throttle "logins/email", :limit => 2, :period => 60 do |req|
+  req.params['email'] unless req.params['email'].blank?
+end
+
+
+# Block cloud IPs from accessing PATH regexp
+Rack::Attack.block "bad_ips from logging in" do |req|
+  req.path =~ /^login/ && bad_ips.include?(req.ip)
+end
+
+# Throttle login attempts
+Rack::Attack.throttle "logins/ip", :limit => 2, :period => 1.second do | req|
+  req.ip if req.post? && req.path_info =~ /^login/
+end
+
+# Whitelist a User-Agent
+Rack::Attack.whitelist 'internal user agent' do |req|
+  req.user_agent =~ 'InternalUserAgent'
+end

lib/rack/attack.rb

@@ -1,6 +1,7 @@
 require 'rack'
 module Rack::Attack
   require 'rack/attack/cache'
+  require 'rack/attack/throttle'
 
   class << self
 
@@ -14,7 +15,8 @@ module Rack::Attack
       (@blocks ||= {})[name] = block
     end
 
-    def throttle
+    def throttle(name, options, &block)
+      (@throttles ||= {})[name] = Throttle.new(name, options, block)
     end
 
     def whitelists; @whitelists ||= {}; end
@@ -22,7 +24,7 @@ module Rack::Attack
     def throttles;  @throttles  ||= {}; end
 
     def new(app)
-      @cache = Cache.new
+      @cache ||= Cache.new
       @notifier = ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
       @app = app
       self
@@ -39,6 +41,7 @@ module Rack::Attack
       if blocked?(req)
         blocked_response
       elsif throttled?(req)
+        throttled_response
       else
         @app.call(env)
       end
@@ -61,7 +64,9 @@ module Rack::Attack
     end
 
     def throttled?(req)
-      false
+      throttles.any? do |name, throttle|
+        throttle[req]
+      end
     end
 
     def instrument(payload)

lib/rack/attack/cache.rb

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

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

lib/rack/attack/whitelist.rb

@@ -0,0 +1,14 @@
+module Rack
+  module Attack
+    class Whitelist
+      attr_reader :name, :block
+      def initialize(name, &block)
+        @name, @block = name, block
+      end
+
+      def [](req)
+      end
+
+    end
+  end
+end

spec/rack_attack_spec.rb

@@ -51,5 +51,27 @@ describe 'Rack::Attack' do
     end
   end
 
+  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
+
 
 end

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"