Include full path while redirecting the user back to secondary

ee/app/controllers/ee/sessions_controller.rb

@@ -26,7 +26,11 @@ module EE
     def gitlab_geo_logout
       return unless ::Gitlab::Geo.secondary?
 
-      oauth = ::Gitlab::Geo::OauthSession.new(access_token: session[:access_token])
+      oauth = ::Gitlab::Geo::OauthSession.new(
+        access_token: session[:access_token],
+        return_to: safe_redirect_path_for_url(request.referer)
+      )
+
       @geo_logout_state = oauth.generate_logout_state # rubocop:disable Gitlab/ModuleWithInstanceVariables
     end
 

ee/app/controllers/oauth/geo_auth_controller.rb

@@ -11,7 +11,6 @@ class Oauth::GeoAuthController < ActionController::Base
     redirect_to oauth.authorize_url(redirect_uri: oauth_geo_callback_url, state: params[:state])
   end
 
-  # rubocop: disable CodeReuse/ActiveRecord
   def callback
     unless oauth.oauth_state_valid?
       redirect_to new_user_session_path
@@ -20,8 +19,7 @@ class Oauth::GeoAuthController < ActionController::Base
 
     token = oauth.get_token(params[:code], redirect_uri: oauth_geo_callback_url)
     remote_user = oauth.authenticate_with_gitlab(token)
-
-    user = User.find_by(id: remote_user['id'])
+    user = UserFinder.new(remote_user['id']).find_by_id
 
     if user && bypass_sign_in(user)
       after_sign_in_with_gitlab(token, oauth.get_oauth_state_return_to)
@@ -29,7 +27,6 @@ class Oauth::GeoAuthController < ActionController::Base
       invalid_credentials
     end
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
   def logout
     logout = Oauth2::LogoutTokenValidationService.new(current_user, params)
@@ -51,6 +48,10 @@ class Oauth::GeoAuthController < ActionController::Base
 
   def after_sign_in_with_gitlab(token, return_to)
     session[:access_token] = token
+
+    # Prevent alert from popping up on the first page shown after authentication.
+    flash[:alert] = nil
+
     redirect_to(return_to || root_path)
   end
 

ee/app/services/oauth2/logout_token_validation_service.rb

@@ -2,21 +2,21 @@ module Oauth2
   class LogoutTokenValidationService < ::BaseService
     include Gitlab::Utils::StrongMemoize
 
-    attr_reader :status
+    attr_reader :state
 
     def initialize(user, params = {})
-      @params = params
       @current_user = user
+      @state = params[:state]
     end
 
     def execute
-      return error('Access token not found') unless access_token
+      return error('Access token not found') unless access_token.present?
 
       status = AccessTokenValidationService.new(access_token).validate
       return error(status) unless status == AccessTokenValidationService::VALID
 
       user = User.find(access_token.resource_owner_id)
-      success(return_to: geo_node_url) if user == current_user
+      success(return_to: user_return_to) if user == current_user
     end
 
     private
@@ -32,7 +32,12 @@ module Oauth2
     end
 
     def oauth_session
-      @oauth_session ||= Gitlab::Geo::OauthSession.new(state: params[:state])
+      @oauth_session ||= Gitlab::Geo::OauthSession.new(state: state)
+    end
+
+    def user_return_to
+      full_path = oauth_session.get_oauth_state_return_to_full_path
+      URI.join(geo_node_url, full_path).to_s
     end
 
     def geo_node_url

ee/lib/gitlab/geo/oauth_session.rb

@@ -11,7 +11,6 @@ module Gitlab
         return false unless state
 
         salt, hmac, return_to = state.split(':', 3)
-
         return false unless return_to
 
         hmac == generate_oauth_hmac(salt, return_to)
@@ -29,7 +28,9 @@ module Gitlab
 
         cipher = logout_token_cipher(oauth_salt, :encrypt)
         encrypted = cipher.update(access_token) + cipher.final
-        self.state = "#{oauth_salt}:#{Base64.urlsafe_encode64(encrypted)}"
+        full_path = ReturnToLocation.new(return_to).full_path
+
+        self.state = "#{oauth_salt}:#{Base64.urlsafe_encode64(encrypted)}:#{full_path}"
       rescue OpenSSL::OpenSSLError
         false
       end
@@ -37,7 +38,8 @@ module Gitlab
       def extract_logout_token
         return unless state.present?
 
-        salt, encrypted = state.split(':', 2)
+        salt, encrypted, _ = state.split(':', 3)
+
         decipher = logout_token_cipher(salt, :decrypt)
         decipher.update(Base64.urlsafe_decode64(encrypted)) + decipher.final
       rescue OpenSSL::OpenSSLError
@@ -48,6 +50,10 @@ module Gitlab
         state.split(':', 3)[2] if state
       end
 
+      def get_oauth_state_return_to_full_path
+        ReturnToLocation.new(get_oauth_state_return_to).full_path
+      end
+
       def authorize_url(params = {})
         oauth_client.auth_code.authorize_url(params)
       end
@@ -65,6 +71,26 @@ module Gitlab
 
       private
 
+      class ReturnToLocation < Struct.new(:location)
+        def full_path
+          uri = parse_uri(location)
+          full_path_for_uri(uri) if uri
+        end
+
+        private
+
+        def parse_uri(location)
+          location && URI.parse(location)
+        rescue URI::InvalidURIError
+          nil
+        end
+
+        def full_path_for_uri(uri)
+          path_with_query = [uri.path, uri.query].compact.join('?')
+          [path_with_query, uri.fragment].compact.join("#")
+        end
+      end
+
       def generate_oauth_hmac(salt, return_to)
         return false unless return_to
 

ee/spec/controllers/oauth/geo_auth_controller_spec.rb

@@ -2,9 +2,11 @@ require 'spec_helper'
 
 describe Oauth::GeoAuthController do
   let(:user) { create(:user) }
-  let(:oauth_app) { create(:doorkeeper_application) }
+  let(:node) { create(:geo_node) }
+  let(:oauth_app) { node.oauth_application }
   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(:oauth_session) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: projects_url) }
+  let(:auth_state) { oauth_session.generate_oauth_state }
   let(:primary_node_url) { 'http://localhost:3001/' }
 
   before do
@@ -30,7 +32,8 @@ describe Oauth::GeoAuthController do
   end
 
   describe 'GET callback' do
-    let(:callback_state) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: projects_url).generate_oauth_state }
+    let(:oauth_session) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: projects_url) }
+    let(:callback_state) { oauth_session.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
@@ -51,6 +54,12 @@ describe Oauth::GeoAuthController do
 
         expect(response).to redirect_to(projects_url)
       end
+
+      it 'does not display a flash message if state is valid' do
+        get :callback, state: callback_state
+
+        expect(controller).to set_flash[:alert].to(nil)
+      end
     end
 
     context 'invalid credentials' do

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

@@ -52,31 +52,62 @@ describe Gitlab::Geo::OauthSession do
   end
 
   describe '#get_oauth_state_return_to' do
-    subject { described_class.new(state: valid_state) }
-
     it 'returns return_to value' do
+      subject = described_class.new(state: valid_state)
+
       expect(subject.get_oauth_state_return_to).to eq(oauth_return_to)
     end
   end
 
+  describe '#get_oauth_state_return_to_full_path' do
+    it 'removes the domain from return_to value' do
+      subject = described_class.new(state: valid_state)
+
+      expect(subject.get_oauth_state_return_to_full_path).to eq('/oauth/geo/callback')
+    end
+  end
+
   describe '#generate_logout_state' do
     it 'returns nil when access_token is not defined' do
       expect(described_class.new.generate_logout_state).to be_nil
     end
 
     it 'returns false when encryptation fails' do
-      allow_any_instance_of(OpenSSL::Cipher::AES).to receive(:final) { raise OpenSSL::OpenSSLError }
+      allow_any_instance_of(OpenSSL::Cipher::AES)
+        .to receive(:final) { raise OpenSSL::OpenSSLError }
+
       expect(subject.generate_logout_state).to be_falsey
     end
 
-    it 'returns a string with salt and encrypted access token colon separated' do
-      state = described_class.new(access_token: access_token).generate_logout_state
+    it 'returns a string with salt, encrypted access token, and return_to colon separated' do
+      subject = described_class.new(access_token: access_token, return_to: oauth_return_to)
+
+      state = subject.generate_logout_state
       expect(state).to be_a String
       expect(state).not_to be_blank
 
-      salt, encrypted = state.split(':', 2)
+      salt, encrypted, return_to = state.split(':', 3)
       expect(salt).not_to be_blank
       expect(encrypted).not_to be_blank
+      expect(return_to).not_to be_blank
+    end
+
+    it 'include a empty value for return_to into state when return_to param is not defined' do
+      subject = described_class.new(access_token: access_token)
+
+      state = subject.generate_logout_state
+      _, _, return_to = state.split(':', 3)
+
+      expect(return_to).to eq ''
+    end
+
+    it 'does not include the host from return_to param into into the state' do
+      subject = described_class.new(access_token: access_token, return_to: oauth_return_to)
+
+      state = subject.generate_logout_state
+      _, _, return_to = state.split(':', 3)
+
+      expect(return_to).to eq '/oauth/geo/callback'
     end
   end
 
@@ -94,8 +125,8 @@ describe Gitlab::Geo::OauthSession do
     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 }
+      allow_any_instance_of(OpenSSL::Cipher::AES)
+        .to receive(:final) { raise OpenSSL::OpenSSLError }
 
       expect(subject.extract_logout_token).to be_falsey
     end

ee/spec/services/oauth2/logout_token_validation_service_spec.rb

@@ -4,7 +4,8 @@ describe Oauth2::LogoutTokenValidationService do
   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(:oauth_return_to) { '/project/test' }
+  let(:oauth_session) { Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: oauth_return_to) }
   let(:logout_state) { oauth_session.generate_logout_state }
 
   context '#execute' do
@@ -15,30 +16,73 @@ describe Oauth2::LogoutTokenValidationService do
     end
 
     it 'returns error when state param is nil' do
-      result = described_class.new(user, { state: nil }).execute
+      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
+      result = described_class.new(user, state: '').execute
 
       expect(result[:status]).to eq(:error)
     end
 
     it 'returns error when incorrect encoding' do
-      invalid_token = "\xD800\xD801\xD802"
-      allow_any_instance_of(Gitlab::Geo::OauthSession).to receive(:extract_logout_token) { invalid_token }
+      allow_any_instance_of(Gitlab::Geo::OauthSession)
+        .to receive(:extract_logout_token)
+        .and_return("\xD800\xD801\xD802")
 
-      result = described_class.new(user, { state: logout_state }).execute
+      result = described_class.new(user, state: logout_state).execute
 
       expect(result[:status]).to eq(:error)
     end
 
-    it 'returns success when token is valid' do
-      result = described_class.new(user, { state: logout_state }).execute
+    context 'when token is valid' do
+      it 'returns success' do
+        result = described_class.new(user, state: logout_state).execute
 
-      expect(result).to eq(status: :success, return_to: node.url)
+        expect(result).to include(status: :success)
+      end
+
+      context 'when OAuth session return_to param is nil' do
+        it 'returns the Geo node URL associated with OAuth application to redirect the user back' do
+          oauth_session = Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: nil)
+          logout_state = oauth_session.generate_logout_state
+
+          result = described_class.new(user, state: logout_state).execute
+
+          expect(result).to include(return_to: node.url)
+        end
+      end
+
+      context 'when OAuth session return_to param is empty' do
+        it 'returns the Geo node URL associated with OAuth application to redirect the user back' do
+          oauth_session = Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: '')
+          logout_state = oauth_session.generate_logout_state
+
+          result = described_class.new(user, state: logout_state).execute
+
+          expect(result).to include(return_to: node.url)
+        end
+      end
+
+      context 'when OAuth session return_to param is set' do
+        it 'returns the fullpath to the Geo node to redirect the user back' do
+          result = described_class.new(user, state: logout_state).execute
+
+          expect(result).to include(return_to: URI.join(node.url, oauth_return_to).to_s)
+        end
+
+        it 'replaces the host with the Geo node associated with OAuth application' do
+          oauth_return_to = 'http://fake-secondary/project/test'
+          oauth_session = Gitlab::Geo::OauthSession.new(access_token: access_token.token, return_to: oauth_return_to)
+          logout_state = oauth_session.generate_logout_state
+
+          result = described_class.new(user, state: logout_state).execute
+
+          expect(result).to include(return_to: URI.join(node.url, '/project/test').to_s)
+        end
+      end
     end
   end
 end