1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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
|
