From e166e87fb961596f34aa7fc2c7a48c96b0c52c07 Mon Sep 17 00:00:00 2001
From: Aaron Suggs <aaron@ktheory.com>
Date: Fri, 27 Jul 2012 17:22:49 -0400
Subject: [PATCH] Add throttle support

---
 examples/rack_attack.rb      | 32 ++++++++++++++++++++++++++++++++
 lib/rack/attack.rb           | 11 ++++++++---
 lib/rack/attack/cache.rb     |  8 +++++---
 lib/rack/attack/throttle.rb  | 31 +++++++++++++++++++++++++++++++
 lib/rack/attack/whitelist.rb | 14 ++++++++++++++
 spec/rack_attack_spec.rb     | 22 ++++++++++++++++++++++
 spec/spec_helper.rb          |  2 ++
 7 files changed, 114 insertions(+), 6 deletions(-)
 create mode 100644 examples/rack_attack.rb
 create mode 100644 lib/rack/attack/throttle.rb
 create mode 100644 lib/rack/attack/whitelist.rb

diff --git a/examples/rack_attack.rb b/examples/rack_attack.rb
new file mode 100644
index 0000000..0ebabd5
--- /dev/null
+++ b/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
diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index c72c62a..c863843 100644
--- a/lib/rack/attack.rb
+++ b/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)
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index d7d7c91..747812d 100644
--- a/lib/rack/attack/cache.rb
+++ b/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)
-        result = store.increment(1, :expires_in => expires_in)
+      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)
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
new file mode 100644
index 0000000..f31d2ad
--- /dev/null
+++ b/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
diff --git a/lib/rack/attack/whitelist.rb b/lib/rack/attack/whitelist.rb
new file mode 100644
index 0000000..93b7e8f
--- /dev/null
+++ b/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
diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb
index 8a986a6..469a748 100644
--- a/spec/rack_attack_spec.rb
+++ b/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
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 8cd7c86..b36fc3c 100644
--- a/spec/spec_helper.rb
+++ b/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"