diff options
| -rw-r--r-- | lib/gitlab/profiler.rb | 39 | ||||
| -rw-r--r-- | spec/lib/gitlab/profiler_spec.rb | 42 |
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( |