Make OmniAuth initializer return Geo proxied URL when it exists

OmniAuth full_host supports a proc, this makes it return a proc
that changes based on the env to verify whether it's a proxied
Geo request or not.

For regular requests, or non-EE requests, this keeps returning
the base URL, while proxied Geo requests return the URL of the
proxying site - to get SSO working with separate URLs.

EE: true
Changelog: changed

config/feature_flags/development/omniauth_initializer_fullhost_proc.yml

+---
+name: omniauth_initializer_fullhost_proc
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/82401
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/355579
+milestone: '14.10'
+type: development
+group: group::geo
+default_enabled: false

config/initializers/omniauth.rb

@@ -11,7 +11,15 @@ if Gitlab::Auth::Ldap::Config.enabled?
   end
 end
 
-OmniAuth.config.full_host = Settings.gitlab['base_url']
+OmniAuth.config.full_host =
+  if Feature.feature_flags_available? && ::Feature.enabled?(:omniauth_initializer_fullhost_proc, default_enabled: :yaml)
+    Gitlab::AppLogger.debug("Using OmniAuth proc initializer")
+    Gitlab::OmniauthInitializer.full_host
+  else
+    Gitlab::AppLogger.debug("Fallback to OmniAuth static full_host")
+    Settings.gitlab['base_url']
+  end
+
 OmniAuth.config.allowed_request_methods = [:post]
 # In case of auto sign-in, the GET method is used (users don't get to click on a button)
 OmniAuth.config.allowed_request_methods << :get if Gitlab.config.omniauth.auto_sign_in_with_provider.present?

ee/app/models/geo_node.rb

@@ -176,6 +176,10 @@ class GeoNode < ApplicationRecord
     @internal_uri ||= URI.parse(internal_url) if internal_url.present?
   end
 
+  def omniauth_host_url
+    read_without_ending_slash(:url)
+  end
+
   # Geo API endpoint for retrieving a replicable item
   #
   # @param [String] replicable_name
@@ -421,6 +425,12 @@ class GeoNode < ApplicationRecord
     add_ending_slash(value)
   end
 
+  def read_without_ending_slash(attribute)
+    value = read_attribute(attribute)
+
+    remove_ending_slash(value)
+  end
+
   def write_with_ending_slash(attribute, value)
     value = add_ending_slash(value)
 
@@ -434,6 +444,12 @@ class GeoNode < ApplicationRecord
     "#{value}/"
   end
 
+  def remove_ending_slash(value)
+    return value if value.blank?
+
+    value.sub(%r{/$}, '')
+  end
+
   def projects_for_selected_namespaces
     Project.where(Project.arel_table.name => { namespace_id: selected_namespaces_and_descendants.select(:id) })
   end

ee/lib/ee/gitlab/omniauth_initializer.rb

@@ -3,8 +3,18 @@
 module EE
   module Gitlab
     module OmniauthInitializer
+      extend ActiveSupport::Concern
       extend ::Gitlab::Utils::Override
 
+      class_methods do
+        extend ::Gitlab::Utils::Override
+
+        override :full_host
+        def full_host
+          proc { |env| ::Gitlab::Geo.proxied_site(env)&.omniauth_host_url || Settings.gitlab['base_url'] }
+        end
+      end
+
       override :build_omniauth_customized_providers
       def build_omniauth_customized_providers
         super.concat(%i[kerberos_spnego group_saml])

ee/lib/gitlab/geo.rb

@@ -16,6 +16,8 @@ module Gitlab
     ).freeze
 
     API_SCOPE = 'geo_api'
+    GEO_PROXIED_HEADER = 'HTTP_GITLAB_WORKHORSE_GEO_PROXY'
+    GEO_PROXIED_EXTRA_DATA_HEADER = 'HTTP_GITLAB_WORKHORSE_GEO_PROXY_EXTRA_DATA'
 
     # TODO: Avoid having to maintain a list. Discussions related to possible
     # solutions can be found at
@@ -119,7 +121,17 @@ module Gitlab
     end
 
     def self.proxied_request?(env)
-      env['HTTP_GITLAB_WORKHORSE_GEO_PROXY'] == '1'
+      env[GEO_PROXIED_HEADER] == '1'
+    end
+
+    def self.proxied_site(env)
+      return unless ::Gitlab::Geo.primary?
+      return unless proxied_request?(env) && env[GEO_PROXIED_EXTRA_DATA_HEADER].present?
+
+      signed_data = Gitlab::Geo::SignedData.new
+      signed_data.decode_data(env[GEO_PROXIED_EXTRA_DATA_HEADER])
+
+      signed_data.geo_node
     end
 
     def self.license_allows?

ee/spec/lib/ee/gitlab/omniauth_initializer_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::OmniauthInitializer do
+  include ::EE::GeoHelpers
+
+  describe '.full_host' do
+    subject { described_class.full_host.call({}) }
+
+    let(:base_url) { 'http://localhost/test' }
+
+    before do
+      allow(Settings).to receive(:gitlab).and_return({ 'base_url' => base_url })
+    end
+
+    context 'with non-proxied request' do
+      it { is_expected.to eq(base_url) }
+    end
+
+    context 'with a proxied request' do
+      context 'for a non-existing node' do
+        before do
+          stub_proxied_site(nil)
+        end
+
+        it { is_expected.to eq(base_url) }
+      end
+
+      context 'for an existing node' do
+        let(:geo_node) { instance_double(GeoNode, omniauth_host_url: 'http://localhost/geonode_url') }
+
+        before do
+          stub_proxied_site(geo_node)
+        end
+
+        it { is_expected.to eq(geo_node.omniauth_host_url) }
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/geo_spec.rb

@@ -254,12 +254,76 @@ RSpec.describe Gitlab::Geo, :geo, :request_store do
       expect(described_class.proxied_request?({ 'HTTP_GITLAB_WORKHORSE_GEO_PROXY' => '1' })).to be_truthy
     end
 
-    it 'returns false when the header is not present or set o an invalid value' do
+    it 'returns false when the header is not present or set to an invalid value' do
       expect(described_class.proxied_request?({})).to be_falsey
       expect(described_class.proxied_request?({ 'HTTP_GITLAB_WORKHORSE_GEO_PROXY' => 'invalid' })).to be_falsey
     end
   end
 
+  describe '.proxied_site' do
+    let(:env) { {} }
+
+    subject { described_class.proxied_site(env) }
+
+    context 'for a non-proxied request' do
+      it { is_expected.to be_nil }
+    end
+
+    context 'without Geo enabled' do
+      it { is_expected.to be_nil }
+    end
+
+    # this should not _really_ get called in a real-life scenario, as
+    # as a secondary should always proxy a primary, so this is nil in
+    # case this somehow happens
+    context 'on a secondary' do
+      before do
+        stub_secondary_node
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'on a primary' do
+      before do
+        stub_primary_node
+      end
+
+      context 'for a proxied request' do
+        before do
+          stub_proxied_request
+        end
+
+        context 'with an absent proxied site ID header' do
+          it { is_expected.to be_nil }
+        end
+
+        context 'with a proxy extra data header' do
+          context 'for an invalid header' do
+            let(:env) do
+              {
+                ::Gitlab::Geo::GEO_PROXIED_EXTRA_DATA_HEADER => "invalid"
+              }
+            end
+
+            it { is_expected.to be_nil }
+          end
+
+          context 'for an existing site' do
+            let(:signed_data) { Gitlab::Geo::SignedData.new(geo_node: secondary_node).sign_and_encode_data({}) }
+            let(:env) do
+              {
+                ::Gitlab::Geo::GEO_PROXIED_EXTRA_DATA_HEADER => signed_data
+              }
+            end
+
+            it { is_expected.to eq(secondary_node) }
+          end
+        end
+      end
+    end
+  end
+
   describe '.enabled?' do
     it_behaves_like 'a Geo cached value', :enabled?, :node_enabled
 

ee/spec/models/geo_node_spec.rb

@@ -547,6 +547,21 @@ RSpec.describe GeoNode, :request_store, :geo, type: :model do
     end
   end
 
+  describe '#omniauth_host_url' do
+    it 'returns a string' do
+      expect(new_node.omniauth_host_url).to be_a String
+    end
+
+    it 'is the URL without the trailing slash' do
+      expect("#{new_node.omniauth_host_url}/").to eq(new_node.url)
+    end
+
+    it 'includes schema home port and relative_url without a terminating /' do
+      expected_url = 'https://localhost:3000/gitlab'
+      expect(new_node.omniauth_host_url).to eq(expected_url)
+    end
+  end
+
   describe '#geo_retrieve_url' do
     let(:retrieve_url) { "https://localhost:3000/gitlab/api/#{api_version}/geo/retrieve/package_file/1" }
 

ee/spec/support/helpers/ee/geo_helpers.rb

@@ -25,6 +25,16 @@ module EE
       allow(::Gitlab::Geo).to receive(:secondary?).and_return(true)
     end
 
+    def stub_proxied_request
+      allow(::Gitlab::Geo).to receive(:proxied_request?).and_return(true)
+    end
+
+    def stub_proxied_site(node)
+      stub_proxied_request
+
+      allow(::Gitlab::Geo).to receive(:proxied_site).and_return(node)
+    end
+
     def create_project_on_shard(shard_name)
       project = create(:project)
 

lib/gitlab/omniauth_initializer.rb

@@ -38,6 +38,10 @@ module Gitlab
         end
       end
 
+      def full_host
+        proc { |_env| Settings.gitlab['base_url'] }
+      end
+
       private
 
       def cas3_signout_handler

spec/initializers/omniauth_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'OmniAuth initializer for GitLab' do
+  let(:load_omniauth_initializer) do
+    load Rails.root.join('config/initializers/omniauth.rb')
+  end
+
+  describe '#full_host' do
+    subject { OmniAuth.config.full_host }
+
+    let(:base_url) { 'http://localhost/test' }
+
+    before do
+      allow(Settings).to receive(:gitlab).and_return({ 'base_url' => base_url })
+      allow(Gitlab::OmniauthInitializer).to receive(:full_host).and_return('proc')
+    end
+
+    context 'with feature flags not available' do
+      before do
+        expect(Feature).to receive(:feature_flags_available?).and_return(false)
+        load_omniauth_initializer
+      end
+
+      it { is_expected.to eq(base_url) }
+    end
+
+    context 'with the omniauth_initializer_fullhost_proc FF disabled' do
+      before do
+        stub_feature_flags(omniauth_initializer_fullhost_proc: false)
+        load_omniauth_initializer
+      end
+
+      it { is_expected.to eq(base_url) }
+    end
+
+    context 'with the omniauth_initializer_fullhost_proc FF disabled' do
+      before do
+        load_omniauth_initializer
+      end
+
+      it { is_expected.to eq('proc') }
+    end
+  end
+end

spec/lib/gitlab/omniauth_initializer_spec.rb

@@ -309,4 +309,16 @@ RSpec.describe Gitlab::OmniauthInitializer do
       subject.execute([conf])
     end
   end
+
+  describe '.full_host' do
+    subject { described_class.full_host.call({}) }
+
+    let(:base_url) { 'http://localhost/test' }
+
+    before do
+      allow(Settings).to receive(:gitlab).and_return({ 'base_url' => base_url })
+    end
+
+    it { is_expected.to eq(base_url) }
+  end
 end