Added Rack::Deflect - DOS protection middleware

Functional, yet expermental middleware for protecting
against Denial-of-service attacks.

README.rdoc

@@ -33,6 +33,7 @@ interface:
 * Rack::NotFound - A default 404 application.
 * Rack::CSSHTTPRequest - Adds CSSHTTPRequest support by encoding responses as
   CSS for cross-site AJAX-style data loading
+* Rack::Deflect - Helps protect against DoS attacks.
 
 === Use
 

lib/rack/contrib.rb

@@ -9,6 +9,7 @@ def self.release
 
   autoload :BounceFavicon,              "rack/contrib/bounce_favicon"
   autoload :CSSHTTPRequest,             "rack/contrib/csshttprequest"
+  autoload :Deflect,                    "rack/contrib/deflect"
   autoload :ETag,                       "rack/contrib/etag"
   autoload :GarbageCollector,           "rack/contrib/garbagecollector"
   autoload :JSONP,                      "rack/contrib/jsonp"

lib/rack/contrib/deflect.rb

@@ -0,0 +1,137 @@
+require 'thread'
+
+# TODO: optional stats
+# TODO: performance
+# TODO: clean up tests
+
+module Rack
+
+  ##
+  # Rack middleware for protecting against Denial-of-service attacks
+  # http://en.wikipedia.org/wiki/Denial-of-service_attack.
+  #
+  # This middleware is designed for small deployments, which most likely
+  # are not utilizing load balancing from other software or hardware. Deflect
+  # current supports the following functionality:
+  #
+  # * Saturation prevention (small DoS attacks, or request abuse)
+  # * Blacklisting of remote addresses
+  # * Whitelisting of remote addresses
+  # * Logging
+  #
+  # === Options:
+  #
+  #   :log                When false logging will be bypassed, otherwise pass an object responding to #puts
+  #   :log_format         Alter the logging format
+  #   :log_date_format    Alter the logging date format
+  #   :request_threshold  Number of requests allowed within the set :interval. Defaults to 100
+  #   :interval           Duration in seconds until the request counter is reset. Defaults to 5
+  #   :block_duration     Duration in seconds that a remote address will be blocked. Defaults to 900 (15 minutes)
+  #   :whitelist          Array of remote addresses which bypass Deflect. NOTE: this does not block others
+  #   :blacklist          Array of remote addresses immediately considered malicious
+  #
+  # === Examples:
+  #
+  #  use Rack::Deflect, :log => $stdout, :request_threshold => 20, :interval => 2, :block_duration => 60
+  #
+  # CREDIT: TJ Holowaychuk <tj@vision-media.ca>
+  #
+
+  class Deflect
+
+    attr_reader :options
+
+    def initialize app, options = {}
+      @mutex = Mutex.new
+      @remote_addr_map = {}
+      @app, @options = app, {
+        :log => false,
+        :log_format => 'deflect(%s): %s',
+        :log_date_format => '%m/%d/%Y',
+        :request_threshold => 100,
+        :interval => 5,
+        :block_duration => 900,
+        :whitelist => [],
+        :blacklist => []
+      }.merge(options)
+    end
+
+    def call env
+      return deflect! if deflect? env
+      status, headers, body = @app.call env
+      [status, headers, body]
+    end
+
+    def deflect!
+      [403, { 'Content-Type' => 'text/html', 'Content-Length' => '0' }, '']
+    end
+
+    def deflect? env
+      @remote_addr = env['REMOTE_ADDR']
+      return false if options[:whitelist].include? @remote_addr
+      return true  if options[:blacklist].include? @remote_addr
+      sync { watch }
+    end
+
+    def log message
+      return unless options[:log]
+      options[:log].puts(options[:log_format] % [Time.now.strftime(options[:log_date_format]), message])
+    end
+
+    def sync &block
+      @mutex.synchronize(&block)
+    end
+
+    def map
+      @remote_addr_map[@remote_addr] ||= {
+        :expires => Time.now + options[:interval],
+        :requests => 0
+      }
+    end
+
+    def watch
+      increment_requests
+      clear! if watch_expired? and not blocked?
+      clear! if blocked? and block_expired?
+      block! if watching? and exceeded_request_threshold?
+      blocked?
+    end
+
+    def block!
+      return if blocked?
+      log "blocked #{@remote_addr}"
+      map[:block_expires] = Time.now + options[:block_duration]
+    end
+
+    def blocked?
+      map.has_key? :block_expires
+    end
+
+    def block_expired?
+      map[:block_expires] < Time.now rescue false
+    end
+
+    def watching?
+      @remote_addr_map.has_key? @remote_addr
+    end
+
+    def clear!
+      return unless watching?
+      log "released #{@remote_addr}" if blocked?
+      @remote_addr_map.delete @remote_addr
+    end
+
+    def increment_requests
+      map[:requests] += 1
+    end
+
+    def exceeded_request_threshold?
+      map[:requests] > options[:request_threshold]
+    end
+
+    def watch_expired?
+      map[:expires] <= Time.now
+    end
+
+  end
+end

test/spec_rack_deflect.rb

@@ -0,0 +1,107 @@
+require 'test/spec'
+require 'rack/mock'
+require 'rack/contrib/deflect'
+
+context "Rack::Deflect" do
+
+  setup do
+    @app = lambda { |env| [200, { 'Content-Type' => 'text/plain' }, 'cookies'] }
+    @mock_addr_1 = '111.111.111.111'
+    @mock_addr_2 = '222.222.222.222'
+    @mock_addr_3 = '333.333.333.333'
+  end
+
+  def mock_env remote_addr, path = '/'
+    Rack::MockRequest.env_for path, { 'REMOTE_ADDR' => remote_addr }
+  end
+
+  def mock_deflect options = {}
+    Rack::Deflect.new @app, options
+  end
+
+  specify "should allow regular requests to follow through" do
+    app = mock_deflect
+    status, headers, body = app.call mock_env(@mock_addr_1)
+    status.should.equal 200
+    body.should.equal 'cookies'
+  end
+
+  specify "should deflect requests exceeding the request threshold" do
+    log = StringIO.new
+    app = mock_deflect :request_threshold => 5, :interval => 10, :block_duration => 10, :log => log
+    env = mock_env @mock_addr_1
+
+    # First 5 should be fine
+    5.times do
+      status, headers, body = app.call env
+      status.should.equal 200
+      body.should.equal 'cookies'
+    end
+
+    # Remaining requests should fail for 10 seconds
+    10.times do
+      status, headers, body = app.call env
+      status.should.equal 403
+      body.should.equal ''
+    end
+
+    # Log should reflect that we have blocked an address
+    log.string.should.match(/^deflect\(\d+\/\d+\/\d+\): blocked 111.111.111.111\n/)
+  end
+
+  specify "should expire blocking" do
+    log = StringIO.new
+    app = mock_deflect :request_threshold => 5, :interval => 2, :block_duration => 2, :log => log
+    env = mock_env @mock_addr_1
+
+    # First 5 should be fine
+    5.times do
+      status, headers, body = app.call env
+      status.should.equal 200
+      body.should.equal 'cookies'
+    end
+
+    # Exceeds request threshold
+    status, headers, body = app.call env
+    status.should.equal 403
+    body.should.equal ''
+
+    # Allow block to expire
+    sleep 3
+
+    # Another 5 is fine now
+    5.times do
+      status, headers, body = app.call env
+      status.should.equal 200
+      body.should.equal 'cookies'
+    end
+
+    # Log should reflect block and release
+    log.string.should.match(/deflect.*: blocked 111\.111\.111\.111\ndeflect.*: released 111\.111\.111\.111\n/)
+  end
+
+  specify "should allow whitelisting of remote addresses" do
+    app = mock_deflect :whitelist => [@mock_addr_1], :request_threshold => 5, :interval => 2
+    env = mock_env @mock_addr_1
+
+    # Whitelisted addresses are always fine
+    10.times do
+      status, headers, body = app.call env
+      status.should.equal 200
+      body.should.equal 'cookies'
+    end
+  end
+
+  specify "should allow blacklisting of remote addresses" do
+    app = mock_deflect :blacklist => [@mock_addr_2]
+
+    status, headers, body = app.call mock_env(@mock_addr_1)
+    status.should.equal 200
+    body.should.equal 'cookies'
+
+    status, headers, body = app.call mock_env(@mock_addr_2)
+    status.should.equal 403
+    body.should.equal ''
+  end
+
+end