discourse/lib/rate_limiter.rb

require_dependency 'rate_limiter/limit_exceeded'
require_dependency 'rate_limiter/on_create_record'

# A redis backed rate limiter.
class RateLimiter

  attr_reader :max, :secs, :user, :key

  def self.key_prefix
    "l-rate-limit:"
  end

  def self.disable
    @disabled = true
  end

  def self.enable
    @disabled = false
  end

  # We don't observe rate limits in test mode
  def self.disabled?
    @disabled || Rails.env.test?
  end

  def self.clear_all!
    $redis.delete_prefixed(RateLimiter.key_prefix)
  end

  def build_key(type)
    "#{RateLimiter.key_prefix}:#{@user && @user.id}:#{type}"
  end

  def initialize(user, type, max, secs)
    @user = user
    @type = type
    @key = build_key(type)
    @max = max
    @secs = secs
  end

  def clear!
    $redis.del(@key)
  end

  def can_perform?
    rate_unlimited? || is_under_limit?
  end

  def performed!
    return if rate_unlimited?

    if is_under_limit?
      # simple ring buffer.
      $redis.lpush(@key, Time.now.to_i)
      $redis.ltrim(@key, 0, @max - 1)

      # let's ensure we expire this key at some point, otherwise we have leaks
      $redis.expire(@key, @secs * 2)
    else
      raise RateLimiter::LimitExceeded.new(seconds_to_wait, @type)
    end
  end

  def rollback!
    return if RateLimiter.disabled?
    $redis.lpop(@key)
  end

  def remaining
    return @max if @user && @user.staff?

    arr = $redis.lrange(@key, 0, @max) || []
    t0 = Time.now.to_i
    arr.reject! { |a| (t0 - a.to_i) > @secs }
    @max - arr.size
  end

  private

  def seconds_to_wait
    @secs - age_of_oldest
  end

  def age_of_oldest
    # age of oldest event in buffer, in seconds
    Time.now.to_i - $redis.lrange(@key, -1, -1).first.to_i
  end

  def is_under_limit?
    # number of events in buffer less than max allowed? OR
    ($redis.llen(@key) < @max) ||
    # age bigger than silding window size?
    (age_of_oldest > @secs)
  end

  def rate_unlimited?
    !!(RateLimiter.disabled? || (@user && @user.staff?))
  end
end