From 53ade6bb2f8f01065b10a5edb2d4121538258615 Mon Sep 17 00:00:00 2001
From: Bobby McDonald <bobbymcwho@gmail.com>
Date: Mon, 10 Aug 2020 20:22:59 -0400
Subject: [PATCH] Minor refactor of pkce

---
 lib/omniauth/strategies/oauth2.rb | 44 ++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 19 deletions(-)

diff --git a/lib/omniauth/strategies/oauth2.rb b/lib/omniauth/strategies/oauth2.rb
index d47137a..3759cbf 100644
--- a/lib/omniauth/strategies/oauth2.rb
+++ b/lib/omniauth/strategies/oauth2.rb
@@ -30,6 +30,16 @@ def self.inherited(subclass)
       option :auth_token_params, {}
       option :provider_ignores_state, false
       option :pkce, false
+      option :pkce_verifier, nil
+      option :pkce_options, {
+        :code_challenge => proc { |verifier|
+          Base64.urlsafe_encode64(
+            Digest::SHA2.digest(verifier),
+            padding: false
+          )
+        },
+        :code_challenge_method => "S256"
+      }
 
       attr_accessor :access_token
 
@@ -50,19 +60,20 @@ def request_phase
       end
 
       def authorize_params
-        verifier = SecureRandom.hex(64)
-
-        pkce_authorize_params!(verifier)
-
         options.authorize_params[:state] = SecureRandom.hex(24)
-        params = options.authorize_params.merge(options_for("authorize"))
 
         if OmniAuth.config.test_mode
           @env ||= {}
           @env["rack.session"] ||= {}
         end
 
-        build_authorize_session!(params, verifier)
+        params = options.authorize_params
+                        .merge(options_for("authorize"))
+                        .merge(pkce_authorize_params)
+
+        session["omniauth.pkce.verifier"] = options.pkce_verifier if options.pkce
+        session["omniauth.state"] = params[:state]
+
         params
       end
 
@@ -91,21 +102,16 @@ def callback_phase # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexi
 
     protected
 
-      def build_authorize_session!(params, verifier)
-        session["omniauth.pkce.verifier"] = verifier if options.pkce
-        session["omniauth.state"] = params[:state]
-      end
-
-      def pkce_authorize_params!(verifier)
-        return unless options.pkce
+      def pkce_authorize_params
+        return {} unless options.pkce
+        options.pkce_verifier = SecureRandom.hex(64)
 
         # NOTE: see https://tools.ietf.org/html/rfc7636#appendix-A
-        challenge = Base64
-                    .urlsafe_encode64(Digest::SHA2.digest(verifier))
-                    .split("=")
-                    .first
-        options.authorize_params[:code_challenge] = challenge
-        options.authorize_params[:code_challenge_method] = "S256"
+        {
+          :code_challenge => options.pkce_options[:code_challenge]
+                                    .call(options.pkce_verifier),
+          :code_challenge_method => options.pkce_options[:code_challenge_method]
+        }
       end
 
       def pkce_token_params