From e51150069644c9c93980e3a30bfc1f8134ebe404 Mon Sep 17 00:00:00 2001
From: Jordan Owens <jkowens@gmail.com>
Date: Tue, 6 Oct 2020 17:57:33 -0400
Subject: [PATCH] Add option to configure token session key

---
 .../lib/rack/protection/authenticity_token.rb        | 12 ++++++++++--
 .../lib/rack/protection/authenticity_token_spec.rb   | 11 +++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/rack-protection/lib/rack/protection/authenticity_token.rb b/rack-protection/lib/rack/protection/authenticity_token.rb
index 23a7784256..13496a7d72 100644
--- a/rack-protection/lib/rack/protection/authenticity_token.rb
+++ b/rack-protection/lib/rack/protection/authenticity_token.rb
@@ -24,6 +24,13 @@ module Protection
     #                                the token on a request. Default value:
     #                                <tt>"authenticity_token"</tt>
     #
+    # [<tt>:key</tt>] the name of the param that should contain
+    #                                the token in the session. Default value:
+    #                                <tt>:csrf</tt>
+    #
+    # [<tt>:allow_if</tt>] a proc for custom allow/deny logic. Default value:
+    #                                <tt>nil</tt>
+    #
     # == Example: Forms application
     #
     # To show what the AuthenticityToken does, this section includes a sample
@@ -85,6 +92,7 @@ class AuthenticityToken < Base
       TOKEN_LENGTH = 32
 
       default_options :authenticity_param => 'authenticity_token',
+                      :key => :csrf,
                       :allow_if => nil
 
       def self.token(session)
@@ -113,7 +121,7 @@ def mask_authenticity_token(session)
       private
 
       def set_token(session)
-        session[:csrf] ||= self.class.random_token
+        session[options[:key]] ||= self.class.random_token
       end
 
       # Checks the client's masked token to see if it matches the
@@ -177,7 +185,7 @@ def compare_with_real_token(token, session)
       end
 
       def real_token(session)
-        decode_token(session[:csrf])
+        decode_token(session[options[:key]])
       end
 
       def encode_token(token)
diff --git a/rack-protection/spec/lib/rack/protection/authenticity_token_spec.rb b/rack-protection/spec/lib/rack/protection/authenticity_token_spec.rb
index 75081f6c89..40eb6b0e66 100644
--- a/rack-protection/spec/lib/rack/protection/authenticity_token_spec.rb
+++ b/rack-protection/spec/lib/rack/protection/authenticity_token_spec.rb
@@ -59,6 +59,17 @@
     expect(env['rack.session'][:csrf]).not_to be_nil
   end
 
+  it "allows for a custom token session key" do
+    mock_app do
+      use Rack::Session::Cookie, :key => 'rack.session'
+      use Rack::Protection::AuthenticityToken, :key => :_csrf
+      run DummyApp
+    end
+
+    get '/'
+    expect(env['rack.session'][:_csrf]).not_to be_nil
+  end
+
   describe ".token" do
     it "returns a unique masked version of the authenticity token" do
       expect(Rack::Protection::AuthenticityToken.token(session)).not_to eq(masked_token)