discourse/lib/rate_limiter.rb

101 lines
2.0 KiB
Ruby
Raw Normal View History

2013-02-05 14:16:51 -05:00
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
2013-02-05 14:16:51 -05:00
# We don't observe rate limits in test mode
def self.disabled?
@disabled || Rails.env.test?
2013-02-05 14:16:51 -05:00
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)
2013-02-05 14:16:51 -05:00
@user = user
@type = type
@key = build_key(type)
2013-02-05 14:16:51 -05:00
@max = max
@secs = secs
end
def clear!
$redis.del(@key)
end
def can_perform?
rate_unlimited? || is_under_limit?
2013-02-05 14:16:51 -05:00
end
def performed!
return if rate_unlimited?
2013-02-05 14:16:51 -05:00
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)
2013-02-05 14:16:51 -05:00
end
end
def rollback!
return if RateLimiter.disabled?
$redis.lpop(@key)
2013-02-05 14:16:51 -05:00
end
def remaining
return @max if @user && @user.staff?
arr = $redis.lrange(@key, 0, @max) || []
t0 = Time.now.to_i
2017-07-27 21:20:09 -04:00
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?
2014-08-17 20:55:30 -04:00
!!(RateLimiter.disabled? || (@user && @user.staff?))
end
2013-02-05 14:16:51 -05:00
end