Perform gitlab-ci-token authentication always using primary

ee/changelogs/unreleased/perform-ci-build-auth-always-on-primary.yml

+---
+title: Perform gitlab-ci-token authentication always using primary
+merge_request:
+author:
+type: fixed

ee/lib/ee/gitlab/auth.rb

@@ -20,6 +20,7 @@ module EE
         end
       end
 
+      override :find_with_user_password
       def find_with_user_password(login, password)
         if Devise.omniauth_providers.include?(:kerberos)
           kerberos_user = ::Gitlab::Kerberos::Authentication.login(login, password)
@@ -28,6 +29,13 @@ module EE
 
         super
       end
+
+      override :find_build_by_token
+      def find_build_by_token(token)
+        ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+          super
+        end
+      end
     end
   end
 end

ee/lib/gitlab/database/load_balancing/session.rb

@@ -30,6 +30,14 @@ module Gitlab
           @use_primary = true
         end
 
+        def use_primary(&blk)
+          used_primary = @use_primary
+          @use_primary = true
+          return yield
+        ensure
+          @use_primary = used_primary || @performed_write
+        end
+
         def write!
           @performed_write = true
           use_primary!

ee/spec/lib/gitlab/auth_spec.rb

@@ -22,4 +22,24 @@ describe Gitlab::Auth do
       expect( gl_auth.find_with_user_password(username, password) ).to eql user
     end
   end
+
+  describe '#build_access_token_check' do
+    subject { gl_auth.find_for_git_client('gitlab-ci-token', build.token, project: build.project, ip: '1.2.3.4') }
+
+    context 'for running build' do
+      let!(:build) { create(:ci_build, :running, user: user) }
+
+      it 'executes query using primary database' do
+        expect(Ci::Build).to receive(:find_by_token).with(build.token).and_wrap_original do |m, *args|
+          expect(::Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
+          m.call(*args)
+        end
+
+        expect(subject).to be_a(Gitlab::Auth::Result)
+        expect(subject.actor).to eq(user)
+        expect(subject.project).to eq(build.project)
+        expect(subject.type).to eq(:build)
+      end
+    end
+  end
 end

ee/spec/lib/gitlab/database/load_balancing/session_spec.rb

@@ -42,6 +42,49 @@ describe Gitlab::Database::LoadBalancing::Session do
     end
   end
 
+  describe '#use_primary' do
+    let(:instance) { described_class.new }
+
+    context 'when primary was used before' do
+      before do
+        instance.write!
+      end
+
+      it 'restores state after use' do
+        expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args
+
+        expect(instance.use_primary?).to eq(true)
+      end
+    end
+
+    context 'when primary was not used' do
+      it 'restores state after use' do
+        expect { |blk| instance.use_primary(&blk) }.to yield_with_no_args
+
+        expect(instance.use_primary?).to eq(false)
+      end
+    end
+
+    it 'uses primary during block' do
+      expect do |blk|
+        instance.use_primary do
+          expect(instance.use_primary?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+    end
+
+    it 'continues using primary when write was performed' do
+      instance.use_primary do
+        instance.write!
+      end
+
+      expect(instance.use_primary?).to eq(true)
+    end
+  end
+
   describe '#performed_write?' do
     it 'returns true if a write was performed' do
       instance = described_class.new

lib/gitlab/auth.rb

@@ -242,7 +242,7 @@ module Gitlab
         return unless login == 'gitlab-ci-token'
         return unless password
 
-        build = ::Ci::Build.running.find_by_token(password)
+        build = find_build_by_token(password)
         return unless build
         return unless build.project.builds_enabled?
 
@@ -303,6 +303,12 @@ module Gitlab
 
         REGISTRY_SCOPES
       end
+
+      private
+
+      def find_build_by_token(token)
+        ::Ci::Build.running.find_by_token(token)
+      end
     end
   end
 end