rack-attack/lib/rack/attack.rb

require 'rack'
module Rack::Attack
  require 'rack/attack/cache'
  require 'rack/attack/throttle'

  class << self

    attr_reader :cache, :notifier

    def whitelist(name, &block)
      (@whitelists ||= {})[name] = block
    end

    def block(name, &block)
      (@blocks ||= {})[name] = block
    end

    def throttle(name, options, &block)
      (@throttles ||= {})[name] = Throttle.new(name, options, block)
    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)
      req = Rack::Request.new(env)

      if whitelisted?(req)
        return @app.call(env)
      end

      if blocked?(req)
        blocked_response
      elsif throttled?(req)
        throttled_response
      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)
      throttles.any? do |name, throttle|
        throttle[req]
      end
    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
end