Allow profiler to authenticate by stubbing users directly54327-profiler-doesn-t-work-with-auth-now

Previously, we used a personal access token. This had a couple of
problems:

1. If the user didn't have a PAT, we couldn't impersonate them.
2. It depended on reading the raw PAT from the database.

Instead, we can monkey-patch the authentication methods on
ApplicationController (overriding the Devise ones), and remove them once
we're done. This does mean that profiles will not profile auth
correctly, so for that, use a PAT directly.

2 files changed, 52 insertions, 29 deletions

diff --git a/lib/gitlab/profiler.rb b/lib/gitlab/profiler.rb
index 4a62f367835..93a9fcf1591 100644
--- a/lib/gitlab/profiler.rb
+++ b/lib/gitlab/profiler.rb
@@ -36,7 +36,6 @@ module Gitlab
     #
     # - private_token: instead of providing a user instance, the token can be
     #   given as a string. Takes precedence over the user option.
-    # rubocop: disable CodeReuse/ActiveRecord
     def self.profile(url, logger: nil, post_data: nil, user: nil, private_token: nil)
       app = ActionDispatch::Integration::Session.new(Rails.application)
       verb = :get
@@ -47,12 +46,11 @@ module Gitlab
         headers['Content-Type'] = 'application/json'
       end
 
-      if user
-        private_token ||= user.personal_access_tokens.active.pluck(:token).first
-        raise 'Your user must have a personal_access_token' unless private_token
+      if private_token
+        headers['Private-Token'] = private_token
+        user = nil # private_token overrides user
       end
 
-      headers['Private-Token'] = private_token if private_token
       logger = create_custom_logger(logger, private_token: private_token)
 
       RequestStore.begin!
@@ -70,7 +68,9 @@ module Gitlab
       app.get('/api/v4/users')
 
       result = with_custom_logger(logger) do
-        RubyProf.profile { app.public_send(verb, url, post_data, headers) } # rubocop:disable GitlabSecurity/PublicSend
+        with_user(user) do
+          RubyProf.profile { app.public_send(verb, url, post_data, headers) } # rubocop:disable GitlabSecurity/PublicSend
+        end
       end
 
       RequestStore.end!
@@ -79,7 +79,6 @@ module Gitlab
 
       result
     end
-    # rubocop: enable CodeReuse/ActiveRecord
 
     def self.create_custom_logger(logger, private_token: nil)
       return unless logger
@@ -130,13 +129,29 @@ module Gitlab
         ActionController::Base.logger = logger
       end
 
-      result = yield
+      yield.tap do
+        ActiveSupport::LogSubscriber.colorize_logging = original_colorize_logging
+        ActiveRecord::Base.logger = original_activerecord_logger
+        ActionController::Base.logger = original_actioncontroller_logger
+      end
+    end
+
+    def self.with_user(user)
+      if user
+        API::Helpers::CommonHelpers.send(:define_method, :find_current_user!) { user } # rubocop:disable GitlabSecurity/PublicSend
+        ApplicationController.send(:define_method, :current_user) { user } # rubocop:disable GitlabSecurity/PublicSend
+        ApplicationController.send(:define_method, :authenticate_user!) { } # rubocop:disable GitlabSecurity/PublicSend
+      end
 
-      ActiveSupport::LogSubscriber.colorize_logging = original_colorize_logging
-      ActiveRecord::Base.logger = original_activerecord_logger
-      ActionController::Base.logger = original_actioncontroller_logger
+      yield.tap do
+        remove_method(API::Helpers::CommonHelpers, :find_current_user!)
+        remove_method(ApplicationController, :current_user)
+        remove_method(ApplicationController, :authenticate_user!)
+      end
+    end
 
-      result
+    def self.remove_method(klass, meth)
+      klass.send(:remove_method, meth) if klass.instance_methods(false).include?(meth) # rubocop:disable GitlabSecurity/PublicSend
     end
 
     # rubocop: disable CodeReuse/ActiveRecord
diff --git a/spec/lib/gitlab/profiler_spec.rb b/spec/lib/gitlab/profiler_spec.rb
index 4059188fba1..8bb0c1a0b8a 100644
--- a/spec/lib/gitlab/profiler_spec.rb
+++ b/spec/lib/gitlab/profiler_spec.rb
@@ -43,31 +43,16 @@ describe Gitlab::Profiler do
 
     it 'uses the user for auth if given' do
       user = double(:user)
-      user_token = 'user'
 
-      allow(user).to receive_message_chain(:personal_access_tokens, :active, :pluck, :first).and_return(user_token)
-
-      expect(app).to receive(:get).with('/', nil, 'Private-Token' => user_token)
-      expect(app).to receive(:get).with('/api/v4/users')
+      expect(described_class).to receive(:with_user).with(user)
 
       described_class.profile('/', user: user)
     end
 
-    context 'when providing a user without a personal access token' do
-      it 'raises an error' do
-        user = double(:user)
-        allow(user).to receive_message_chain(:personal_access_tokens, :active, :pluck).and_return([])
-
-        expect { described_class.profile('/', user: user) }.to raise_error('Your user must have a personal_access_token')
-      end
-    end
-
     it 'uses the private_token for auth if both it and user are set' do
       user = double(:user)
-      user_token = 'user'
-
-      allow(user).to receive_message_chain(:personal_access_tokens, :active, :pluck, :first).and_return(user_token)
 
+      expect(described_class).to receive(:with_user).with(nil).and_call_original
       expect(app).to receive(:get).with('/', nil, 'Private-Token' => private_token)
       expect(app).to receive(:get).with('/api/v4/users')
 
@@ -210,6 +195,29 @@ describe Gitlab::Profiler do
     end
   end
 
+  describe '.with_user' do
+    context 'when the user is set' do
+      let(:user) { double(:user) }
+
+      it 'overrides auth in ApplicationController to use the given user' do
+        expect(described_class.with_user(user) { ApplicationController.new.current_user }).to eq(user)
+      end
+
+      it 'cleans up ApplicationController afterwards' do
+        expect { described_class.with_user(user) { } }
+          .to not_change { ActionController.instance_methods(false) }
+      end
+    end
+
+    context 'when the user is nil' do
+      it 'does not define methods on ApplicationController' do
+        expect(ApplicationController).not_to receive(:define_method)
+
+        described_class.with_user(nil) { }
+      end
+    end
+  end
+
   describe '.log_load_times_by_model' do
     it 'logs the model, query count, and time by slowest first' do
       expect(null_logger).to receive(:load_times_by_model).and_return(