Redirect user back to secondary after logout/login

Redirect user back to the secondary after a logout & re-login
via the primary.

ee/app/controllers/oauth/geo_auth_controller.rb

@@ -3,7 +3,6 @@ class Oauth::GeoAuthController < ActionController::Base
   rescue_from OAuth2::Error, with: :auth
 
   def auth
-    oauth = Gitlab::Geo::OauthSession.new(state: params[:state])
     unless oauth.oauth_state_valid?
       redirect_to root_url
       return
@@ -14,7 +13,6 @@ class Oauth::GeoAuthController < ActionController::Base
 
   # rubocop: disable CodeReuse/ActiveRecord
   def callback
-    oauth = Gitlab::Geo::OauthSession.new(state: params[:state])
     unless oauth.oauth_state_valid?
       redirect_to new_user_session_path
       return
@@ -25,7 +23,7 @@ class Oauth::GeoAuthController < ActionController::Base
 
     user = User.find_by(id: remote_user['id'])
 
-    if user && sign_in(user, bypass: true)
+    if user && bypass_sign_in(user)
       after_sign_in_with_gitlab(token, oauth.get_oauth_state_return_to)
     else
       invalid_credentials
@@ -36,9 +34,10 @@ class Oauth::GeoAuthController < ActionController::Base
   def logout
     logout = Oauth2::LogoutTokenValidationService.new(current_user, params)
     result = logout.execute
+
     if result[:status] == :success
       sign_out current_user
-      redirect_to root_path
+      after_sign_out_with_gitlab(result[:return_to])
     else
       access_token_error(result[:message])
     end
@@ -46,11 +45,20 @@ class Oauth::GeoAuthController < ActionController::Base
 
   private
 
+  def oauth
+    @oauth ||= Gitlab::Geo::OauthSession.new(state: params[:state])
+  end
+
   def after_sign_in_with_gitlab(token, return_to)
     session[:access_token] = token
     redirect_to(return_to || root_path)
   end
 
+  def after_sign_out_with_gitlab(return_to)
+    session[:user_return_to] = return_to
+    redirect_to(root_path)
+  end
+
   def invalid_credentials
     @error = 'Cannot find user to login. Your account may have been deleted.'
     render :error, layout: 'errors'

ee/app/models/geo_node.rb

@@ -89,6 +89,11 @@ class GeoNode < ActiveRecord::Base
       left_join_status.minimum(:cursor_last_event_id)
     end
 
+    # Tries to find a GeoNode by oauth_application_id, returning nil if none could be found.
+    def find_by_oauth_application_id(oauth_application_id)
+      where(oauth_application_id: oauth_application_id).take
+    end
+
     private
 
     def left_join_status

ee/app/services/oauth2/logout_token_validation_service.rb

 module Oauth2
   class LogoutTokenValidationService < ::BaseService
+    include Gitlab::Utils::StrongMemoize
+
     attr_reader :status
 
     def initialize(user, params = {})
@@ -8,32 +10,33 @@ module Oauth2
     end
 
     def execute
-      return error('access token not found') unless access_token
+      return error('Access token not found') unless access_token
 
       status = AccessTokenValidationService.new(access_token).validate
+      return error(status) unless status == AccessTokenValidationService::VALID
 
-      if status == AccessTokenValidationService::VALID
-        user = User.find(access_token.resource_owner_id)
-
-        if current_user == user
-          success
-        end
-      else
-        error(status)
-      end
+      user = User.find(access_token.resource_owner_id)
+      success(return_to: geo_node_url) if user == current_user
     end
 
-    def access_token
-      @access_token ||= begin
-        return unless  params[:state] && !params[:state].empty?
-
-        oauth_session = Gitlab::Geo::OauthSession.new(state: params[:state])
+    private
 
+    def access_token
+      strong_memoize(:access_token) do
         logout_token = oauth_session.extract_logout_token
-        return unless logout_token && logout_token.is_utf8?
 
-        Doorkeeper::AccessToken.by_token(logout_token)
+        if logout_token&.is_utf8?
+          Doorkeeper::AccessToken.by_token(logout_token)
+        end
       end
     end
+
+    def oauth_session
+      @oauth_session ||= Gitlab::Geo::OauthSession.new(state: params[:state])
+    end
+
+    def geo_node_url
+      GeoNode.find_by_oauth_application_id(access_token.application_id)&.url
+    end
   end
 end

ee/changelogs/unreleased/4459-redirect-users-back-to-secondary-after-logout-login.yml

+---
+title: Geo - Redirect user back to the secondary after a logout & re-login via the primary
+merge_request: 8157
+author:
+type: fixed

ee/lib/gitlab/geo/oauth_session.rb

@@ -35,7 +35,7 @@ module Gitlab
       end
 
       def extract_logout_token
-        return unless state
+        return unless state.present?
 
         salt, encrypted = state.split(':', 2)
         decipher = logout_token_cipher(salt, :decrypt)

ee/spec/controllers/oauth/geo_auth_controller_spec.rb

@@ -3,8 +3,8 @@ require 'spec_helper'
 describe Oauth::GeoAuthController do
   let(:user) { create(:user) }
   let(:oauth_app) { create(:doorkeeper_application) }
-  let(:access_token) { create(:doorkeeper_access_token, resource_owner_id: user.id).token }
-  let(:auth_state) { Gitlab::Geo::OauthSession.new(access_token: access_token, return_to: projects_url).generate_oauth_state }
+  let(:access_token) { create(:doorkeeper_access_token, resource_owner_id: user.id, application: oauth_app) }
+  let(:auth_state) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: projects_url).generate_oauth_state }
   let(:primary_node_url) { 'http://localhost:3001/' }
 
   before do
@@ -30,7 +30,7 @@ describe Oauth::GeoAuthController do
   end
 
   describe 'GET callback' do
-    let(:callback_state) { Gitlab::Geo::OauthSession.new(access_token: access_token, return_to: projects_url).generate_oauth_state }
+    let(:callback_state) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: projects_url).generate_oauth_state }
     let(:primary_node_oauth_endpoint) { Gitlab::Geo::OauthSession.new.authorize_url(redirect_uri: oauth_geo_callback_url, state: callback_state) }
 
     context 'redirection' do
@@ -58,7 +58,7 @@ describe Oauth::GeoAuthController do
       let(:oauth_error) { OAuth2::Error.new(OAuth2::Response.new(fake_response)) }
 
       before do
-        expect_any_instance_of(Gitlab::Geo::OauthSession).to receive(:get_token) { access_token }
+        expect_any_instance_of(Gitlab::Geo::OauthSession).to receive(:get_token) { access_token.token }
         expect_any_instance_of(Gitlab::Geo::OauthSession).to receive(:authenticate_with_gitlab).and_raise(oauth_error)
       end
 
@@ -87,23 +87,27 @@ describe Oauth::GeoAuthController do
   end
 
   describe 'GET logout' do
-    let(:logout_state) { Gitlab::Geo::OauthSession.new(access_token: access_token).generate_logout_state }
+    let(:oauth_session) { Gitlab::Geo::OauthSession.new(access_token: access_token.token) }
+    let(:logout_state) { oauth_session.generate_logout_state }
 
-    context 'access_token error' do
-      render_views
+    render_views
 
-      before do
-        sign_in(user)
-      end
+    before do
+      sign_in(user)
+    end
 
-      it 'logs out when correct access_token is informed' do
+    context 'when access_token is valid' do
+      it 'logs out and redirects to the root_url' do
         get :logout, state: logout_state
 
         expect(response).to redirect_to root_url
       end
+    end
 
+    context 'when access_token is invalid' do
       it 'handles access token problems' do
         allow_any_instance_of(Oauth2::LogoutTokenValidationService).to receive(:execute) { { status: :error, message: :expired } }
+
         get :logout, state: logout_state
 
         expect(response.body).to include("There is a problem with the OAuth access_token: expired")

ee/spec/lib/gitlab/geo/oauth_session_spec.rb

@@ -87,6 +87,12 @@ describe Gitlab::Geo::OauthSession do
       expect(subject.extract_logout_token).to be_nil
     end
 
+    it 'returns nil when state is empty' do
+      subject.state = ''
+
+      expect(subject.extract_logout_token).to be_nil
+    end
+
     it 'returns false when decryptation fails' do
       subject.generate_logout_state
       allow_any_instance_of(OpenSSL::Cipher::AES).to receive(:final) { raise OpenSSL::OpenSSLError }

ee/spec/models/geo_node_spec.rb

@@ -158,6 +158,24 @@ describe GeoNode, type: :model do
     end
   end
 
+  describe '.find_by_oauth_application_id' do
+    context 'when the Geo node exists' do
+      it 'returns the Geo node' do
+        found = described_class.find_by_oauth_application_id(node.oauth_application_id)
+
+        expect(found).to eq(node)
+      end
+    end
+
+    context 'when the Geo node does not exist' do
+      it 'returns nil' do
+        found = described_class.find_by_oauth_application_id(-1)
+
+        expect(found).to be_nil
+      end
+    end
+  end
+
   describe '#repair' do
     it 'creates an oauth application for a Geo secondary node' do
       stub_current_geo_node(node)

ee/spec/services/oauth2/logout_token_validation_service_spec.rb

 require 'spec_helper'
 
 describe Oauth2::LogoutTokenValidationService do
-  let(:user) { FactoryBot.create(:user) }
-  let(:access_token) { FactoryBot.create(:doorkeeper_access_token, resource_owner_id: user.id).token }
-  let(:logout_state) { Gitlab::Geo::OauthSession.new(access_token: access_token).generate_logout_state }
+  let(:user) { create(:user) }
+  let(:node) { create(:geo_node) }
+  let(:access_token) { create(:doorkeeper_access_token, resource_owner_id: user.id, application_id: node.oauth_application_id) }
+  let(:oauth_session) { Gitlab::Geo::OauthSession.new(access_token: access_token.token) }
+  let(:logout_state) { oauth_session.generate_logout_state }
 
   context '#execute' do
     it 'return error when params are empty' do
       result = described_class.new(user, {}).execute
+
       expect(result[:status]).to eq(:error)
     end
 
-    it 'returns error when state param is empty' do
+    it 'returns error when state param is nil' do
       result = described_class.new(user, { state: nil }).execute
+
       expect(result[:status]).to eq(:error)
+    end
 
+    it 'returns error when state param is empty' do
       result = described_class.new(user, { state: '' }).execute
+
       expect(result[:status]).to eq(:error)
     end
 
@@ -24,12 +31,14 @@ describe Oauth2::LogoutTokenValidationService do
       allow_any_instance_of(Gitlab::Geo::OauthSession).to receive(:extract_logout_token) { invalid_token }
 
       result = described_class.new(user, { state: logout_state }).execute
+
       expect(result[:status]).to eq(:error)
     end
 
-    it 'returns true when token is valid' do
+    it 'returns success when token is valid' do
       result = described_class.new(user, { state: logout_state }).execute
-      expect(result[:status]).to eq(:success)
+
+      expect(result).to eq(status: :success, return_to: node.url)
     end
   end
 end