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-limit3:"
  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
  end

  # Only used in test, only clears current namespace, does not clear globals
  def self.clear_all!
    $redis.delete_prefixed(RateLimiter.key_prefix)
  end

  def self.clear_all_global!
    $redis.without_namespace.keys("GLOBAL::#{key_prefix}*").each do |k|
      $redis.without_namespace.del k
    end
  end

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

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

  def clear!
    redis.del(prefixed_key)
  end

  def can_perform?
    rate_unlimited? || is_under_limit?
  end

  # reloader friendly
  unless defined? PERFORM_LUA
    PERFORM_LUA = <<~LUA
      local now = tonumber(ARGV[1])
      local secs = tonumber(ARGV[2])
      local max = tonumber(ARGV[3])

      local key = KEYS[1]


      if ((tonumber(redis.call("LLEN", key)) < max) or
          (now - tonumber(redis.call("LRANGE", key, -1, -1)[1])) > secs) then
        redis.call("LPUSH", key, now)
        redis.call("LTRIM", key, 0, max - 1)
        redis.call("EXPIRE", key, secs * 2)

        return 1
      else
        return 0
      end
    LUA

    PERFORM_LUA_SHA = Digest::SHA1.hexdigest(PERFORM_LUA)
  end

  def performed!
    return if rate_unlimited?
    now = Time.now.to_i
    if eval_lua(PERFORM_LUA, PERFORM_LUA_SHA, [prefixed_key], [now, @secs, @max]) == 0
      raise RateLimiter::LimitExceeded.new(seconds_to_wait, @type)
    end
  rescue Redis::CommandError => e
    if e.message =~ /READONLY/
      # TODO,switch to in-memory rate limiter
    else
      raise
    end
  end

  def rollback!
    return if RateLimiter.disabled?
    redis.lpop(prefixed_key)
  end

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

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

  private

  def prefixed_key
    if @global
      "GLOBAL::#{key}"
    else
      $redis.namespace_key(key)
    end
  end

  def redis
    $redis.without_namespace
  end

  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(prefixed_key, -1, -1).first.to_i
  end

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

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

  def eval_lua(lua, sha, keys, args)
    redis.evalsha(sha, keys, args)
  rescue Redis::CommandError => e
    if e.to_s =~ /^NOSCRIPT/
      redis.eval(lua, keys, args)
    else
      raise
    end
  end
end