Redirect to space on subdomain following successful omniauth auth

Author: fbacall

app/controllers/callbacks_controller.rb

@@ -11,6 +11,9 @@ class CallbacksController < Devise::OmniauthCallbacksController
 
   def handle_callback(provider, config)
     @user = User.from_omniauth(request.env["omniauth.auth"])
+    if request.env['omniauth.params'] && request.env['omniauth.params']['space_id']
+      space = Space.find_by_id(request.env['omniauth.params']['space_id'])
+    end
 
     if @user.new_record?
       # new user
@@ -27,13 +30,28 @@ def handle_callback(provider, config)
 
         sign_in @user
         flash[:notice] = "#{I18n.t('devise.registrations.signed_up')} Please ensure your profile is correct."
-        redirect_to edit_user_path(@user)
+        redirect_to_space(edit_user_path(@user), space)
       rescue Exception => e
         flash[:notice] = "Login failed: #{e.message.to_s}"
-        redirect_to new_user_session_path
+        redirect_to_space(new_user_session_path, space)
       end
     else
-      sign_in_and_redirect @user
+      scope = Devise::Mapping.find_scope!(@user)
+      sign_in(scope, resource, {})
+      redirect_to_space(after_sign_in_path_for(@user), space)
+    end
+  end
+
+  private
+
+  def redirect_to_space(path, space)
+    if space && space.is_subdomain?(request.domain)
+      port_part = ''
+      port_part = ":#{request.port}" if (request.protocol == "http://" && request.port != 80) ||
+                                        (request.protocol == "https://" && request.port != 443)
+      redirect_to URI.join("#{request.protocol}#{space.host}#{port_part}", path).to_s, allow_other_host: true
+    else
+      redirect_to path
     end
   end
 

app/helpers/application_helper.rb

@@ -700,11 +700,12 @@ def theme_path
   end
 
   def omniauth_login_link(provider, config)
+    params = Space.current_space&.default? ? {} : { space_id: Space.current_space.id }
     link_to(
       t('authentication.omniauth.log_in_with',
         provider: config.options[:label] ||
                   t("authentication.omniauth.providers.#{provider}", default: provider.to_s.titleize)),
-      omniauth_authorize_path('user', provider),
+      omniauth_authorize_path('user', provider, **params),
       method: :post
     )
   end

app/models/space.rb

@@ -65,6 +65,10 @@ def enabled_features
     (FEATURES - disabled_features)
   end
 
+  def is_subdomain?(domain)
+    (host == domain || host.ends_with?(".#{domain}"))
+  end
+
   private
 
   def disabled_features_valid?

test/integration/omniauth_test.rb

@@ -399,4 +399,64 @@ class OmniauthTest < ActionDispatch::IntegrationTest
     assert_response :not_found
   end
 
+  test 'authentication redirects users back to origin space on subdomain' do
+    space = spaces(:astro)
+    space.update!(host: 'space.example.com')
+
+    OmniAuth.config.mock_auth[:oidc] = OmniAuth::AuthHash.new(
+      {
+        provider: 'oidc',
+        uid: '0123456789abcdcef',
+        info: {
+          email: 'aai@example.com',
+          nickname: 'aaf_user',
+          first_name: 'AAF',
+          last_name: 'User'
+        }
+      })
+
+    post user_oidc_omniauth_authorize_url(space_id: space.id)
+    follow_redirect! # OmniAuth redirect
+    assert_equal "http://space.example.com/users/aaf_user/edit", response.headers['Location']
+    end
+
+  test 'authentication does not redirect user to entirely different domain' do
+    space = spaces(:astro)
+    space.update!(host: 'my-cool-space-host.com')
+
+    OmniAuth.config.mock_auth[:oidc] = OmniAuth::AuthHash.new(
+      {
+        provider: 'oidc',
+        uid: '0123456789abcdcef',
+        info: {
+          email: 'aai@example.com',
+          nickname: 'aaf_user',
+          first_name: 'AAF',
+          last_name: 'User'
+        }
+      })
+
+    post user_oidc_omniauth_authorize_url(space_id: space.id)
+    follow_redirect! # OmniAuth redirect
+    assert_equal "http://www.example.com/users/aaf_user/edit", response.headers['Location']
+  end
+
+  test 'invalid space is ignored when redirecting' do
+    OmniAuth.config.mock_auth[:oidc] = OmniAuth::AuthHash.new(
+      {
+        provider: 'oidc',
+        uid: '0123456789abcdcef',
+        info: {
+          email: 'aai@example.com',
+          nickname: 'aaf_user',
+          first_name: 'AAF',
+          last_name: 'User'
+        }
+      })
+
+    post user_oidc_omniauth_authorize_url(space_id: 'ufhgfsdkhgskdjfhsdkjfhsdkjfhsd')
+    follow_redirect! # OmniAuth redirect
+    assert_equal "http://www.example.com/users/aaf_user/edit", response.headers['Location']
+  end
+
 end

test/models/space_test.rb

@@ -128,4 +128,12 @@ class SpaceTest < ActiveSupport::TestCase
     @space.enabled_features = Space::FEATURES
     assert_equal [], @space.disabled_features
   end
+
+  test 'check subdomain of' do
+    assert @space.is_subdomain?('mytess.training')
+    refute @space.is_subdomain?('amytess.training')
+    refute @space.is_subdomain?('mytess.com')
+    refute @space.is_subdomain?('mytess.training.com')
+    refute @space.is_subdomain?('space.mytess.training')
+  end
 end