Commit 725660cd authored by Matthias Käppler's avatar Matthias Käppler

Merge branch...

Merge branch '347550-gitlab-applicationratelimiter-declared-as-class-but-used-as-a-module' into 'master'

Simplify ApplicationRateLimiter

See merge request gitlab-org/gitlab!78561
parents b99ebc5d 1192fb33
......@@ -17,7 +17,7 @@ module RateLimitedService
end
def log_request(request, current_user)
rate_limiter.class.log_request(request, "#{key}_request_limit".to_sym, current_user)
rate_limiter.log_request(request, "#{key}_request_limit".to_sym, current_user)
end
private
......@@ -26,20 +26,19 @@ module RateLimitedService
end
class RateLimiterScopedAndKeyed
attr_reader :key, :opts, :rate_limiter_klass
attr_reader :key, :opts, :rate_limiter
def initialize(key:, opts:, rate_limiter_klass:)
def initialize(key:, opts:, rate_limiter:)
@key = key
@opts = opts
@rate_limiter_klass = rate_limiter_klass
@rate_limiter = rate_limiter
end
def rate_limit!(service)
evaluated_scope = evaluated_scope_for(service)
return if feature_flag_disabled?(evaluated_scope[:project])
rate_limiter = new_rate_limiter(evaluated_scope)
if rate_limiter.throttled?
if rate_limiter.throttled?(key, **opts.merge(scope: evaluated_scope.values, users_allowlist: users_allowlist))
raise RateLimitedError.new(key: key, rate_limiter: rate_limiter), _('This endpoint has been requested too many times. Try again later.')
end
end
......@@ -59,20 +58,16 @@ module RateLimitedService
def feature_flag_disabled?(project)
Feature.disabled?("rate_limited_service_#{key}", project, default_enabled: :yaml)
end
def new_rate_limiter(evaluated_scope)
rate_limiter_klass.new(key, **opts.merge(scope: evaluated_scope.values, users_allowlist: users_allowlist))
end
end
prepended do
attr_accessor :rate_limiter_bypassed
cattr_accessor :rate_limiter_scoped_and_keyed
def self.rate_limit(key:, opts:, rate_limiter_klass: ::Gitlab::ApplicationRateLimiter)
def self.rate_limit(key:, opts:, rate_limiter: ::Gitlab::ApplicationRateLimiter)
self.rate_limiter_scoped_and_keyed = RateLimiterScopedAndKeyed.new(key: key,
opts: opts,
rate_limiter_klass: rate_limiter_klass)
rate_limiter: rate_limiter)
end
end
......
......@@ -5,7 +5,7 @@ module API
# == RateLimiter
#
# Helper that checks if the rate limit for a given endpoint is throttled by calling the
# Gitlab::ApplicationRateLimiter class. If the action is throttled for the current user, the request
# Gitlab::ApplicationRateLimiter module. If the action is throttled for the current user, the request
# will be logged and an error message will be rendered with a Too Many Requests response status.
# See app/controllers/concerns/check_rate_limit.rb for Rails controllers version
module RateLimiter
......
# frozen_string_literal: true
module Gitlab
# This class implements a simple rate limiter that can be used to throttle
# This module implements a simple rate limiter that can be used to throttle
# certain actions. Unlike Rack Attack and Rack::Throttle, which operate at
# the middleware level, this can be used at the controller or API level.
# See CheckRateLimit concern for usage.
class ApplicationRateLimiter
module ApplicationRateLimiter
InvalidKeyError = Class.new(StandardError)
def initialize(key, **options)
@key = key
@options = options
end
def throttled?
self.class.throttled?(key, **options)
end
def threshold_value
options[:threshold] || self.class.threshold(key)
end
def interval_value
self.class.interval(key)
end
class << self
# Application rate limits
#
......@@ -201,9 +184,5 @@ module Gitlab
scoped_user.username.downcase.in?(users_allowlist)
end
end
private
attr_reader :key, :options
end
end
......@@ -6,11 +6,10 @@ RSpec.describe RateLimitedService do
let(:key) { :issues_create }
let(:scope) { [:project, :current_user] }
let(:opts) { { scope: scope, users_allowlist: -> { [User.support_bot.username] } } }
let(:rate_limiter_klass) { ::Gitlab::ApplicationRateLimiter }
let(:rate_limiter_instance) { rate_limiter_klass.new(key, **opts) }
let(:rate_limiter) { ::Gitlab::ApplicationRateLimiter }
describe 'RateLimitedError' do
subject { described_class::RateLimitedError.new(key: key, rate_limiter: rate_limiter_instance) }
subject { described_class::RateLimitedError.new(key: key, rate_limiter: rate_limiter) }
describe '#headers' do
it 'returns a Hash of HTTP headers' do
......@@ -26,7 +25,7 @@ RSpec.describe RateLimitedService do
request = instance_double(Grape::Request)
user = instance_double(User)
expect(rate_limiter_klass).to receive(:log_request).with(request, "#{key}_request_limit".to_sym, user)
expect(rate_limiter).to receive(:log_request).with(request, "#{key}_request_limit".to_sym, user)
subject.log_request(request, user)
end
......@@ -34,7 +33,7 @@ RSpec.describe RateLimitedService do
end
describe 'RateLimiterScopedAndKeyed' do
subject { described_class::RateLimiterScopedAndKeyed.new(key: key, opts: opts, rate_limiter_klass: rate_limiter_klass) }
subject { described_class::RateLimiterScopedAndKeyed.new(key: key, opts: opts, rate_limiter: rate_limiter) }
describe '#rate_limit!' do
let(:project_with_feature_enabled) { create(:project) }
......@@ -49,13 +48,12 @@ RSpec.describe RateLimitedService do
let(:rate_limited_service_issues_create_feature_enabled) { nil }
before do
allow(rate_limiter_klass).to receive(:new).with(key, **evaluated_opts).and_return(rate_limiter_instance)
stub_feature_flags(rate_limited_service_issues_create: rate_limited_service_issues_create_feature_enabled)
end
shared_examples 'a service that does not attempt to throttle' do
it 'does not attempt to throttle' do
expect(rate_limiter_instance).not_to receive(:throttled?)
expect(rate_limiter).not_to receive(:throttled?)
expect(subject.rate_limit!(service)).to be_nil
end
......@@ -63,7 +61,7 @@ RSpec.describe RateLimitedService do
shared_examples 'a service that does attempt to throttle' do
before do
allow(rate_limiter_instance).to receive(:throttled?).and_return(throttled)
allow(rate_limiter).to receive(:throttled?).and_return(throttled)
end
context 'when rate limiting is not in effect' do
......@@ -134,7 +132,7 @@ RSpec.describe RateLimitedService do
end
before do
allow(RateLimitedService::RateLimiterScopedAndKeyed).to receive(:new).with(key: key, opts: opts, rate_limiter_klass: rate_limiter_klass).and_return(rate_limiter_scoped_and_keyed)
allow(RateLimitedService::RateLimiterScopedAndKeyed).to receive(:new).with(key: key, opts: opts, rate_limiter: rate_limiter).and_return(rate_limiter_scoped_and_keyed)
end
context 'bypasses rate limiting' do
......@@ -173,12 +171,12 @@ RSpec.describe RateLimitedService do
end
before do
allow(RateLimitedService::RateLimiterScopedAndKeyed).to receive(:new).with(key: key, opts: opts, rate_limiter_klass: rate_limiter_klass).and_return(rate_limiter_scoped_and_keyed)
allow(RateLimitedService::RateLimiterScopedAndKeyed).to receive(:new).with(key: key, opts: opts, rate_limiter: rate_limiter).and_return(rate_limiter_scoped_and_keyed)
end
context 'and applies rate limiting' do
it 'raises an RateLimitedService::RateLimitedError exception' do
expect(rate_limiter_scoped_and_keyed).to receive(:rate_limit!).with(subject).and_raise(RateLimitedService::RateLimitedError.new(key: key, rate_limiter: rate_limiter_instance))
expect(rate_limiter_scoped_and_keyed).to receive(:rate_limit!).with(subject).and_raise(RateLimitedService::RateLimitedError.new(key: key, rate_limiter: rate_limiter))
expect { subject.execute }.to raise_error(RateLimitedService::RateLimitedError)
end
......
......@@ -17,7 +17,7 @@ RSpec.describe Issues::CreateService do
expect(described_class.rate_limiter_scoped_and_keyed.key).to eq(:issues_create)
expect(described_class.rate_limiter_scoped_and_keyed.opts[:scope]).to eq(%i[project current_user external_author])
expect(described_class.rate_limiter_scoped_and_keyed.rate_limiter_klass).to eq(Gitlab::ApplicationRateLimiter)
expect(described_class.rate_limiter_scoped_and_keyed.rate_limiter).to eq(Gitlab::ApplicationRateLimiter)
end
end
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment