From 2f37b6e86a6a619dd4cc941b7706119462cc5764 Mon Sep 17 00:00:00 2001
From: Lynn Hurley <lynn.dylan.hurley@gmail.com>
Date: Sun, 14 Aug 2016 13:56:26 -0500
Subject: [PATCH] fix #696

---
 lib/devise_token_auth/rails/routes.rb         | 16 ++++++-
 .../omniauth_callbacks_controller_test.rb     | 42 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/lib/devise_token_auth/rails/routes.rb b/lib/devise_token_auth/rails/routes.rb
index 06818cf53..c5db0d887 100644
--- a/lib/devise_token_auth/rails/routes.rb
+++ b/lib/devise_token_auth/rails/routes.rb
@@ -73,8 +73,22 @@ def mount_devise_token_auth_for(resource, opts)
 
               set_omniauth_path_prefix!(DeviseTokenAuth.omniauth_prefix)
 
+              redirect_params = {}.tap {|hash| qs.each{|k, v| hash[k] = v.first}}
+
+              if DeviseTokenAuth.redirect_whitelist
+                redirect_url = request.params['auth_origin_url']
+                unless DeviseTokenAuth.redirect_whitelist.include?(redirect_url)
+                  message = I18n.t(
+                    'devise_token_auth.registrations.redirect_url_not_allowed',
+                    redirect_url: redirect_url
+                  )
+                  redirect_params['message'] = message
+                  next "#{::OmniAuth.config.path_prefix}/failure?#{redirect_params.to_param}"
+                end
+              end
+
               # re-construct the path for omniauth
-              "#{::OmniAuth.config.path_prefix}/#{params[:provider]}?#{{}.tap {|hash| qs.each{|k, v| hash[k] = v.first}}.to_param}"
+              "#{::OmniAuth.config.path_prefix}/#{params[:provider]}?#{redirect_params.to_param}"
             }, via: [:get]
           end
         end
diff --git a/test/controllers/devise_token_auth/omniauth_callbacks_controller_test.rb b/test/controllers/devise_token_auth/omniauth_callbacks_controller_test.rb
index f81acb6f1..91c85835e 100644
--- a/test/controllers/devise_token_auth/omniauth_callbacks_controller_test.rb
+++ b/test/controllers/devise_token_auth/omniauth_callbacks_controller_test.rb
@@ -279,4 +279,46 @@ def get_success(params = {})
       }
     end
   end
+
+  describe 'Using redirect_whitelist' do
+    before do
+      @user_email = 'slemp.diggler@sillybandz.gov'
+      OmniAuth.config.mock_auth[:facebook] = OmniAuth::AuthHash.new(
+        provider: 'facebook',
+        uid: '123545',
+        info: {
+          name: 'chong',
+          email: @user_email
+        }
+      )
+      @good_redirect_url = Faker::Internet.url
+      @bad_redirect_url = Faker::Internet.url
+      DeviseTokenAuth.redirect_whitelist = [@good_redirect_url]
+    end
+
+    teardown do
+      DeviseTokenAuth.redirect_whitelist = nil
+    end
+
+    test 'request using non-whitelisted redirect fail' do
+      get_via_redirect '/auth/facebook',
+                       auth_origin_url: @bad_redirect_url,
+                       omniauth_window_type: 'newWindow'
+
+      data_json = @response.body.match(/var data \= (.+)\;/)[1]
+      data = ActiveSupport::JSON.decode(data_json)
+      assert_equal "Redirect to '#{@bad_redirect_url}' not allowed.",
+                   data['error']
+    end
+
+    test 'request to whitelisted redirect should succeed' do
+      get_via_redirect '/auth/facebook',
+                       auth_origin_url: @good_redirect_url,
+                       omniauth_window_type: 'newWindow'
+
+      data_json = @response.body.match(/var data \= (.+)\;/)[1]
+      data = ActiveSupport::JSON.decode(data_json)
+      assert_equal @user_email, data['email']
+    end
+  end
 end