diff options
Diffstat (limited to 'lib/gitlab/application_rate_limiter.rb')
-rw-r--r-- | lib/gitlab/application_rate_limiter.rb | 34 |
1 files changed, 20 insertions, 14 deletions
diff --git a/lib/gitlab/application_rate_limiter.rb b/lib/gitlab/application_rate_limiter.rb index 7c37f67b766..3db2f1295f9 100644 --- a/lib/gitlab/application_rate_limiter.rb +++ b/lib/gitlab/application_rate_limiter.rb @@ -11,6 +11,8 @@ module Gitlab # redirect_to(edit_project_path(@project), status: :too_many_requests) # end class ApplicationRateLimiter + InvalidKeyError = Class.new(StandardError) + def initialize(key, **options) @key = key @options = options @@ -64,39 +66,43 @@ module Gitlab # @param key [Symbol] Key attribute registered in `.rate_limits` # @option scope [Array<ActiveRecord>] Array of ActiveRecord models to scope throttling to a specific request (e.g. per user per project) # @option threshold [Integer] Optional threshold value to override default one registered in `.rate_limits` - # @option interval [Integer] Optional interval value to override default one registered in `.rate_limits` # @option users_allowlist [Array<String>] Optional list of usernames to exclude from the limit. This param will only be functional if Scope includes a current user. # # @return [Boolean] Whether or not a request should be throttled def throttled?(key, **options) - return unless rate_limits[key] + raise InvalidKeyError unless rate_limits[key] return if scoped_user_in_allowlist?(options) threshold_value = options[:threshold] || threshold(key) threshold_value > 0 && - increment(key, options[:scope], options[:interval]) > threshold_value + increment(key, options[:scope]) > threshold_value end - # Increments the given cache key and increments the value by 1 with the - # expiration interval defined in `.rate_limits`. + # Increments a cache key that is based on the current time and interval. + # So that when time passes to the next interval, the key changes and the count starts again from 0. + # + # Based on https://github.com/rack/rack-attack/blob/886ba3a18d13c6484cd511a4dc9b76c0d14e5e96/lib/rack/attack/cache.rb#L63-L68 # # @param key [Symbol] Key attribute registered in `.rate_limits` # @option scope [Array<ActiveRecord>] Array of ActiveRecord models to scope throttling to a specific request (e.g. per user per project) - # @option interval [Integer] Optional interval value to override default one registered in `.rate_limits` # # @return [Integer] incremented value - def increment(key, scope, interval = nil) - value = 0 - interval_value = interval || interval(key) + def increment(key, scope) + interval_value = interval(key) + + period_key, time_elapsed_in_period = Time.now.to_i.divmod(interval_value) + + cache_key = "#{action_key(key, scope)}:#{period_key}" + # We add a 1 second buffer to avoid timing issues when we're at the end of a period + expiry = interval_value - time_elapsed_in_period + 1 ::Gitlab::Redis::RateLimiting.with do |redis| - cache_key = action_key(key, scope) - value = redis.incr(cache_key) - redis.expire(cache_key, interval_value) if value == 1 + redis.pipelined do + redis.incr(cache_key) + redis.expire(cache_key, expiry) + end.first end - - value end # Logs request using provided logger |