discourse/lib/rate_limiter.rb

155 lines
3.2 KiB
Ruby

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