summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--lib/gitlab/profiler.rb39
-rw-r--r--spec/lib/gitlab/profiler_spec.rb42
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(