Refactor OmniauthCallbacksController to remove duplication

7 files changed, 113 insertions, 30 deletions

diff --git a/app/controllers/omniauth_callbacks_controller.rb b/app/controllers/omniauth_callbacks_controller.rb
index d631d09f1b8..f62bc9c5736 100644
--- a/app/controllers/omniauth_callbacks_controller.rb
+++ b/app/controllers/omniauth_callbacks_controller.rb
@@ -32,9 +32,8 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
   # if the authentication to LDAP was successful.
   def ldap
     ldap_user = Gitlab::LDAP::User.new(oauth)
-    ldap_user.save if ldap_user.changed? # will also save new users
 
-    @user = ldap_user.gl_user
+    @user = ldap_user.find_and_update!
     @user.remember_me = params[:remember_me] if ldap_user.persisted?
 
     # Do additional LDAP checks for the user filter and EE features
@@ -51,25 +50,7 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
   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::Saml::User.new(oauth)
-      saml_user.save if saml_user.changed?
-      @user = saml_user.gl_user
-
-      continue_login_process
-    end
-  rescue Gitlab::OAuth::SignupDisabledError
-    handle_signup_error
+    omniauth_flow(Gitlab::Saml::CallbackHandler.new(self, oauth))
   end
 
   def omniauth_error
@@ -98,17 +79,14 @@ class OmniauthCallbacksController < Devise::OmniauthCallbacksController
   private
 
   def handle_omniauth
+    omniauth_flow(Gitlab::OAuth::CallbackHandler.new(self, oauth))
+  end
+
+  def omniauth_flow(callback_handler)
     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'
+      callback_handler.link_provider!
     else
-      oauth_user = Gitlab::OAuth::User.new(oauth)
-      oauth_user.save
-      @user = oauth_user.gl_user
+      @user = callback_handler.find_and_update_user!
 
       continue_login_process
     end
diff --git a/lib/gitlab/ldap/user.rb b/lib/gitlab/ldap/user.rb
index 84ee94e38e4..37d09c35237 100644
--- a/lib/gitlab/ldap/user.rb
+++ b/lib/gitlab/ldap/user.rb
@@ -7,6 +7,8 @@
 module Gitlab
   module LDAP
     class User < Gitlab::OAuth::User
+      extend ::Gitlab::Utils::Override
+
       class << self
         def find_by_uid_and_provider(uid, provider)
           identity = ::Identity.with_extern_uid(provider, uid).take
@@ -32,6 +34,11 @@ module Gitlab
         gl_user.changed? || gl_user.identities.any?(&:changed?)
       end
 
+      override :omniauth_should_save?
+      def omniauth_should_save?
+        changed?
+      end
+
       def block_after_signup?
         ldap_config.block_auto_created_users
       end
diff --git a/lib/gitlab/o_auth/callback_handler.rb b/lib/gitlab/o_auth/callback_handler.rb
new file mode 100644
index 00000000000..f248e6bb3ca
--- /dev/null
+++ b/lib/gitlab/o_auth/callback_handler.rb
@@ -0,0 +1,19 @@
+module Gitlab
+  module OAuth
+    class CallbackHandler < OmniauthCallbackHandlerBase
+      def link_provider!
+        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'])
+        notify_authentication_updated!
+      end
+
+      protected
+
+      def linked_user_class
+        Gitlab::OAuth::User
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/o_auth/user.rb b/lib/gitlab/o_auth/user.rb
index fff9360ea27..4c9193b4aa3 100644
--- a/lib/gitlab/o_auth/user.rb
+++ b/lib/gitlab/o_auth/user.rb
@@ -60,8 +60,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/omniauth_callback_handler_base.rb b/lib/gitlab/omniauth_callback_handler_base.rb
new file mode 100644
index 00000000000..9dd1bd81b03
--- /dev/null
+++ b/lib/gitlab/omniauth_callback_handler_base.rb
@@ -0,0 +1,37 @@
+module Gitlab
+  class OmniauthCallbackHandlerBase
+    include Gitlab::Routing
+
+    delegate :redirect_to, :current_user, to: :@controller
+
+    attr_reader :oauth
+
+    def initialize(controller, oauth)
+      @controller = controller
+      @oauth = oauth
+    end
+
+    def link_provider!
+      raise NotImplementedError
+    end
+
+    def find_and_update_user!
+      linked_user_class.new(oauth).find_and_update!
+    end
+
+    protected
+
+    def notify_authentication_updated!
+      redirect_to profile_account_path, notice: 'Authentication method updated'
+    end
+
+    def linked_user_class
+      raise NotImplementedError
+    end
+
+    def log_audit_event(user, options = {})
+      AuditEventService.new(user, user, options)
+        .for_authentication.security_event
+    end
+  end
+end
diff --git a/lib/gitlab/saml/callback_handler.rb b/lib/gitlab/saml/callback_handler.rb
new file mode 100644
index 00000000000..106ac3696fa
--- /dev/null
+++ b/lib/gitlab/saml/callback_handler.rb
@@ -0,0 +1,25 @@
+module Gitlab
+  module Saml
+    class CallbackHandler < OmniauthCallbackHandlerBase
+      def link_provider!
+        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)
+          notify_authentication_updated!
+        else
+          redirect_to after_sign_in_path_for(current_user)
+        end
+      end
+
+      protected
+
+      def linked_user_class
+        Gitlab::OAuth::User
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/saml/user.rb b/lib/gitlab/saml/user.rb
index d8faf7aad8c..346253583ca 100644
--- a/lib/gitlab/saml/user.rb
+++ b/lib/gitlab/saml/user.rb
@@ -6,6 +6,8 @@
 module Gitlab
   module Saml
     class User < Gitlab::OAuth::User
+      extend ::Gitlab::Utils::Override
+
       def save
         super('SAML')
       end
@@ -32,6 +34,11 @@ module Gitlab
         gl_user.changed? || gl_user.identities.any?(&:changed?)
       end
 
+      override :omniauth_should_save?
+      def omniauth_should_save?
+        changed?
+      end
+
       protected
 
       def auto_link_saml_user?