From f10c999bca2b5b37b068ff3680a6e35a6707828d Mon Sep 17 00:00:00 2001
From: James Edwards-Jones <jedwardsjones@gitlab.com>
Date: Wed, 18 Apr 2018 15:03:27 +0100
Subject: [PATCH] Refactor OmniauthCallbacksController to remove duplication

Moves LDAP to its own controller with tests
Provides path forward for implementing GroupSaml
---
 .../concerns/authenticates_with_two_factor.rb |   3 +
 .../ldap/omniauth_callbacks_controller.rb     |  33 +++++
 .../omniauth_callbacks_controller.rb          | 134 +++++++-----------
 config/routes/user.rb                         |  18 +++
 .../ee/ldap/omniauth_callbacks_controller.rb  |  22 +++
 .../omniauth_callbacks_controller_spec.rb     |  29 ++++
 lib/gitlab/auth/ldap/user.rb                  |  11 ++
 lib/gitlab/auth/o_auth/identity_linker.rb     |  15 ++
 lib/gitlab/auth/o_auth/user.rb                |  14 ++
 .../auth/omniauth_identity_linker_base.rb     |  21 +++
 lib/gitlab/auth/saml/identity_linker.rb       |  27 ++++
 lib/gitlab/auth/saml/user.rb                  |  15 +-
 .../omniauth_callbacks_controller_spec.rb     |  58 ++++++++
 spec/features/oauth_login_spec.rb             |  49 ++++---
 .../auth/o_auth/identity_linker_spec.rb       |  42 ++++++
 .../gitlab/auth/saml/identity_linker_spec.rb  |  48 +++++++
 ...uth_callbacks_controller_shared_context.rb |  33 +++++
 spec/support/ldap_helpers.rb                  |   4 +
 spec/support/login_helpers.rb                 |   4 +-
 19 files changed, 474 insertions(+), 106 deletions(-)
 create mode 100644 app/controllers/ldap/omniauth_callbacks_controller.rb
 create mode 100644 ee/app/controllers/ee/ldap/omniauth_callbacks_controller.rb
 create mode 100644 ee/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
 create mode 100644 lib/gitlab/auth/o_auth/identity_linker.rb
 create mode 100644 lib/gitlab/auth/omniauth_identity_linker_base.rb
 create mode 100644 lib/gitlab/auth/saml/identity_linker.rb
 create mode 100644 spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
 create mode 100644 spec/lib/gitlab/auth/o_auth/identity_linker_spec.rb
 create mode 100644 spec/lib/gitlab/auth/saml/identity_linker_spec.rb
 create mode 100644 spec/support/controllers/ldap_omniauth_callbacks_controller_shared_context.rb

diff --git a/app/controllers/concerns/authenticates_with_two_factor.rb b/app/controllers/concerns/authenticates_with_two_factor.rb
index 2fdf346ef44..69a053d4246 100644
--- a/app/controllers/concerns/authenticates_with_two_factor.rb
+++ b/app/controllers/concerns/authenticates_with_two_factor.rb
@@ -23,6 +23,9 @@ module AuthenticatesWithTwoFactor
   #
   # Returns nil
   def prompt_for_two_factor(user)
+    # Set @user for Devise views
+    @user = user # rubocop:disable Gitlab/ModuleWithInstanceVariables
+
     return locked_user_redirect(user) unless user.can?(:log_in)
 
     session[:otp_user_id] = user.id
diff --git a/app/controllers/ldap/omniauth_callbacks_controller.rb b/app/controllers/ldap/omniauth_callbacks_controller.rb
new file mode 100644
index 00000000000..e9219274182
--- /dev/null
+++ b/app/controllers/ldap/omniauth_callbacks_controller.rb
@@ -0,0 +1,33 @@
+class Ldap::OmniauthCallbacksController < OmniauthCallbacksController
+  extend ::Gitlab::Utils::Override
+
+  def self.define_providers!
+    if Gitlab::Auth::LDAP::Config.enabled?
+      Gitlab::Auth::LDAP::Config.available_servers.each do |server|
+        define_method server['provider_name'] do
+          ldap
+        end
+      end
+    end
+  end
+
+  define_providers!
+
+  # We only find ourselves here
+  # if the authentication to LDAP was successful.
+  def ldap
+    sign_in_user_flow(Gitlab::Auth::LDAP::User)
+  end
+
+  override :set_remember_me
+  def set_remember_me(user)
+    user.remember_me = params[:remember_me] if user.persisted?
+  end
+
+  override :fail_login
+  def fail_login(user)
+    flash[:alert] = 'Access denied for your LDAP account.'
+
+    redirect_to new_user_session_path
+  end
+end
diff --git a/app/controllers/omniauth_callbacks_controller.rb b/app/controllers/omniauth_callbacks_controller.rb
index 5e6676ea513..190ad55859b 100644
--- a/app/controllers/omniauth_callbacks_controller.rb
+++ b/app/controllers/omniauth_callbacks_controller.rb
@@ -10,14 +10,6 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
     end
   end
 
-  if Gitlab::Auth::LDAP::Config.enabled?
-    Gitlab::Auth::LDAP::Config.available_servers.each do |server|
-      define_method server['provider_name'] do
-        ldap
-      end
-    end
-  end
-
   # Extend the standard implementation to also increment
   # the number of failed sign in attempts
   def failure
@@ -37,51 +29,12 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
     error ||= exception.error        if exception.respond_to?(:error)
     error ||= exception.message      if exception.respond_to?(:message)
     error ||= env["omniauth.error.type"].to_s
-    error.to_s.humanize if error
-  end
-
-  # We only find ourselves here
-  # if the authentication to LDAP was successful.
-  def ldap
-    ldap_user = Gitlab::Auth::LDAP::User.new(oauth)
-    ldap_user.save if ldap_user.changed? # will also save new users
 
-    @user = ldap_user.gl_user
-    @user.remember_me = params[:remember_me] if ldap_user.persisted?
-
-    # Do additional LDAP checks for the user filter and EE features
-    if ldap_user.allowed?
-      if @user.two_factor_enabled?
-        prompt_for_two_factor(@user)
-      else
-        log_audit_event(@user, with: oauth['provider'])
-        sign_in_and_redirect(@user)
-      end
-    else
-      fail_ldap_login
-    end
+    error.to_s.humanize if error
   end
 
   def saml
-    if current_user
-      log_audit_event(current_user, with: :saml)
-      # Update SAML identity if data has changed.
-      identity = current_user.identities.with_extern_uid(:saml, oauth['uid']).take
-      if identity.nil?
-        current_user.identities.create(extern_uid: oauth['uid'], provider: :saml)
-        redirect_to profile_account_path, notice: 'Authentication method updated'
-      else
-        redirect_to after_sign_in_path_for(current_user)
-      end
-    else
-      saml_user = Gitlab::Auth::Saml::User.new(oauth)
-      saml_user.save if saml_user.changed?
-      @user = saml_user.gl_user
-
-      continue_login_process
-    end
-  rescue Gitlab::Auth::OAuth::User::SignupDisabledError
-    handle_signup_error
+    omniauth_flow(Gitlab::Auth::Saml)
   end
 
   def omniauth_error
@@ -118,24 +71,33 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
   private
 
   def handle_omniauth
+    omniauth_flow(Gitlab::Auth::OAuth)
+  end
+
+  def omniauth_flow(auth_module, identity_linker: nil)
     if current_user
-      # Add new authentication method
-      current_user.identities
-                  .with_extern_uid(oauth['provider'], oauth['uid'])
-                  .first_or_create(extern_uid: oauth['uid'])
       log_audit_event(current_user, with: oauth['provider'])
-      redirect_to profile_account_path, notice: 'Authentication method updated'
-    else
-      oauth_user = Gitlab::Auth::OAuth::User.new(oauth)
-      oauth_user.save
-      @user = oauth_user.gl_user
 
-      continue_login_process
+      identity_linker ||= auth_module::IdentityLinker.new(current_user, oauth)
+
+      identity_linker.create_or_update
+
+      if identity_linker.created?
+        redirect_identity_linked
+      else
+        redirect_identity_exists
+      end
+    else
+      sign_in_user_flow(auth_module::User)
     end
-  rescue Gitlab::Auth::OAuth::User::SigninDisabledForProviderError
-    handle_disabled_provider
-  rescue Gitlab::Auth::OAuth::User::SignupDisabledError
-    handle_signup_error
+  end
+
+  def redirect_identity_exists
+    redirect_to after_sign_in_path_for(current_user)
+  end
+
+  def redirect_identity_linked
+    redirect_to profile_account_path, notice: 'Authentication method updated'
   end
 
   def handle_service_ticket(provider, ticket)
@@ -144,21 +106,27 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
     session[:service_tickets][provider] = ticket
   end
 
-  def continue_login_process
-    # Only allow properly saved users to login.
-    if @user.persisted? && @user.valid?
-      log_audit_event(@user, with: oauth['provider'])
+  def sign_in_user_flow(auth_user_class)
+    auth_user = auth_user_class.new(oauth)
+    user = auth_user.find_and_update!
+
+    if auth_user.valid_sign_in?
+      log_audit_event(user, with: oauth['provider'])
+
+      set_remember_me(user)
 
-      if @user.two_factor_enabled?
-        params[:remember_me] = '1' if remember_me?
-        prompt_for_two_factor(@user)
+      if user.two_factor_enabled?
+        prompt_for_two_factor(user)
       else
-        remember_me(@user) if remember_me?
-        sign_in_and_redirect(@user)
+        sign_in_and_redirect(user)
       end
     else
-      fail_login
+      fail_login(user)
     end
+  rescue Gitlab::Auth::OAuth::User::SigninDisabledForProviderError
+    handle_disabled_provider
+  rescue Gitlab::Auth::OAuth::User::SignupDisabledError
+    handle_signup_error
   end
 
   def handle_signup_error
@@ -178,18 +146,12 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
     @oauth ||= request.env['omniauth.auth']
   end
 
-  def fail_login
-    error_message = @user.errors.full_messages.to_sentence
+  def fail_login(user)
+    error_message = user.errors.full_messages.to_sentence
 
     return redirect_to omniauth_error_path(oauth['provider'], error: error_message)
   end
 
-  def fail_ldap_login
-    flash[:alert] = 'Access denied for your LDAP account.'
-
-    redirect_to new_user_session_path
-  end
-
   def fail_auth0_login
     flash[:alert] = 'Wrong extern UID provided. Make sure Auth0 is configured correctly.'
 
@@ -208,6 +170,16 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
       .for_authentication.security_event
   end
 
+  def set_remember_me(user)
+    return unless remember_me?
+
+    if user.two_factor_enabled?
+      params[:remember_me] = '1'
+    else
+      remember_me(user)
+    end
+  end
+
   def remember_me?
     request_params = request.env['omniauth.params']
     (request_params['remember_me'] == '1') if request_params.present?
diff --git a/config/routes/user.rb b/config/routes/user.rb
index 57fb37530bb..f8677693fab 100644
--- a/config/routes/user.rb
+++ b/config/routes/user.rb
@@ -1,3 +1,21 @@
+# Allows individual providers to be directed to a chosen controller
+# Call from inside devise_scope
+def override_omniauth(provider, controller, path_prefix = '/users/auth')
+  match "#{path_prefix}/#{provider}/callback",
+    to: "#{controller}##{provider}",
+    as: "#{provider}_omniauth_callback",
+    via: [:get, :post]
+end
+
+# Use custom controller for LDAP omniauth callback
+if Gitlab::Auth::LDAP::Config.enabled?
+  devise_scope :user do
+    Gitlab::Auth::LDAP::Config.available_servers.each do |server|
+      override_omniauth(server['provider_name'], 'ldap/omniauth_callbacks')
+    end
+  end
+end
+
 devise_for :users, controllers: { omniauth_callbacks: :omniauth_callbacks,
                                   registrations: :registrations,
                                   passwords: :passwords,
diff --git a/ee/app/controllers/ee/ldap/omniauth_callbacks_controller.rb b/ee/app/controllers/ee/ldap/omniauth_callbacks_controller.rb
new file mode 100644
index 00000000000..f1e851a210b
--- /dev/null
+++ b/ee/app/controllers/ee/ldap/omniauth_callbacks_controller.rb
@@ -0,0 +1,22 @@
+module EE
+  module Ldap
+    module OmniauthCallbacksController
+      extend ::Gitlab::Utils::Override
+
+      override :sign_in_and_redirect
+      def sign_in_and_redirect(user)
+        # The counter gets incremented in `sign_in_and_redirect`
+        show_ldap_sync_flash if user.sign_in_count == 0
+
+        super
+      end
+
+      private
+
+      def show_ldap_sync_flash
+        flash[:notice] = 'LDAP sync in progress. This could take a few minutes. '\
+                         'Refresh the page to see the changes.'
+      end
+    end
+  end
+end
diff --git a/ee/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb b/ee/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
new file mode 100644
index 00000000000..0835ff35846
--- /dev/null
+++ b/ee/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
@@ -0,0 +1,29 @@
+require 'spec_helper'
+
+describe Ldap::OmniauthCallbacksController do
+  include_context 'Ldap::OmniauthCallbacksController'
+
+  it "displays LDAP sync flash on first sign in" do
+    post provider
+
+    expect(flash[:notice]).to match(/LDAP sync in progress*/)
+  end
+
+  it "skips LDAP sync flash on subsequent sign ins" do
+    user.update!(sign_in_count: 1)
+
+    post provider
+
+    expect(flash[:notice]).to eq nil
+  end
+
+  context 'access denied' do
+    let(:valid_login?) { false }
+
+    it 'logs a failure event' do
+      stub_licensed_features(extended_audit_events: true)
+
+      expect { post provider }.to change(SecurityEvent, :count).by(1)
+    end
+  end
+end
diff --git a/lib/gitlab/auth/ldap/user.rb b/lib/gitlab/auth/ldap/user.rb
index 068212d9a21..604c2d222e9 100644
--- a/lib/gitlab/auth/ldap/user.rb
+++ b/lib/gitlab/auth/ldap/user.rb
@@ -8,6 +8,8 @@ module Gitlab
   module Auth
     module LDAP
       class User < Gitlab::Auth::OAuth::User
+        extend ::Gitlab::Utils::Override
+
         class << self
           def find_by_uid_and_provider(uid, provider)
             identity = ::Identity.with_extern_uid(provider, uid).take
@@ -33,6 +35,11 @@ module Gitlab
           gl_user.changed? || gl_user.identities.any?(&:changed?)
         end
 
+        override :omniauth_should_save?
+        def omniauth_should_save?
+          changed? && super
+        end
+
         def block_after_signup?
           ldap_config.block_auto_created_users
         end
@@ -41,6 +48,10 @@ module Gitlab
           Gitlab::Auth::LDAP::Access.allowed?(gl_user)
         end
 
+        def valid_sign_in?
+          allowed?
+        end
+
         def ldap_config
           Gitlab::Auth::LDAP::Config.new(auth_hash.provider)
         end
diff --git a/lib/gitlab/auth/o_auth/identity_linker.rb b/lib/gitlab/auth/o_auth/identity_linker.rb
new file mode 100644
index 00000000000..cfa83ba2a55
--- /dev/null
+++ b/lib/gitlab/auth/o_auth/identity_linker.rb
@@ -0,0 +1,15 @@
+module Gitlab
+  module Auth
+    module OAuth
+      class IdentityLinker < OmniauthIdentityLinkerBase
+        def create_or_update
+          current_user.identities
+                      .with_extern_uid(oauth['provider'], oauth['uid'])
+                      .first_or_create(extern_uid: oauth['uid'])
+
+          @created = true
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/auth/o_auth/user.rb b/lib/gitlab/auth/o_auth/user.rb
index d0c6b0386ba..f2923c29163 100644
--- a/lib/gitlab/auth/o_auth/user.rb
+++ b/lib/gitlab/auth/o_auth/user.rb
@@ -30,6 +30,10 @@ module Gitlab
           gl_user.try(:valid?)
         end
 
+        def valid_sign_in?
+          valid? && persisted?
+        end
+
         def save(provider = 'OAuth')
           raise SigninDisabledForProviderError if oauth_provider_disabled?
           raise SignupDisabledError unless gl_user
@@ -64,8 +68,18 @@ module Gitlab
           user
         end
 
+        def find_and_update!
+          save if omniauth_should_save?
+
+          gl_user
+        end
+
         protected
 
+        def omniauth_should_save?
+          true
+        end
+
         def add_or_update_user_identities
           return unless gl_user
 
diff --git a/lib/gitlab/auth/omniauth_identity_linker_base.rb b/lib/gitlab/auth/omniauth_identity_linker_base.rb
new file mode 100644
index 00000000000..c60d9f70a99
--- /dev/null
+++ b/lib/gitlab/auth/omniauth_identity_linker_base.rb
@@ -0,0 +1,21 @@
+module Gitlab
+  module Auth
+    class OmniauthIdentityLinkerBase
+      attr_reader :current_user, :oauth
+
+      def initialize(current_user, oauth)
+        @current_user = current_user
+        @oauth = oauth
+        @created = false
+      end
+
+      def created?
+        @created
+      end
+
+      def create_or_update
+        raise NotImplementedError
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/auth/saml/identity_linker.rb b/lib/gitlab/auth/saml/identity_linker.rb
new file mode 100644
index 00000000000..d5f97f01df3
--- /dev/null
+++ b/lib/gitlab/auth/saml/identity_linker.rb
@@ -0,0 +1,27 @@
+module Gitlab
+  module Auth
+    module Saml
+      class IdentityLinker < OmniauthIdentityLinkerBase
+        def create_or_update
+          if find_saml_identity.nil?
+            create_saml_identity
+
+            @created = true
+          else
+            @created = false
+          end
+        end
+
+        protected
+
+        def find_saml_identity
+          current_user.identities.with_extern_uid(:saml, oauth['uid']).take
+        end
+
+        def create_saml_identity
+          current_user.identities.create(extern_uid: oauth['uid'], provider: :saml)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/auth/saml/user.rb b/lib/gitlab/auth/saml/user.rb
index d4024e9ec39..557e6aa21a4 100644
--- a/lib/gitlab/auth/saml/user.rb
+++ b/lib/gitlab/auth/saml/user.rb
@@ -7,6 +7,8 @@ module Gitlab
   module Auth
     module Saml
       class User < Gitlab::Auth::OAuth::User
+        extend ::Gitlab::Utils::Override
+
         def save
           super('SAML')
         end
@@ -21,7 +23,7 @@ module Gitlab
           if external_users_enabled? && user
             # Check if there is overlap between the user's groups and the external groups
             # setting then set user as external or internal.
-            user.external = !(auth_hash.groups & Gitlab::Auth::Saml::Config.external_groups).empty?
+            user.external = !(auth_hash.groups & saml_config.external_groups).empty?
           end
 
           user
@@ -33,14 +35,23 @@ module Gitlab
           gl_user.changed? || gl_user.identities.any?(&:changed?)
         end
 
+        override :omniauth_should_save?
+        def omniauth_should_save?
+          changed? && super
+        end
+
         protected
 
+        def saml_config
+          Gitlab::Auth::Saml::Config
+        end
+
         def auto_link_saml_user?
           Gitlab.config.omniauth.auto_link_saml_user
         end
 
         def external_users_enabled?
-          !Gitlab::Auth::Saml::Config.external_groups.nil?
+          !saml_config.external_groups.nil?
         end
 
         def auth_hash=(auth_hash)
diff --git a/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb b/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
new file mode 100644
index 00000000000..87c10a86cdd
--- /dev/null
+++ b/spec/controllers/ldap/omniauth_callbacks_controller_spec.rb
@@ -0,0 +1,58 @@
+require 'spec_helper'
+
+describe Ldap::OmniauthCallbacksController do
+  include_context 'Ldap::OmniauthCallbacksController'
+
+  it 'allows sign in' do
+    post provider
+
+    expect(request.env['warden']).to be_authenticated
+  end
+
+  it 'respects remember me checkbox' do
+    expect do
+      post provider, remember_me: '1'
+    end.to change { user.reload.remember_created_at }.from(nil)
+  end
+
+  context 'with 2FA' do
+    let(:user) { create(:omniauth_user, :two_factor_via_otp, extern_uid: uid, provider: provider) }
+
+    it 'passes remember_me to the Devise view' do
+      post provider, remember_me: '1'
+
+      expect(assigns[:user].remember_me).to eq '1'
+    end
+  end
+
+  context 'access denied' do
+    let(:valid_login?) { false }
+
+    it 'warns the user' do
+      post provider
+
+      expect(flash[:alert]).to match(/Access denied for your LDAP account*/)
+    end
+
+    it "doesn't authenticate user" do
+      post provider
+
+      expect(request.env['warden']).not_to be_authenticated
+      expect(response).to redirect_to(new_user_session_path)
+    end
+  end
+
+  context 'sign up' do
+    let(:user) { double(email: 'new@example.com') }
+
+    before do
+      stub_omniauth_setting(block_auto_created_users: false)
+    end
+
+    it 'is allowed' do
+      post provider
+
+      expect(request.env['warden']).to be_authenticated
+    end
+  end
+end
diff --git a/spec/features/oauth_login_spec.rb b/spec/features/oauth_login_spec.rb
index a5e325ee2e3..013cdaa6479 100644
--- a/spec/features/oauth_login_spec.rb
+++ b/spec/features/oauth_login_spec.rb
@@ -28,35 +28,46 @@ feature 'OAuth Login', :js, :allow_forgery_protection do
     OmniAuth.config.full_host = @omniauth_config_full_host
   end
 
+  def login_with_provider(provider, enter_two_factor: false)
+    login_via(provider.to_s, user, uid, remember_me: remember_me)
+    enter_code(user.current_otp) if enter_two_factor
+  end
+
   providers.each do |provider|
     context "when the user logs in using the #{provider} provider" do
+      let(:uid) { 'my-uid' }
+      let(:remember_me) { false }
+      let(:user) { create(:omniauth_user, extern_uid: uid, provider: provider.to_s) }
+      let(:two_factor_user) { create(:omniauth_user, :two_factor, extern_uid: uid, provider: provider.to_s) }
+
+      before do
+        stub_omniauth_config(provider)
+      end
+
       context 'when two-factor authentication is disabled' do
         it 'logs the user in' do
-          stub_omniauth_config(provider)
-          user = create(:omniauth_user, extern_uid: 'my-uid', provider: provider.to_s)
-          login_via(provider.to_s, user, 'my-uid')
+          login_with_provider(provider)
 
           expect(current_path).to eq root_path
         end
       end
 
       context 'when two-factor authentication is enabled' do
+        let(:user) { two_factor_user }
+
         it 'logs the user in' do
-          stub_omniauth_config(provider)
-          user = create(:omniauth_user, :two_factor, extern_uid: 'my-uid', provider: provider.to_s)
-          login_via(provider.to_s, user, 'my-uid')
+          login_with_provider(provider, enter_two_factor: true)
 
-          enter_code(user.current_otp)
           expect(current_path).to eq root_path
         end
       end
 
       context 'when "remember me" is checked' do
+        let(:remember_me) { true }
+
         context 'when two-factor authentication is disabled' do
           it 'remembers the user after a browser restart' do
-            stub_omniauth_config(provider)
-            user = create(:omniauth_user, extern_uid: 'my-uid', provider: provider.to_s)
-            login_via(provider.to_s, user, 'my-uid', remember_me: true)
+            login_with_provider(provider)
 
             clear_browser_session
 
@@ -66,11 +77,10 @@ feature 'OAuth Login', :js, :allow_forgery_protection do
         end
 
         context 'when two-factor authentication is enabled' do
+          let(:user) { two_factor_user }
+
           it 'remembers the user after a browser restart' do
-            stub_omniauth_config(provider)
-            user = create(:omniauth_user, :two_factor, extern_uid: 'my-uid', provider: provider.to_s)
-            login_via(provider.to_s, user, 'my-uid', remember_me: true)
-            enter_code(user.current_otp)
+            login_with_provider(provider, enter_two_factor: true)
 
             clear_browser_session
 
@@ -83,9 +93,7 @@ feature 'OAuth Login', :js, :allow_forgery_protection do
       context 'when "remember me" is not checked' do
         context 'when two-factor authentication is disabled' do
           it 'does not remember the user after a browser restart' do
-            stub_omniauth_config(provider)
-            user = create(:omniauth_user, extern_uid: 'my-uid', provider: provider.to_s)
-            login_via(provider.to_s, user, 'my-uid', remember_me: false)
+            login_with_provider(provider)
 
             clear_browser_session
 
@@ -95,11 +103,10 @@ feature 'OAuth Login', :js, :allow_forgery_protection do
         end
 
         context 'when two-factor authentication is enabled' do
+          let(:user) { two_factor_user }
+
           it 'does not remember the user after a browser restart' do
-            stub_omniauth_config(provider)
-            user = create(:omniauth_user, :two_factor, extern_uid: 'my-uid', provider: provider.to_s)
-            login_via(provider.to_s, user, 'my-uid', remember_me: false)
-            enter_code(user.current_otp)
+            login_with_provider(provider, enter_two_factor: true)
 
             clear_browser_session
 
diff --git a/spec/lib/gitlab/auth/o_auth/identity_linker_spec.rb b/spec/lib/gitlab/auth/o_auth/identity_linker_spec.rb
new file mode 100644
index 00000000000..dc6aa5de53a
--- /dev/null
+++ b/spec/lib/gitlab/auth/o_auth/identity_linker_spec.rb
@@ -0,0 +1,42 @@
+require 'spec_helper'
+
+describe Gitlab::Auth::OAuth::IdentityLinker do
+  let(:user) { create(:user) }
+  let(:provider) { 'twitter' }
+  let(:uid) { user.email }
+  let(:oauth) { { 'provider' => provider, 'uid' => uid } }
+
+  subject { described_class.new(user, oauth) }
+
+  context 'linked identity exists' do
+    let!(:identity) { user.identities.create!(provider: provider, extern_uid: uid) }
+
+    it "doesn't create new identity" do
+      expect { subject.create_or_update }.not_to change { Identity.count }
+    end
+  end
+
+  context 'identity needs to be created' do
+    it 'creates linked identity' do
+      expect { subject.create_or_update }.to change { user.identities.count }
+    end
+
+    it 'sets identity provider' do
+      subject.create_or_update
+
+      expect(user.identities.last.provider).to eq provider
+    end
+
+    it 'sets identity extern_uid' do
+      subject.create_or_update
+
+      expect(user.identities.last.extern_uid).to eq uid
+    end
+
+    it 'sets #created? to true' do
+      subject.create_or_update
+
+      expect(subject).to be_created
+    end
+  end
+end
diff --git a/spec/lib/gitlab/auth/saml/identity_linker_spec.rb b/spec/lib/gitlab/auth/saml/identity_linker_spec.rb
new file mode 100644
index 00000000000..11e8469e9f4
--- /dev/null
+++ b/spec/lib/gitlab/auth/saml/identity_linker_spec.rb
@@ -0,0 +1,48 @@
+require 'spec_helper'
+
+describe Gitlab::Auth::Saml::IdentityLinker do
+  let(:user) { create(:user) }
+  let(:provider) { 'saml' }
+  let(:uid) { user.email }
+  let(:oauth) { { 'provider' => provider, 'uid' => uid } }
+
+  subject { described_class.new(user, oauth) }
+
+  context 'linked identity exists' do
+    let!(:identity) { user.identities.create!(provider: provider, extern_uid: uid) }
+
+    it "doesn't create new identity" do
+      expect { subject.create_or_update }.not_to change { Identity.count }
+    end
+
+    it 'sets #created? to false' do
+      subject.create_or_update
+
+      expect(subject).not_to be_created
+    end
+  end
+
+  context 'identity needs to be created' do
+    it 'creates linked identity' do
+      expect { subject.create_or_update }.to change { user.identities.count }
+    end
+
+    it 'sets identity provider' do
+      subject.create_or_update
+
+      expect(user.identities.last.provider).to eq provider
+    end
+
+    it 'sets identity extern_uid' do
+      subject.create_or_update
+
+      expect(user.identities.last.extern_uid).to eq uid
+    end
+
+    it 'sets #created? to true' do
+      subject.create_or_update
+
+      expect(subject).to be_created
+    end
+  end
+end
diff --git a/spec/support/controllers/ldap_omniauth_callbacks_controller_shared_context.rb b/spec/support/controllers/ldap_omniauth_callbacks_controller_shared_context.rb
new file mode 100644
index 00000000000..72912ffb89d
--- /dev/null
+++ b/spec/support/controllers/ldap_omniauth_callbacks_controller_shared_context.rb
@@ -0,0 +1,33 @@
+require 'spec_helper'
+
+shared_context 'Ldap::OmniauthCallbacksController' do
+  include LoginHelpers
+  include LdapHelpers
+
+  let(:uid) { 'my-uid' }
+  let(:provider) { 'ldapmain' }
+  let(:valid_login?) { true }
+  let(:user) { create(:omniauth_user, extern_uid: uid, provider: provider) }
+  let(:ldap_server_config) do
+    { main: ldap_config_defaults(:main) }
+  end
+
+  def ldap_config_defaults(key, hash = {})
+    {
+      provider_name: "ldap#{key}",
+      attributes: {},
+      encryption: 'plain'
+    }.merge(hash)
+  end
+
+  before do
+    stub_ldap_setting(enabled: true, servers: ldap_server_config)
+    described_class.define_providers!
+    Rails.application.reload_routes!
+
+    mock_auth_hash(provider.to_s, uid, user.email)
+    stub_omniauth_provider(provider, context: request)
+
+    allow(Gitlab::Auth::LDAP::Access).to receive(:allowed?).and_return(valid_login?)
+  end
+end
diff --git a/spec/support/ldap_helpers.rb b/spec/support/ldap_helpers.rb
index 0e87b3d359d..b90bbc4b106 100644
--- a/spec/support/ldap_helpers.rb
+++ b/spec/support/ldap_helpers.rb
@@ -18,6 +18,10 @@ module LdapHelpers
     allow_any_instance_of(::Gitlab::Auth::LDAP::Config).to receive_messages(messages)
   end
 
+  def stub_ldap_setting(messages)
+    allow(Gitlab.config.ldap).to receive_messages(to_settings(messages))
+  end
+
   # Stub an LDAP person search and provide the return entry. Specify `nil` for
   # `entry` to simulate when an LDAP person is not found
   #
diff --git a/spec/support/login_helpers.rb b/spec/support/login_helpers.rb
index db34090e971..72e5c2d66dd 100644
--- a/spec/support/login_helpers.rb
+++ b/spec/support/login_helpers.rb
@@ -112,7 +112,7 @@ module LoginHelpers
         }
       }
     })
-    Rails.application.env_config['omniauth.auth'] = OmniAuth.config.mock_auth[:saml]
+    Rails.application.env_config['omniauth.auth'] = OmniAuth.config.mock_auth[provider.to_sym]
   end
 
   def mock_saml_config
@@ -129,7 +129,7 @@ module LoginHelpers
     env = env_from_context(context)
 
     set_devise_mapping(context: context)
-    env['omniauth.auth'] = OmniAuth.config.mock_auth[provider]
+    env['omniauth.auth'] = OmniAuth.config.mock_auth[provider.to_sym]
   end
 
   def stub_omniauth_saml_config(messages)
-- 
GitLab