From 507b859fa64b81226a593384afaef932d87b92b7 Mon Sep 17 00:00:00 2001
From: Coatezy <2954283+coatezy@users.noreply.github.com>
Date: Wed, 27 Mar 2024 10:52:45 +0000
Subject: [PATCH] fix: use client.site instead of the callback_url

By using client.site instead of callback_url we can preserve any query string parameters that could be included within the site key and also configured in the callback url within the user pool client app
---
 lib/omniauth/strategies/cognito_idp.rb       | 2 +-
 spec/omniauth/strategies/cognito_idp_spec.rb | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/lib/omniauth/strategies/cognito_idp.rb b/lib/omniauth/strategies/cognito_idp.rb
index 50b9a4f..b0071e6 100644
--- a/lib/omniauth/strategies/cognito_idp.rb
+++ b/lib/omniauth/strategies/cognito_idp.rb
@@ -67,7 +67,7 @@ class CognitoIdP < OmniAuth::Strategies::OAuth2
       def build_access_token
         client.auth_code.get_token(
           request.params['code'],
-          { redirect_uri: callback_url.split('?').first }.merge(token_params.to_hash(symbolize_keys: true)),
+          { redirect_uri: client.site }.merge(token_params.to_hash(symbolize_keys: true)),
           deep_symbolize(options.auth_token_params)
         )
       end
diff --git a/spec/omniauth/strategies/cognito_idp_spec.rb b/spec/omniauth/strategies/cognito_idp_spec.rb
index f162315..7c9c0ae 100644
--- a/spec/omniauth/strategies/cognito_idp_spec.rb
+++ b/spec/omniauth/strategies/cognito_idp_spec.rb
@@ -11,6 +11,7 @@
   let(:options) { {} }
   let(:client_id) { 'ABCDE' }
   let(:client_secret) { '987654321' }
+  let(:site) { "http://localhost/auth/cognito-idp/callback" }
 
   around do |example|
     OmniAuth.config.test_mode = true
@@ -31,10 +32,10 @@
       end
     end
 
-    let(:oauth_client) { double('OAuth2::Client', auth_code: auth_code) }
+    let(:oauth_client) { double('OAuth2::Client', auth_code: auth_code, site: site) }
     let(:auth_code) { double('OAuth2::AuthCode', get_token: access_token_object) }
     let(:access_token_object) { double('OAuth2::AccessToken') }
-    let(:callback_url) { 'http://localhost/auth/cognito-idp/callback?code=1234' }
+    let(:callback_url) { "#{site}?code=1234" }
 
     let(:request) { double('Rack::Request', params: params) }
     let(:params) { { 'code' => '12345' } }
@@ -57,7 +58,7 @@
     let(:env) { {} }
     let(:request) { double('Rack::Request', params: {'state' => strategy.session['omniauth.state']}) }
     let(:session) { { 'omniauth.state' => 'some_state' } }
-    let(:oauth_client) { double('OAuth2::Client', auth_code: auth_code) }
+    let(:oauth_client) { double('OAuth2::Client', auth_code: auth_code, site: site) }
     let(:auth_code) { double('OAuth2::AuthCode') }
     let(:access_token_object) { OAuth2::AccessToken.from_hash(oauth_client, token_hash) }