From 198ef91997c031511ee113f0658a1316dcac4499 Mon Sep 17 00:00:00 2001
From: Konrad Jurkowski <konrad@jurkowski.me>
Date: Wed, 12 Oct 2016 12:07:17 +0200
Subject: [PATCH] Allow to customize value passed to 2fa remember cookie

---
 .../two_factor_authentication_controller.rb   | 20 +++++++++++--------
 lib/two_factor_authentication.rb              |  3 +++
 .../hooks/two_factor_authenticatable.rb       |  2 +-
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/app/controllers/devise/two_factor_authentication_controller.rb b/app/controllers/devise/two_factor_authentication_controller.rb
index 3904b53..3823c7c 100644
--- a/app/controllers/devise/two_factor_authentication_controller.rb
+++ b/app/controllers/devise/two_factor_authentication_controller.rb
@@ -23,14 +23,7 @@ class Devise::TwoFactorAuthenticationController < DeviseController
   private
 
   def after_two_factor_success_for(resource)
-    expires_seconds = resource.class.remember_otp_session_for_seconds
-
-    if expires_seconds && expires_seconds > 0
-      cookies.signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME] = {
-          value: "#{resource.class}-#{resource.id}",
-          expires: expires_seconds.from_now
-      }
-    end
+    set_remember_tfa_cookie(resource)
 
     warden.session(resource_name)[TwoFactorAuthentication::NEED_AUTHENTICATION] = false
     bypass_sign_in(resource, scope: resource_name)
@@ -40,6 +33,17 @@ class Devise::TwoFactorAuthenticationController < DeviseController
     redirect_to after_two_factor_success_path_for(resource)
   end
 
+  def set_remember_tfa_cookie(resource)
+    expires_seconds = resource.class.remember_otp_session_for_seconds
+
+    if expires_seconds && expires_seconds > 0
+      cookies.signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME] = {
+          value: "#{resource.class}-#{resource.public_send(Devise.second_factor_resource_id)}",
+          expires: expires_seconds.from_now
+      }
+    end
+  end
+
   def after_two_factor_success_path_for(resource)
     stored_location_for(resource_name) || :root
   end
diff --git a/lib/two_factor_authentication.rb b/lib/two_factor_authentication.rb
index 8da594d..59ffa03 100644
--- a/lib/two_factor_authentication.rb
+++ b/lib/two_factor_authentication.rb
@@ -27,6 +27,9 @@ module Devise
 
   mattr_accessor :otp_secret_encryption_key
   @@otp_secret_encryption_key = ''
+
+  mattr_accessor :second_factor_resource_id
+  @@second_factor_resource_id = 'id'
 end
 
 module TwoFactorAuthentication
diff --git a/lib/two_factor_authentication/hooks/two_factor_authenticatable.rb b/lib/two_factor_authentication/hooks/two_factor_authenticatable.rb
index 208c0e2..6ecfd70 100644
--- a/lib/two_factor_authentication/hooks/two_factor_authenticatable.rb
+++ b/lib/two_factor_authentication/hooks/two_factor_authenticatable.rb
@@ -1,6 +1,6 @@
 Warden::Manager.after_authentication do |user, auth, options|
   if auth.env["action_dispatch.cookies"]
-    expected_cookie_value = "#{user.class}-#{user.id}"
+    expected_cookie_value = "#{user.class}-#{user.public_send(Devise.second_factor_resource_id)}",
     actual_cookie_value = auth.env["action_dispatch.cookies"].signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME]
     bypass_by_cookie = actual_cookie_value == expected_cookie_value
   end