Add latest changes from gitlab-org/gitlab@master

1 files changed, 12 insertions, 5 deletions

diff --git a/app/controllers/omniauth_callbacks_controller.rb b/app/controllers/omniauth_callbacks_controller.rb
index bc3308fd6c6..d82a46e57ea 100644
--- a/app/controllers/omniauth_callbacks_controller.rb
+++ b/app/controllers/omniauth_callbacks_controller.rb
@@ -2,6 +2,7 @@
 
 class OmniauthCallbacksController < Devise::OmniauthCallbacksController
   include AuthenticatesWithTwoFactor
+  include Authenticates2FAForAdminMode
   include Devise::Controllers::Rememberable
   include AuthHelper
   include InitializesCurrentUserMode
@@ -97,7 +98,7 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
       log_audit_event(current_user, with: oauth['provider'])
 
       if Feature.enabled?(:user_mode_in_session)
-        return admin_mode_flow if current_user_mode.admin_mode_requested?
+        return admin_mode_flow(auth_module::User) if current_user_mode.admin_mode_requested?
       end
 
       identity_linker ||= auth_module::IdentityLinker.new(current_user, oauth, session)
@@ -245,13 +246,19 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
     end
   end
 
-  def admin_mode_flow
-    if omniauth_identity_matches_current_user?
+  def admin_mode_flow(auth_user_class)
+    auth_user = build_auth_user(auth_user_class)
+
+    return fail_admin_mode_invalid_credentials unless omniauth_identity_matches_current_user?
+
+    if current_user.two_factor_enabled? && !auth_user.bypass_two_factor?
+      admin_mode_prompt_for_two_factor(current_user)
+    else
+      # Can only reach here if the omniauth identity matches current user
+      # and current_user is an admin that requested admin mode
       current_user_mode.enable_admin_mode!(skip_password_validation: true)
 
       redirect_to stored_location_for(:redirect) || admin_root_path, notice: _('Admin mode enabled')
-    else
-      fail_admin_mode_invalid_credentials
     end
   end