101 lines
2.0 KiB
Ruby
101 lines
2.0 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-limit:"
|
|
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 || Rails.env.test?
|
|
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)
|
|
@user = user
|
|
@type = type
|
|
@key = build_key(type)
|
|
@max = max
|
|
@secs = secs
|
|
end
|
|
|
|
def clear!
|
|
$redis.del(@key)
|
|
end
|
|
|
|
def can_perform?
|
|
rate_unlimited? || is_under_limit?
|
|
end
|
|
|
|
def performed!
|
|
return if rate_unlimited?
|
|
|
|
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)
|
|
end
|
|
end
|
|
|
|
def rollback!
|
|
return if RateLimiter.disabled?
|
|
$redis.lpop(@key)
|
|
end
|
|
|
|
def remaining
|
|
return @max if @user && @user.staff?
|
|
|
|
arr = $redis.lrange(@key, 0, @max) || []
|
|
t0 = Time.now.to_i
|
|
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?
|
|
!!(RateLimiter.disabled? || (@user && @user.staff?))
|
|
end
|
|
end
|