example_usage.ru

require 'rubygems'
require 'rack'
require 'redis_request_limiter'

# Give 5 requests a minute to make testing really easy
use RedisAPILimiter, :per_user_limit => 5, :reset_requests_every => 60
run proc { |env| [200, {"Content-Type" => "text/html", "Content-Length" => "5"}, ["Hello"]]}

redis_request_limiter.rb

require 'redis'

class RedisRequestLimiter
  
  DEFAULT_OPTIONS = {
    :per_user_limit       => 100,
    :reset_requests_every => 300,
    :override_limit_proc  => proc { |env| [nil, nil] }, # limit, every
    :redis_rate_key       => "api-requests".freeze,
    :header_status_prefix => "X-API-Limit".freeze,
    :env_to_api_key_proc  => proc { |env| env["REMOTE_ADDR"] },
    :limiter_app          => proc do |env|
      headers = {"Content-Type" => "text/html"}
      resets_in = env['x-rack.rate-limiter.resets-at'] - Time.now.to_i
      body = "You are currently limited and will be reset in appx. %s seconds." % resets_in
      headers["Content-Length"] = body.length.to_s
      return 503, headers, [body]
    end
  }
  
  attr_accessor :per_user_limit, :reset_requests_every, :redis_rate_key, :override_limit_proc,
                :header_status_prefix, :limiter_app, :redis, :env_to_api_key_proc
  
  def initialize(app, opts = {})
    @app = app
    # Setup options w/ defaults
    real_opts = DEFAULT_OPTIONS.merge(opts)
    real_opts.each_pair do |key, value|
      send(:"#{key}=", value) if respond_to?(:"#{key}=")
    end
  end
  
  def call(env)
    self.dup._call(env)
  end
  
  def _call(env)
    self.preload_request_info(env)
    env = self.env_with_limit_information(env)
    if self.under_request_limit?(env)
      count_request!(env)
      status, headers, body = @app.call(env)
      self.add_headers(headers)
      [status, headers, body]
    else
      status, headers, body = @limiter_app.call(env)
      self.add_headers(headers)
      [status, headers, body]
    end
  end
  
  def self.redis
    @@redis ||= Redis.new
  end
  
  def self.redis=(value)
    @@redis = value
  end
  
  protected
  
  def env_with_limit_information(env)
    env = env.dup
    env['x-rack.rate-limiter.current-count'] = self.current_request_count
    env['x-rack.rate-limiter.last-period']   = self.last_period
    env['x-rack.rate-limiter.resets-at']     = self.last_period + @reset_requests_every
    env['x-rack.rate-limiter.request-limit'] = @per_user_limit
    return env
  end
  
  def add_headers(header_hash)
    remaining_count = [(@per_user_limit - self.current_request_count), 0].max
    header_hash["#{@header_status_prefix}-Requests-Used"]      = @current_request_count.to_s
    header_hash["#{@header_status_prefix}-Requests-Remaining"] = remaining_count.to_s
    header_hash["#{@header_status_prefix}-Period"]             = self.last_period.to_s
    header_hash["#{@header_status_prefix}-Seconds-Remaining"]  = ((self.last_period + @reset_requests_every) - Time.now.to_i).to_s
    header_hash["#{@header_status_prefix}-Is-Limited"]         = (remaining_count > 0 ? "No" : "Yes")
  end
  
  def under_request_limit?(env)
    if self.last_period == self.current_period
      self.current_request_count < @per_user_limit
    else
      @current_request_count = 0
      @last_period = current_period
      @redis.set(@count_env_key, "0")
      @redis.set(@period_env_key, @current_period)
      true
    end
  end
  
  def count_request!(env)
    @current_request_count = @redis.incr(@count_env_key)
  end
  
  def env_rate_limit_key(env, key = "count")
    "#{@redis_rate_key}:#{@env_to_api_key_proc.call(env)}:#{key}"
  end
  
  def preload_request_info(env)
    current_request_limit, current_reset_every = @override_limit_proc.call(env)
    # Since we dup, we can change ivars
    @reset_requests_every = current_reset_every   if !current_reset_every.nil?
    @per_user_limit       = current_request_limit if !current_request_limit.nil?
    @count_env_key        = env_rate_limit_key(env)
    @period_env_key       = env_rate_limit_key(env, "current-limit-period")
    @redis              ||= self.class.redis
  end
  
  def last_period
    @last_period ||= begin
      raw_stored_period = @redis.get(@period_env_key).to_i
      (raw_stored_period / @reset_requests_every) * @reset_requests_every
    end
  end
  
  def current_period
    @current_period ||= (Time.now.to_i / @reset_requests_every) * @reset_requests_every 
  end
  
  def current_request_count
    @current_request_count ||= @redis.get(@count_env_key).to_i
  end
  
end