From 4009905531f28ebd7ecab22f898b5d6180eefb4a Mon Sep 17 00:00:00 2001
From: Geoff Harcourt <geoff.harcourt@gmail.com>
Date: Fri, 6 Nov 2020 13:41:26 -0500
Subject: [PATCH] Allow wrapping of lockable Warden after sign-in hook

For applications using a primary/replica multiple database setup, Rails
6's default connection switching can lead to the `Lockable` after
sign-in hook firing on `GET` requests, getting the read-only database,
and then attempting to write while zeroing out the user's failed attempts.

This change adds a customizable wrapper proc that can allow applications
to customize this behavior. For a primary/replica setup using
ActiveRecord's default connection switching, it might look like this:

```ruby
Devise.setup do |config|
  config.warden_hook_save_wrapper = Proc.new do |hook|
    ApplicationRecord.connected_to(role: :writing) do
      hook.call
    end
  end
end
```

Fix #5264
---
 lib/devise.rb                | 3 +++
 lib/devise/hooks/lockable.rb | 8 ++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/lib/devise.rb b/lib/devise.rb
index 0451876df9..dcdacc151a 100644
--- a/lib/devise.rb
+++ b/lib/devise.rb
@@ -297,6 +297,9 @@ module Test
   mattr_accessor :sign_in_after_change_password
   @@sign_in_after_change_password = true
 
+  mattr_accessor :warden_hook_save_wrapper
+  @@warden_hook_save_wrapper = Proc.new { |hook| hook.call }
+
   def self.activerecord51? # :nodoc:
     defined?(ActiveRecord) && ActiveRecord.gem_version >= Gem::Version.new("5.1.x")
   end
diff --git a/lib/devise/hooks/lockable.rb b/lib/devise/hooks/lockable.rb
index a73a1752e2..fe964db951 100644
--- a/lib/devise/hooks/lockable.rb
+++ b/lib/devise/hooks/lockable.rb
@@ -5,8 +5,12 @@
 Warden::Manager.after_set_user except: :fetch do |record, warden, options|
   if record.respond_to?(:failed_attempts) && warden.authenticated?(options[:scope])
     unless record.failed_attempts.to_i.zero?
-      record.failed_attempts = 0
-      record.save(validate: false)
+      wipe_attempts = Proc.new do
+        record.failed_attempts = 0
+        record.save(validate: false)
+      end
+
+      Devise.warden_hook_save_wrapper.call(wipe_attempts)
     end
   end
 end