155 lines
3.2 KiB
Ruby
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
|