From f40405c138af28ef6f571d77312373382e0dc744 Mon Sep 17 00:00:00 2001
From: Jack McCracken <jack.mccracken@shopify.com>
Date: Wed, 13 Oct 2021 11:19:20 -0400
Subject: [PATCH] Add support for custom CSRF strategies.

---
 .../metal/request_forgery_protection.rb       | 33 ++++++++++++++++---
 .../request_forgery_protection_test.rb        | 27 +++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/actionpack/lib/action_controller/metal/request_forgery_protection.rb b/actionpack/lib/action_controller/metal/request_forgery_protection.rb
index 9fbb41c031..c0b779e9be 100644
--- a/actionpack/lib/action_controller/metal/request_forgery_protection.rb
+++ b/actionpack/lib/action_controller/metal/request_forgery_protection.rb
@@ -124,10 +124,26 @@ module ActionController # :nodoc:
       #   If you need to add verification to the beginning of the callback chain, use <tt>prepend: true</tt>.
       # * <tt>:with</tt> - Set the method to handle unverified request.
       #
-      # Valid unverified request handling methods are:
+      # Built-in unverified request handling methods are:
       # * <tt>:exception</tt> - Raises ActionController::InvalidAuthenticityToken exception.
       # * <tt>:reset_session</tt> - Resets the session.
       # * <tt>:null_session</tt> - Provides an empty session during request but doesn't reset it completely. Used as default if <tt>:with</tt> option is not specified.
+      #
+      # You can also implement custom strategy classes for unverified request handling:
+      #
+      #    class CustomStrategy
+      #      def initialize(controller)
+      #        @controller = controller
+      #      end
+      #
+      #      def handle_unverified_request
+      #        # Custom behaviour for unverfied request
+      #      end
+      #    end
+      #
+      #    class ApplicationController < ActionController:x:Base
+      #      protect_from_forgery with: CustomStrategy
+      #    end
       def protect_from_forgery(options = {})
         options = options.reverse_merge(prepend: false)
 
@@ -148,9 +164,18 @@ module ActionController # :nodoc:
 
       private
         def protection_method_class(name)
-          ActionController::RequestForgeryProtection::ProtectionMethods.const_get(name.to_s.classify)
-        rescue NameError
-          raise ArgumentError, "Invalid request forgery protection method, use :null_session, :exception, or :reset_session"
+          return name if name.is_a?(Class)
+
+          case name
+          when :null_session
+            ProtectionMethods::NullSession
+          when :reset_session
+            ProtectionMethods::ResetSession
+          when :exception
+            ProtectionMethods::Exception
+          else
+            raise ArgumentError, "Invalid request forgery protection method, use :null_session, :exception, :reset_session, or a custom forgery protection class."
+          end
         end
     end
 
diff --git a/actionpack/test/controller/request_forgery_protection_test.rb b/actionpack/test/controller/request_forgery_protection_test.rb
index 265dae382a..f300c81751 100644
--- a/actionpack/test/controller/request_forgery_protection_test.rb
+++ b/actionpack/test/controller/request_forgery_protection_test.rb
@@ -103,6 +103,24 @@ class RequestForgeryProtectionControllerUsingNullSession < ActionController::Bas
   end
 end
 
+class RequestForgeryProtectionControllerUsingCustomStrategy < ActionController::Base
+  include RequestForgeryProtectionActions
+
+  class FakeException < Exception; end
+
+  class CustomStrategy
+    def initialize(controller)
+      @controller = controller
+    end
+
+    def handle_unverified_request
+      raise FakeException, "Raised a fake exception."
+    end
+  end
+
+  protect_from_forgery only: %w(index meta same_origin_js negotiate_same_origin), with: CustomStrategy
+end
+
 class PrependProtectForgeryBaseController < ActionController::Base
   before_action :custom_action
   attr_accessor :called_callbacks
@@ -699,6 +717,7 @@ end
 
 class RequestForgeryProtectionControllerUsingExceptionTest < ActionController::TestCase
   include RequestForgeryProtectionTests
+
   def assert_blocked(&block)
     assert_raises(ActionController::InvalidAuthenticityToken, &block)
   end
@@ -720,6 +739,14 @@ class RequestForgeryProtectionControllerUsingExceptionTest < ActionController::T
   end
 end
 
+class RequestForgeryProtectionControllerUsingCustomStrategyTest < ActionController::TestCase
+  include RequestForgeryProtectionTests
+
+  def assert_blocked(&block)
+    assert_raises(RequestForgeryProtectionControllerUsingCustomStrategy::FakeException, &block)
+  end
+end
+
 class PrependProtectForgeryBaseControllerTest < ActionController::TestCase
   PrependTrueController = Class.new(PrependProtectForgeryBaseController) do
     protect_from_forgery prepend: true