From e6fcaadd456c1ac02c8d32e4e968577e683f8e65 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Fri, 16 Sep 2016 13:48:50 +1000
Subject: [PATCH] FIX: redirects back to origin for SSO and omniauth login

---
 app/controllers/session_controller.rb            | 16 +++++++++-------
 app/controllers/user_api_keys_controller.rb      |  7 ++++++-
 .../users/omniauth_callbacks_controller.rb       |  9 +++++++--
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/app/controllers/session_controller.rb b/app/controllers/session_controller.rb
index e0344c3cc6a..463c830e3ab 100644
--- a/app/controllers/session_controller.rb
+++ b/app/controllers/session_controller.rb
@@ -11,15 +11,17 @@ class SessionController < ApplicationController
   end
 
   def sso
-    return_path = if params[:return_path]
-      params[:return_path]
-    elsif session[:destination_url]
-      uri = URI::parse(session[:destination_url])
-      "#{uri.path}#{uri.query ? "?" << uri.query : ""}"
-    else
-      path('/')
+    destination_url = cookies[:destination_url] || session[:destination_url]
+    return_path = params[:return_path] || path('/')
+
+    if destination_url && return_path == path('/')
+      uri = URI::parse(destination_url)
+      return_path = "#{uri.path}#{uri.query ? "?" << uri.query : ""}"
     end
 
+    session.delete(:destination_url)
+    cookies.delete(:destination_url)
+
     if SiteSetting.enable_sso?
       sso = DiscourseSingleSignOn.generate_sso(return_path)
       if SiteSetting.verbose_sso_logging
diff --git a/app/controllers/user_api_keys_controller.rb b/app/controllers/user_api_keys_controller.rb
index 8b532a1009c..20ba2d84aa5 100644
--- a/app/controllers/user_api_keys_controller.rb
+++ b/app/controllers/user_api_keys_controller.rb
@@ -20,7 +20,12 @@ class UserApiKeysController < ApplicationController
 
     unless current_user
       cookies[:destination_url] = request.fullpath
-      redirect_to path('/login')
+
+      if SiteSetting.enable_sso?
+        redirect_to path('/session/sso')
+      else
+        redirect_to path('/login')
+      end
       return
     end
 
diff --git a/app/controllers/users/omniauth_callbacks_controller.rb b/app/controllers/users/omniauth_callbacks_controller.rb
index ffef101a67d..5295b3a3fba 100644
--- a/app/controllers/users/omniauth_callbacks_controller.rb
+++ b/app/controllers/users/omniauth_callbacks_controller.rb
@@ -39,10 +39,15 @@ class Users::OmniauthCallbacksController < ApplicationController
     @auth_result = authenticator.after_authenticate(auth)
 
     origin = request.env['omniauth.origin']
+    if cookies[:destination_url].present?
+      origin = cookies[:destination_url]
+      cookies.delete(:destination_url)
+    end
+
     if origin.present?
-      parsed = URI.parse(@origin) rescue nil
+      parsed = URI.parse(origin) rescue nil
       if parsed
-        @origin = parsed.path
+        @origin = "#{parsed.path}?#{parsed.query}"
       end
     end