Reuse callback_url between request phase and callback phase

lib/omniauth/strategies/oauth2.rb

@@ -73,6 +73,8 @@ def authorize_params # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
 
         session["omniauth.pkce.verifier"] = options.pkce_verifier if options.pkce
         session["omniauth.state"] = params[:state]
+        # store the callback_url in session during request phase to be re-used during callback phase
+        session["omniauth.callback_url"] = callback_url
 
         params
       end
@@ -121,9 +123,17 @@ def pkce_token_params
         {:code_verifier => session.delete("omniauth.pkce.verifier")}
       end
 
+      def callback_url_from_session
+        session.delete("omniauth.callback_url")
+      end
+
       def build_access_token
         verifier = request.params["code"]
-        client.auth_code.get_token(verifier, {:redirect_uri => callback_url}.merge(token_params.to_hash(:symbolize_keys => true)), deep_symbolize(options.auth_token_params))
+        client
+          .auth_code
+          .get_token(verifier,
+                     {:redirect_uri => callback_url_from_session}.merge(token_params.to_hash(:symbolize_keys => true)),
+                     deep_symbolize(options.auth_token_params))
       end
 
       def deep_symbolize(options)

spec/omniauth/strategies/oauth2_spec.rb

@@ -61,6 +61,12 @@ def app
       expect(instance.session["omniauth.state"]).not_to be_empty
     end
 
+    it "sets callback_url in session" do
+      instance = subject.new("abc", "def")
+      instance.authorize_params
+      expect(instance.session["omniauth.callback_url"]).not_to be_empty
+    end
+
     it "includes custom state in the authorize params" do
       instance = subject.new("abc", "def", :state => proc { "qux" })
       expect(instance.authorize_params.keys).to eq(["state"])
@@ -96,6 +102,28 @@ def app
     end
   end
 
+  describe "#build_access_token" do
+    subject { fresh_strategy }
+    let(:request_double) { double("Request", params: {}) }
+    let(:client_double) { double }
+    let(:strategy_double) { double }
+    let(:callback_url) { "http://abc.def/auth/abc/callback" }
+
+    it "uses callback_url stored in session" do
+      instance = subject.new("abc", "def")
+
+      allow(instance).to receive(:request).and_return(request_double)
+      allow(instance).to receive(:session).and_return({ "omniauth.callback_url" => callback_url })
+      allow(instance).to receive(:client).and_return(client_double)
+      allow(client_double).to receive(:auth_code).and_return(strategy_double)
+
+      expect(strategy_double).to receive(:get_token)
+        .with(nil, hash_including(redirect_uri: callback_url), {})
+
+      instance.send(:build_access_token)
+    end
+  end
+
   describe "#callback_phase" do
     subject { fresh_strategy }
     it "calls fail with the client error received" do