Move Geo stuff from sessions_controller to EE mixin

app/controllers/sessions_controller.rb

@@ -9,9 +9,7 @@ class SessionsController < Devise::SessionsController
   prepend_before_action :check_initial_setup, only: [:new]
   prepend_before_action :authenticate_with_two_factor,
     if: :two_factor_enabled?, only: [:create]
-  prepend_before_action :store_redirect_path, only: [:new]
-  before_action :gitlab_geo_login, only: [:new]
-  before_action :gitlab_geo_logout, only: [:destroy]
+  prepend_before_action :store_redirect_uri, only: [:new]
   before_action :auto_sign_in_with_provider, only: [:new]
   before_action :load_recaptcha
 
@@ -88,7 +86,11 @@ class SessionsController < Devise::SessionsController
     end
   end
 
-  def store_redirect_path
+  def stored_redirect_uri
+    @redirect_to ||= stored_location_for(:redirect)
+  end
+
+  def store_redirect_uri
     redirect_uri =
       if request.referer.present? && (params['redirect_to_referer'] == 'yes')
         URI(request.referer)
@@ -98,40 +100,22 @@ class SessionsController < Devise::SessionsController
 
     # Prevent a 'you are already signed in' message directly after signing:
     # we should never redirect to '/users/sign_in' after signing in successfully.
-    if redirect_uri.path == new_user_session_path
-      return true
-    elsif redirect_uri.host == Gitlab.config.gitlab.host && redirect_uri.port == Gitlab.config.gitlab.port
-      redirect_to = redirect_uri.to_s
-    elsif Gitlab::Geo.geo_node?(host: redirect_uri.host, port: redirect_uri.port)
-      redirect_to = redirect_uri.to_s
-    end
+    return true if redirect_uri.path == new_user_session_path
+
+    redirect_to = redirect_uri.to_s if redirect_allowed_to?(redirect_uri)
 
     @redirect_to = redirect_to
     store_location_for(:redirect, redirect_to)
   end
 
-  def two_factor_enabled?
-    find_user.try(:two_factor_enabled?)
-  end
-
-  def gitlab_geo_login
-    return unless Gitlab::Geo.secondary? # TODO?
-    return if signed_in?
-
-    oauth = Gitlab::Geo::OauthSession.new
-
-    # share full url with primary node by oauth state
-    user_return_to = URI.join(root_url, session[:user_return_to].to_s).to_s
-    oauth.return_to = @redirect_to || user_return_to
-
-    redirect_to oauth_geo_auth_url(state: oauth.generate_oauth_state)
+  # Overridden in EE
+  def redirect_allowed_to?(uri)
+    uri.host == Gitlab.config.gitlab.host &&
+      uri.port == Gitlab.config.gitlab.port
   end
 
-  def gitlab_geo_logout
-    return unless Gitlab::Geo.secondary? # TODO?
-
-    oauth = Gitlab::Geo::OauthSession.new(access_token: session[:access_token])
-    @geo_logout_state = oauth.generate_logout_state
+  def two_factor_enabled?
+    find_user&.two_factor_enabled?
   end
 
   def auto_sign_in_with_provider

ee/app/controllers/ee/sessions_controller.rb

@@ -2,13 +2,45 @@ module EE
   module SessionsController
     extend ActiveSupport::Concern
 
+    prepended do
+      before_action :gitlab_geo_login, only: [:new]
+      before_action :gitlab_geo_logout, only: [:destroy]
+    end
+
     private
 
+    def gitlab_geo_login
+      return unless ::Gitlab::Geo.secondary?
+      return if signed_in?
+
+      oauth = ::Gitlab::Geo::OauthSession.new
+
+      # share full url with primary node by oauth state
+      user_return_to = URI.join(root_url, session[:user_return_to].to_s).to_s
+      oauth.return_to = stored_redirect_uri || user_return_to
+
+      redirect_to oauth_geo_auth_url(state: oauth.generate_oauth_state)
+    end
+
+    def gitlab_geo_logout
+      return unless ::Gitlab::Geo.secondary?
+
+      oauth = ::Gitlab::Geo::OauthSession.new(access_token: session[:access_token])
+      @geo_logout_state = oauth.generate_logout_state
+    end
+
     def log_failed_login
       ::AuditEventService.new(request.filtered_parameters['user']['login'], nil, ip_address: request.remote_ip)
           .for_failed_login.unauth_security_event
 
       super
     end
+
+    def redirect_allowed_to?(uri)
+      raise NotImplementedError unless defined?(super)
+
+      # Redirect is not only allowed to current host, but also to other Geo nodes
+      super || ::Gitlab::Geo.geo_node?(host: uri.host, port: uri.port)
+    end
   end
 end