94 lines
2.7 KiB
Ruby
94 lines
2.7 KiB
Ruby
|
# frozen_string_literal: true
|
||
|
|
|
||
|
|
module RateLimitedService
|
||
|
|
extend ActiveSupport::Concern
|
||
|
|
|
||
|
|
RateLimitedNotSetupError = Class.new(StandardError)
|
||
|
|
|
||
|
|
class RateLimitedError < StandardError
|
||
|
|
def initialize(key:, rate_limiter:)
|
||
|
|
@key = key
|
||
|
|
@rate_limiter = rate_limiter
|
||
|
|
end
|
||
|
|
|
||
|
|
def headers
|
||
|
|
# TODO: This will be fleshed out in https://gitlab.com/gitlab-org/gitlab/-/issues/342370
|
||
|
|
{}
|
||
|
|
end
|
||
|
|
|
||
|
|
def log_request(request, current_user)
|
||
|
|
rate_limiter.class.log_request(request, "#{key}_request_limit".to_sym, current_user)
|
||
|
|
end
|
||
|
|
|
||
|
|
private
|
||
|
|
|
||
|
|
attr_reader :key, :rate_limiter
|
||
|
|
end
|
||
|
|
|
||
|
|
class RateLimiterScopedAndKeyed
|
||
|
|
attr_reader :key, :opts, :rate_limiter_klass
|
||
|
|
|
||
|
|
def initialize(key:, opts:, rate_limiter_klass:)
|
||
|
|
@key = key
|
||
|
|
@opts = opts
|
||
|
|
@rate_limiter_klass = rate_limiter_klass
|
||
|
|
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?
|
||
|
|
raise RateLimitedError.new(key: key, rate_limiter: rate_limiter), _('This endpoint has been requested too many times. Try again later.')
|
||
|
|
end
|
||
|
|
end
|
||
|
|
|
||
|
|
private
|
||
|
|
|
||
|
|
def users_allowlist
|
||
|
|
@users_allowlist ||= opts[:users_allowlist] ? opts[:users_allowlist].call : []
|
||
|
|
end
|
||
|
|
|
||
|
|
def evaluated_scope_for(service)
|
||
|
|
opts[:scope].each_with_object({}) do |var, all|
|
||
|
|
all[var] = service.public_send(var) # rubocop: disable GitlabSecurity/PublicSend
|
||
|
|
end
|
||
|
|
end
|
||
|
|
|
||
|
|
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)
|
||
|
|
self.rate_limiter_scoped_and_keyed = RateLimiterScopedAndKeyed.new(key: key,
|
||
|
|
opts: opts,
|
||
|
|
rate_limiter_klass: rate_limiter_klass)
|
||
|
|
end
|
||
|
|
end
|
||
|
|
|
||
|
|
def execute_without_rate_limiting(*args, **kwargs)
|
||
|
|
self.rate_limiter_bypassed = true
|
||
|
|
execute(*args, **kwargs)
|
||
|
|
ensure
|
||
|
|
self.rate_limiter_bypassed = false
|
||
|
|
end
|
||
|
|
|
||
|
|
def execute(*args, **kwargs)
|
||
|
|
raise RateLimitedNotSetupError if rate_limiter_scoped_and_keyed.nil?
|
||
|
|
|
||
|
|
rate_limiter_scoped_and_keyed.rate_limit!(self) unless rate_limiter_bypassed
|
||
|
|
|
||
|
|
super
|
||
|
|
end
|
||
|
|
end
|