Add rate limiting bypass mechanism triggered by HTTP header

changelogs/unreleased/jv-add-rate-limit-bypass.yml

+---
+title: Add rate limit bypass
+merge_request: 46259
+author:
+type: changed

config/initializers/rack_attack.rb

+# frozen_string_literal: true
+
 # Specs for this file can be found on:
 # * spec/lib/gitlab/throttle_spec.rb
 # * spec/requests/rack_attack_global_spec.rb
@@ -15,6 +17,13 @@ module Gitlab::Throttle
     Rack::Attack.throttles.key?('protected paths')
   end
 
+  def self.bypass_header
+    env_value = ENV['GITLAB_THROTTLE_BYPASS_HEADER']
+    return unless env_value.present?
+
+    "HTTP_#{env_value.upcase.tr('-', '_')}"
+  end
+
   def self.unauthenticated_options
     limit_proc = proc { |req| settings.throttle_unauthenticated_requests_per_period }
     period_proc = proc { |req| settings.throttle_unauthenticated_period_in_seconds.seconds }
@@ -112,6 +121,11 @@ class Rack::Attack
     end
   end
 
+  safelist('throttle_bypass_header') do |req|
+    Gitlab::Throttle.bypass_header.present? &&
+      req.get_header(Gitlab::Throttle.bypass_header) == '1'
+  end
+
   class Request
     def unauthenticated?
       !(authenticated_user_id([:api, :rss, :ics]) || authenticated_runner_id)

config/initializers/rack_attack_logging.rb

@@ -5,7 +5,8 @@
 ActiveSupport::Notifications.subscribe(/rack_attack/) do |name, start, finish, request_id, payload|
   req = payload[:request]
 
-  if [:throttle, :blocklist].include? req.env['rack.attack.match_type']
+  case req.env['rack.attack.match_type']
+  when :throttle, :blocklist
     rack_attack_info = {
       message: 'Rack_Attack',
       env: req.env['rack.attack.match_type'],
@@ -31,5 +32,7 @@ ActiveSupport::Notifications.subscribe(/rack_attack/) do |name, start, finish, r
     end
 
     Gitlab::AuthLogger.error(rack_attack_info)
+  when :safelist
+    Gitlab::Instrumentation::Throttle.safelist = req.env['rack.attack.matched']
   end
 end

doc/user/admin_area/settings/user_and_ip_rate_limits.md

@@ -22,6 +22,43 @@ These limits are disabled by default.
 
 ![user-and-ip-rate-limits](img/user_and_ip_rate_limits.png)
 
+## Use an HTTP header to bypass rate limiting
+
+> [Introduced](https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/622) in GitLab 13.6.
+
+Depending on the needs of your organization, you may want to enable rate limiting
+but have some requests bypass the rate limiter.
+
+You can do this by marking requests that should bypass the rate limiter with a custom
+header. You must do this somewhere in a load balancer or reverse proxy in front of
+GitLab. For example:
+
+1. Pick a name for your bypass header. For example, `Gitlab-Bypass-Rate-Limiting`.
+1. Configure your load balancer to set `Gitlab-Bypass-Rate-Limiting: 1` on requests
+   that should bypass GitLab rate limiting.
+1. Configure your load balancer to either:
+   - Erase `Gitlab-Bypass-Rate-Limiting`.
+   - Set `Gitlab-Bypass-Rate-Limiting` to a value other than `1` on all requests that
+     should be affected by rate limiting.
+1. Set the environment variable  `GITLAB_THROTTLE_BYPASS_HEADER`.
+   - For [Omnibus](https://docs.gitlab.com/omnibus/settings/environment-variables.html),
+     set `'GITLAB_THROTTLE_BYPASS_HEADER' => 'Gitlab-Bypass-Rate-Limiting'` in `gitlab_rails['env']`.
+   - For source installations, set `export GITLAB_THROTTLE_BYPASS_HEADER=Gitlab-Bypass-Rate-Limiting`
+     in `/etc/default/gitlab`.
+
+It is important that your load balancer erases or overwrites the bypass
+header on all incoming traffic, because otherwise you must trust your
+users to not set that header and bypass the GitLab rate limiter.
+
+Note that the bypass only works if the header is set to `1`.
+
+Requests that bypassed the rate limiter because of the bypass header
+will be marked with `"throttle_safelist":"throttle_bypass_header"` in
+[`production_json.log`](../../../administration/logs.md#production_jsonlog).
+
+To disable the bypass mechanism, make sure the environment variable
+`GITLAB_THROTTLE_BYPASS_HEADER` is unset or empty.
+
 <!-- ## Troubleshooting
 
 Include any troubleshooting steps that you can foresee. If you know beforehand what issues

lib/gitlab/instrumentation/throttle.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Instrumentation
+    class Throttle
+      KEY = :instrumentation_throttle_safelist
+
+      def self.safelist
+        Gitlab::SafeRequestStore[KEY]
+      end
+
+      def self.safelist=(name)
+        Gitlab::SafeRequestStore[KEY] = name
+      end
+    end
+  end
+end

lib/gitlab/instrumentation_helper.rb

@@ -21,6 +21,7 @@ module Gitlab
       instrument_rugged(payload)
       instrument_redis(payload)
       instrument_elasticsearch(payload)
+      instrument_throttle(payload)
     end
 
     def instrument_gitaly(payload)
@@ -56,6 +57,11 @@ module Gitlab
       payload[:elasticsearch_duration_s] = Gitlab::Instrumentation::ElasticsearchTransport.query_time
     end
 
+    def instrument_throttle(payload)
+      safelist = Gitlab::Instrumentation::Throttle.safelist
+      payload[:throttle_safelist] = safelist if safelist.present?
+    end
+
     # Returns the queuing duration for a Sidekiq job in seconds, as a float, if the
     # `enqueued_at` field or `created_at` field is available.
     #

spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -97,6 +97,16 @@ RSpec.describe Gitlab::InstrumentationHelper do
         expect(payload[:gitaly_duration]).to be_nil
       end
     end
+
+    context 'when the request matched a Rack::Attack safelist' do
+      it 'logs the safelist name' do
+        Gitlab::Instrumentation::Throttle.safelist = 'foobar'
+
+        subject
+
+        expect(payload[:throttle_safelist]).to eq('foobar')
+      end
+    end
   end
 
   describe '.queue_duration_for_job' do

spec/lib/gitlab/throttle_spec.rb

@@ -12,4 +12,22 @@ RSpec.describe Gitlab::Throttle do
       subject
     end
   end
+
+  describe '.bypass_header' do
+    subject { described_class.bypass_header }
+
+    it 'is nil' do
+      expect(subject).to be_nil
+    end
+
+    context 'when a header is configured' do
+      before do
+        stub_env('GITLAB_THROTTLE_BYPASS_HEADER', 'My-Custom-Header')
+      end
+
+      it 'is a funny upper case rack key' do
+        expect(subject).to eq('HTTP_MY_CUSTOM_HEADER')
+      end
+    end
+  end
 end

spec/requests/rack_attack_global_spec.rb

@@ -320,4 +320,62 @@ RSpec.describe 'Rack Attack global throttles' do
       it_behaves_like 'rate-limited web authenticated requests'
     end
   end
+
+  describe 'throttle bypass header' do
+    let(:headers) { {} }
+    let(:bypass_header) { 'gitlab-bypass-rate-limiting' }
+
+    def do_request
+      get '/users/sign_in', headers: headers
+    end
+
+    before do
+      # Disabling protected paths throttle, otherwise requests to
+      # '/users/sign_in' are caught by this throttle.
+      settings_to_set[:throttle_protected_paths_enabled] = false
+
+      # Set low limits
+      settings_to_set[:throttle_unauthenticated_requests_per_period] = requests_per_period
+      settings_to_set[:throttle_unauthenticated_period_in_seconds] = period_in_seconds
+
+      stub_env('GITLAB_THROTTLE_BYPASS_HEADER', bypass_header)
+      settings_to_set[:throttle_unauthenticated_enabled] = true
+
+      stub_application_setting(settings_to_set)
+    end
+
+    shared_examples 'reject requests over the rate limit' do
+      it 'rejects requests over the rate limit' do
+        # At first, allow requests under the rate limit.
+        requests_per_period.times do
+          do_request
+          expect(response).to have_gitlab_http_status(:ok)
+        end
+
+        # the last straw
+        expect_rejection { do_request }
+      end
+    end
+
+    context 'without the bypass header set' do
+      it_behaves_like 'reject requests over the rate limit'
+    end
+
+    context 'with bypass header set to 1' do
+      let(:headers) { { bypass_header => '1' } }
+
+      it 'does not throttle' do
+        (1 + requests_per_period).times do
+          do_request
+          expect(response).to have_gitlab_http_status(:ok)
+        end
+      end
+    end
+
+    context 'with bypass header set to some other value' do
+      let(:headers) { { bypass_header => 'some other value' } }
+
+      it_behaves_like 'reject requests over the rate limit'
+    end
+  end
 end