diff options
Diffstat (limited to 'app/services/concerns/rate_limited_service.rb')
-rw-r--r-- | app/services/concerns/rate_limited_service.rb | 93 |
1 files changed, 93 insertions, 0 deletions
diff --git a/app/services/concerns/rate_limited_service.rb b/app/services/concerns/rate_limited_service.rb new file mode 100644 index 00000000000..87cba7814fe --- /dev/null +++ b/app/services/concerns/rate_limited_service.rb @@ -0,0 +1,93 @@ +# 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 |