Revert "Revert "PERF: improve speed of rate limiter""

This reverts commit 2373d85239.
This commit is contained in:
Sam 2017-12-04 21:23:11 +11:00
parent 2373d85239
commit dd70ef3abf
7 changed files with 185 additions and 82 deletions

View File

@ -20,9 +20,10 @@ class RateLimiter
# We don't observe rate limits in test mode
def self.disabled?
@disabled || Rails.env.test?
@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
@ -31,65 +32,102 @@ class RateLimiter
"#{RateLimiter.key_prefix}:#{@user && @user.id}:#{type}"
end
def initialize(user, type, max, secs)
def initialize(user, type, max, secs, global: false)
@user = user
@type = type
@key = build_key(type)
@max = max
@secs = secs
@global = false
end
def clear!
$redis.del(@key)
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?
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
now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
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(@key)
redis.lpop(prefixed_key)
end
def remaining
return @max if @user && @user.staff?
arr = $redis.lrange(@key, 0, @max) || []
t0 = Time.now.to_i
arr = redis.lrange(prefixed_key, 0, @max) || []
t0 = Process.clock_gettime(Process::CLOCK_MONOTONIC)
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(@key, -1, -1).first.to_i
Process.clock_gettime(Process::CLOCK_MONOTONIC) - 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(@key) < @max) ||
(redis.llen(prefixed_key) < @max) ||
# age bigger than silding window size?
(age_of_oldest > @secs)
end
@ -97,4 +135,14 @@ class RateLimiter
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

View File

@ -172,38 +172,47 @@ describe Auth::DefaultCurrentUserProvider do
end
it "can only try 10 bad cookies a minute" do
user = Fabricate(:user)
token = UserAuthToken.generate!(user_id: user.id)
context "rate limiting" do
provider('/').log_on_user(user, {}, {})
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.new(nil, "cookie_auth_10.0.0.1", 10, 60).clear!
RateLimiter.new(nil, "cookie_auth_10.0.0.2", 10, 60).clear!
ip = "10.0.0.1"
env = { "HTTP_COOKIE" => "_t=#{SecureRandom.hex}", "REMOTE_ADDR" => ip }
10.times do
provider('/', env).current_user
before do
RateLimiter.enable
end
expect {
provider('/', env).current_user
}.to raise_error(Discourse::InvalidAccess)
after do
RateLimiter.disable
end
expect {
env["HTTP_COOKIE"] = "_t=#{token.unhashed_auth_token}"
provider("/", env).current_user
}.to raise_error(Discourse::InvalidAccess)
it "can only try 10 bad cookies a minute" do
user = Fabricate(:user)
token = UserAuthToken.generate!(user_id: user.id)
env["REMOTE_ADDR"] = "10.0.0.2"
provider('/').log_on_user(user, {}, {})
expect {
provider('/', env).current_user
}.not_to raise_error
RateLimiter.new(nil, "cookie_auth_10.0.0.1", 10, 60).clear!
RateLimiter.new(nil, "cookie_auth_10.0.0.2", 10, 60).clear!
ip = "10.0.0.1"
env = { "HTTP_COOKIE" => "_t=#{SecureRandom.hex}", "REMOTE_ADDR" => ip }
10.times do
provider('/', env).current_user
end
expect {
provider('/', env).current_user
}.to raise_error(Discourse::InvalidAccess)
expect {
env["HTTP_COOKIE"] = "_t=#{token.unhashed_auth_token}"
provider("/", env).current_user
}.to raise_error(Discourse::InvalidAccess)
env["REMOTE_ADDR"] = "10.0.0.2"
expect {
provider('/', env).current_user
}.not_to raise_error
end
end
it "correctly removes invalid cookies" do
@ -295,44 +304,53 @@ describe Auth::DefaultCurrentUserProvider do
end
it "rate limits api usage" do
context "rate limiting" do
RateLimiter.stubs(:disabled?).returns(false)
limiter1 = RateLimiter.new(nil, "user_api_day_#{api_key.key}", 10, 60)
limiter2 = RateLimiter.new(nil, "user_api_min_#{api_key.key}", 10, 60)
limiter1.clear!
limiter2.clear!
SiteSetting.max_user_api_reqs_per_day = 3
SiteSetting.max_user_api_reqs_per_minute = 4
params = {
"REQUEST_METHOD" => "GET",
"HTTP_USER_API_KEY" => api_key.key,
}
3.times do
provider("/", params).current_user
before do
RateLimiter.enable
end
expect {
provider("/", params).current_user
}.to raise_error(RateLimiter::LimitExceeded)
SiteSetting.max_user_api_reqs_per_day = 4
SiteSetting.max_user_api_reqs_per_minute = 3
limiter1.clear!
limiter2.clear!
3.times do
provider("/", params).current_user
after do
RateLimiter.disable
end
expect {
provider("/", params).current_user
}.to raise_error(RateLimiter::LimitExceeded)
it "rate limits api usage" do
limiter1 = RateLimiter.new(nil, "user_api_day_#{api_key.key}", 10, 60)
limiter2 = RateLimiter.new(nil, "user_api_min_#{api_key.key}", 10, 60)
limiter1.clear!
limiter2.clear!
SiteSetting.max_user_api_reqs_per_day = 3
SiteSetting.max_user_api_reqs_per_minute = 4
params = {
"REQUEST_METHOD" => "GET",
"HTTP_USER_API_KEY" => api_key.key,
}
3.times do
provider("/", params).current_user
end
expect {
provider("/", params).current_user
}.to raise_error(RateLimiter::LimitExceeded)
SiteSetting.max_user_api_reqs_per_day = 4
SiteSetting.max_user_api_reqs_per_minute = 3
limiter1.clear!
limiter2.clear!
3.times do
provider("/", params).current_user
end
expect {
provider("/", params).current_user
}.to raise_error(RateLimiter::LimitExceeded)
end
end
end
end

View File

@ -8,11 +8,14 @@ describe RateLimiter do
context 'disabled' do
before do
RateLimiter.stubs(:disabled?).returns(true)
rate_limiter.performed!
rate_limiter.performed!
end
it "should be disabled" do
expect(RateLimiter.disabled?).to eq(true)
end
it "returns true for can_perform?" do
expect(rate_limiter.can_perform?).to eq(true)
end
@ -25,10 +28,41 @@ describe RateLimiter do
context 'enabled' do
before do
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.enable
rate_limiter.clear!
end
after do
RateLimiter.disable
end
context 'global rate limiter' do
it 'can operate in global mode' do
limiter = RateLimiter.new(nil, "test", 2, 10, global: true)
limiter.clear!
limiter.performed!
limiter.performed!
expect { limiter.performed! }.to raise_error(RateLimiter::LimitExceeded)
end
end
context 'handles readonly' do
before do
$redis.without_namespace.slaveof '10.0.0.1', '99999'
end
after do
$redis.without_namespace.slaveof 'no', 'one'
end
it 'does not explode' do
expect { rate_limiter.performed! }.not_to raise_error
end
end
context 'never done' do
it "should perform right away" do
expect(rate_limiter.can_perform?).to eq(true)

View File

@ -756,7 +756,7 @@ describe SessionController do
context 'rate limited' do
it 'rate limits login' do
SiteSetting.max_logins_per_ip_per_hour = 2
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.enable
RateLimiter.clear_all!
2.times do

View File

@ -20,7 +20,7 @@ describe PostAction do
it "limits redo/undo" do
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.enable
PostAction.act(eviltrout, post, PostActionType.types[:like])
PostAction.remove_act(eviltrout, post, PostActionType.types[:like])

View File

@ -579,7 +579,7 @@ describe Topic do
it "rate limits topic invitations" do
SiteSetting.max_topic_invitations_per_day = 2
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.enable
RateLimiter.clear_all!
start = Time.now.tomorrow.beginning_of_day
@ -1694,7 +1694,7 @@ describe Topic do
SiteSetting.max_replies_in_first_day = 1
SiteSetting.stubs(:client_settings_json).returns(SiteSetting.client_settings_json_uncached)
RateLimiter.stubs(:rate_limit_create_topic).returns(100)
RateLimiter.stubs(:disabled?).returns(false)
RateLimiter.enable
RateLimiter.clear_all!
end

View File

@ -106,7 +106,8 @@ RSpec.configure do |config|
# perf benefit seems low (shaves 20 secs off a 4 minute test suite)
#
# $redis = DiscourseMockRedis.new
#
RateLimiter.disable
PostActionNotifier.disable
SearchIndexer.disable
UserActionCreator.disable
@ -202,6 +203,7 @@ def freeze_time(now = Time.now)
Time.stubs(:now).returns(time)
Date.stubs(:today).returns(datetime.to_date)
TrackTimeStub.stubs(:stubbed).returns(true)
Process.stubs(:clock_gettime).with(Process::CLOCK_MONOTONIC).returns(datetime.to_f)
if block_given?
begin
@ -217,6 +219,7 @@ def unfreeze_time
Time.unstub(:now)
Date.unstub(:today)
TrackTimeStub.unstub(:stubbed)
Process.unstub(:clock_gettime)
end
def file_from_fixtures(filename, directory = "images")