mirror of
https://github.com/discourse/discourse.git
synced 2024-12-11 21:53:44 +08:00
208 lines
4.6 KiB
Ruby
208 lines
4.6 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
# A redis backed rate limiter.
|
|
class RateLimiter
|
|
attr_reader :max, :secs, :user, :key, :error_code
|
|
|
|
def self.key_prefix
|
|
"l-rate-limit3:"
|
|
end
|
|
|
|
def self.disable
|
|
@disabled = true
|
|
end
|
|
|
|
def self.enable
|
|
@disabled = false
|
|
end
|
|
|
|
disable if Rails.env.profile?
|
|
|
|
# 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!
|
|
Discourse.redis.delete_prefixed(RateLimiter.key_prefix)
|
|
end
|
|
|
|
def self.clear_all_global!
|
|
Discourse
|
|
.redis
|
|
.without_namespace
|
|
.keys("GLOBAL::#{key_prefix}*")
|
|
.each { |k| Discourse.redis.without_namespace.del k }
|
|
end
|
|
|
|
def build_key(type)
|
|
"#{RateLimiter.key_prefix}:#{@user && @user.id}:#{type}"
|
|
end
|
|
|
|
def initialize(
|
|
user,
|
|
type,
|
|
max,
|
|
secs,
|
|
global: false,
|
|
aggressive: false,
|
|
error_code: nil,
|
|
apply_limit_to_staff: false,
|
|
staff_limit: { max: nil, secs: nil }
|
|
)
|
|
@user = user
|
|
@type = type
|
|
@key = build_key(type)
|
|
@max = max
|
|
@secs = secs
|
|
@global = global
|
|
@aggressive = aggressive
|
|
@error_code = error_code
|
|
@apply_limit_to_staff = apply_limit_to_staff
|
|
@staff_limit = staff_limit
|
|
|
|
# override the default values if staff user, and staff specific max is passed
|
|
if @user&.staff? && !@apply_limit_to_staff && @staff_limit[:max].present?
|
|
@max = @staff_limit[:max]
|
|
@secs = @staff_limit[:secs]
|
|
end
|
|
end
|
|
|
|
def clear!
|
|
redis.del(prefixed_key)
|
|
end
|
|
|
|
def can_perform?
|
|
rate_unlimited? || is_under_limit?
|
|
end
|
|
|
|
def seconds_to_wait(now = Time.now.to_i)
|
|
@secs - age_of_oldest(now)
|
|
end
|
|
|
|
# reloader friendly
|
|
PERFORM_LUA = DiscourseRedis::EvalHelper.new <<~LUA unless defined?(PERFORM_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
|
|
|
|
unless defined?(PERFORM_LUA_AGGRESSIVE)
|
|
PERFORM_LUA_AGGRESSIVE = DiscourseRedis::EvalHelper.new <<~LUA
|
|
local now = tonumber(ARGV[1])
|
|
local secs = tonumber(ARGV[2])
|
|
local max = tonumber(ARGV[3])
|
|
|
|
local key = KEYS[1]
|
|
|
|
local return_val = 0
|
|
|
|
if ((tonumber(redis.call("LLEN", key)) < max) or
|
|
(now - tonumber(redis.call("LRANGE", key, -1, -1)[1])) >= secs) then
|
|
return_val = 1
|
|
else
|
|
return_val = 0
|
|
end
|
|
|
|
redis.call("LPUSH", key, now)
|
|
redis.call("LTRIM", key, 0, max - 1)
|
|
redis.call("EXPIRE", key, secs * 2)
|
|
|
|
return return_val
|
|
LUA
|
|
end
|
|
|
|
def performed!(raise_error: true)
|
|
return true if rate_unlimited?
|
|
now = Time.now.to_i
|
|
if ((@max || 0) <= 0) || rate_limiter_allowed?(now)
|
|
raise RateLimiter::LimitExceeded.new(seconds_to_wait(now), @type, @error_code) if raise_error
|
|
false
|
|
else
|
|
true
|
|
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)
|
|
rescue Redis::CommandError => e
|
|
if e.message =~ /READONLY/
|
|
# TODO,switch to in-memory rate limiter
|
|
else
|
|
raise
|
|
end
|
|
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 rate_limiter_allowed?(now)
|
|
lua, lua_sha = nil
|
|
eval_helper = @aggressive ? PERFORM_LUA_AGGRESSIVE : PERFORM_LUA
|
|
|
|
eval_helper.eval(redis, [prefixed_key], [now, @secs, @max]) == 0
|
|
end
|
|
|
|
def prefixed_key
|
|
if @global
|
|
"GLOBAL::#{key}"
|
|
else
|
|
Discourse.redis.namespace_key(key)
|
|
end
|
|
end
|
|
|
|
def redis
|
|
Discourse.redis.without_namespace
|
|
end
|
|
|
|
def age_of_oldest(now)
|
|
# age of oldest event in buffer, in seconds
|
|
now - redis.lrange(prefixed_key, -1, -1).first.to_i
|
|
end
|
|
|
|
def is_under_limit?
|
|
now = Time.now.to_i
|
|
|
|
# number of events in buffer less than max allowed? OR
|
|
(redis.llen(prefixed_key) < @max) ||
|
|
# age bigger than sliding window size?
|
|
(age_of_oldest(now) >= @secs)
|
|
end
|
|
|
|
def rate_unlimited?
|
|
!!(
|
|
RateLimiter.disabled? || (@user&.staff? && !@apply_limit_to_staff && @staff_limit[:max].nil?)
|
|
)
|
|
end
|
|
end
|