module Rack; module Throttle
  ##
  # This is the base class for rate limiter implementations.
  #
  # @example Defining a rate limiter subclass
  #   class MyLimiter < Limiter
  #     def allowed?(request)
  #       # TODO: custom logic goes here
  #     end
  #   end
  #
  class Limiter
    attr_reader :app, :options, :matchers

    ##
    # @param  [#call]                       app
    # @param  [Hash{Symbol => Object}]      options
    # @option options [String]  :cache      (Hash.new)
    # @option options [String]  :key        (nil)
    # @option options [String]  :key_prefix (nil)
    # @option options [Integer] :code       (403)
    # @option options [String]  :message    ("Rate Limit Exceeded")
    def initialize(app, options = {})
      rules = options.delete(:rules) || {}
      @app, @options, @matchers = app, options, []
      @matchers += Array(rules[:ip]).map { |rule| IpMatcher.new(rule) } if rules[:ip]
      @matchers += Array(rules[:url]).map { |rule| UrlMatcher.new(rule) } if rules[:url]
      @matchers += Array(rules[:user_agent]).map { |rule| UserAgentMatcher.new(rule) } if rules[:user_agent]
      @matchers += Array(rules[:method]).map { |rule| MethodMatcher.new(rule) } if rules[:method]
    end

    ##
    # @param  [Hash{String => String}] env
    # @return [Array(Integer, Hash, #each)]
    # @see    http://rack.rubyforge.org/doc/SPEC.html
    def call(env)
      request = Rack::Request.new(env)
      match_results = @matchers.map { |m| m.match?(request) }.uniq
      applicable = @matchers.empty? || match_results == [true]
      if applicable and !allowed?(request)
        rate_limit_exceeded
      else
        app.call(env)
      end
    end

    ##
    # Returns `true` if no :url_rule regex or if the request path
    # matches the :url regex, `false` otherwise.
    #
    # You can override this class, though that might be weird.
    #
    # @param  [String] path
    # @return [Boolean]
    def restricted_url?(path)
      options[:url_rule].nil? || options[:url_rule].match(path)
    end

    ##
    # Returns `false` if the rate limit has been exceeded for the given
    # `request`, or `true` otherwise.
    #
    # Override this method in subclasses that implement custom rate limiter
    # strategies.
    #
    # @param  [Rack::Request] request
    # @return [Boolean]
    def allowed?(request)
      case
        when whitelisted?(request) then true
        when blacklisted?(request) then false
        else true # override in subclasses
      end
    end

    ##
    # Returns `true` if the originator of the given `request` is whitelisted
    # (not subject to further rate limits).
    #
    # The default implementation always returns `false`. Override this
    # method in a subclass to implement custom whitelisting logic.
    #
    # @param  [Rack::Request] request
    # @return [Boolean]
    # @abstract
    def whitelisted?(request)
      false
    end

    ##
    # Returns `true` if the originator of the given `request` is blacklisted
    # (not honoring rate limits, and thus permanently forbidden access
    # without the need to maintain further rate limit counters).
    #
    # The default implementation always returns `false`. Override this
    # method in a subclass to implement custom blacklisting logic.
    #
    # @param  [Rack::Request] request
    # @return [Boolean]
    # @abstract
    def blacklisted?(request)
      false
    end

    protected

    ##
    # @return [Hash]
    def cache
      case cache = (options[:cache] ||= {})
        when Proc then cache.call
        else cache
      end
    end

    ##
    # @param  [String] key
    def cache_has?(key)
      case
        when cache.respond_to?(:has_key?)
          cache.has_key?(key)
        when cache.respond_to?(:get)
          cache.get(key) rescue false
        else false
      end
    end

    ##
    # @param  [String] key
    # @return [Object]
    def cache_get(key, default = nil)
      case
      when cache.respond_to?(:[])
        cache[key] || default
      when cache.respond_to?(:get)
        cache.get(key) || default
      end
    end

    ##
    # @param  [String] key
    # @param  [Object] value
    # @return [void]
    def cache_set(key, value)
      case
      when cache.respond_to?(:[]=)
        begin
          cache[key] = value
        rescue TypeError => e
          # GDBM throws a "TypeError: can't convert Float into String"
          # exception when trying to store a Float. On the other hand, we
          # don't want to unnecessarily coerce the value to a String for
          # any stores that do support other data types (e.g. in-memory
          # hash objects). So, this is a compromise.
          cache[key] = value.to_s
        end
      when cache.respond_to?(:set)
        cache.set(key, value)
      end
    end

    ##
    # @param  [Rack::Request] request
    # @return [String]
    def cache_key(request)
      id = client_identifier(request)
      id = options[:key].call(request) if options.has_key?(:key)
      id = [options[:key_prefix], id].join(':') if options.has_key?(:key_prefix)
      @matchers.each do |matcher|
        id += ":#{matcher.identifier}"
      end

      id
    end

    ##
    # @param  [Rack::Request] request
    # @return [String]
    def client_identifier(request)
      request.ip.to_s
    end

    ##
    # @param  [Rack::Request] request
    # @return [Float]
    def request_start_time(request)
      case
      when request.env.has_key?('HTTP_X_REQUEST_START')
        request.env['HTTP_X_REQUEST_START'].to_f / 1000
      else
        Time.now.to_f
      end
    end

    ##
    # Outputs a `Rate Limit Exceeded` error.
    #
    # @return [Array(Integer, Hash, #each)]
    def rate_limit_exceeded
      headers = respond_to?(:retry_after) ? {'Retry-After' => retry_after.to_f.ceil.to_s} : {}
      http_error(options[:code] || 403, options[:message] || 'Rate Limit Exceeded', headers)
    end

    ##
    # Outputs an HTTP `4xx` or `5xx` response.
    #
    # @param  [Integer]                code
    # @param  [String, #to_s]          message
    # @param  [Hash{String => String}] headers
    # @return [Array(Integer, Hash, #each)]
    def http_error(code, message = nil, headers = {})
      [ code,
        { 'Content-Type' => 'text/plain; charset=utf-8' }.merge(headers),
        Array( http_status(code) + (message.nil? ? "\n" : " (#{message})\n") )
      ]
    end

    ##
    # Returns the standard HTTP status message for the given status `code`.
    #
    # @param  [Integer] code
    # @return [String]
    def http_status(code)
      [code, Rack::Utils::HTTP_STATUS_CODES[code]].join(' ')
    end
  end
end; end