debian-mirror-gitlab/app/services/concerns/rate_limited_service.rb

# 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.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

    def initialize(key:, opts:, rate_limiter:)
      @key = key
      @opts = opts
      @rate_limiter = rate_limiter
    end

    def rate_limit!(service)
      evaluated_scope = evaluated_scope_for(service)

      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

    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
  end

  prepended do
    attr_accessor :rate_limiter_bypassed

    cattr_accessor :rate_limiter_scoped_and_keyed

    def self.rate_limit(key:, opts:, rate_limiter: ::Gitlab::ApplicationRateLimiter)
      self.rate_limiter_scoped_and_keyed = RateLimiterScopedAndKeyed.new(key: key,
                                                                         opts: opts,
                                                                         rate_limiter: rate_limiter)
    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