Optional subdomain aware redirects and scope based cookie names

app/controllers/devise/two_factor_authentication_controller.rb

@@ -27,7 +27,7 @@ def resend_code
   def after_two_factor_success_for(resource)
     set_remember_two_factor_cookie(resource)
 
-    warden.session(resource_name)[TwoFactorAuthentication::NEED_AUTHENTICATION] = false
+    warden.session(resource_name)[TwoFactorAuthentication::name_for(:need_authentication, resource_name)] = false
     # For compatability with devise versions below v4.2.0
     # https://github.com/plataformatec/devise/commit/2044fffa25d781fcbaf090e7728b48b65c854ccb
     if respond_to?(:bypass_sign_in)
@@ -45,7 +45,7 @@ def set_remember_two_factor_cookie(resource)
     expires_seconds = resource.class.remember_otp_session_for_seconds
 
     if expires_seconds && expires_seconds > 0
-      cookies.signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME] = {
+      cookies.signed[TwoFactorAuthentication::name_for(:remember_tfa_cookie_name, Devise::Mapping.find_scope!(resource))] = {
           value: "#{resource.class}-#{resource.public_send(Devise.second_factor_resource_id)}",
           expires: expires_seconds.seconds.from_now
       }

lib/two_factor_authentication.rb

@@ -38,6 +38,10 @@ module TwoFactorAuthentication
   NEED_AUTHENTICATION = 'need_two_factor_authentication'
   REMEMBER_TFA_COOKIE_NAME = "remember_tfa"
 
+  def self.name_for(sym, scope)
+    "#{self.const_get(sym.upcase.to_s)}_#{scope.to_s}"
+  end
+
   autoload :Schema, 'two_factor_authentication/schema'
   module Controllers
     autoload :Helpers, 'two_factor_authentication/controllers/helpers'

lib/two_factor_authentication/controllers/helpers.rb

@@ -12,7 +12,9 @@ module Helpers
       def handle_two_factor_authentication
         unless devise_controller?
           Devise.mappings.keys.flatten.any? do |scope|
-            if signed_in?(scope) and warden.session(scope)[TwoFactorAuthentication::NEED_AUTHENTICATION]
+            if signed_in?(scope) and warden.session(scope)[TwoFactorAuthentication::name_for(:need_authentication,
+                                                                                             scope)] and
+                public_send("current_#{scope}").class.subdomain_in_scope?(request.subdomain)
               handle_failed_second_factor(scope)
             end
           end
@@ -41,8 +43,9 @@ def two_factor_authentication_path_for(resource_or_scope = nil)
 module Devise
   module Controllers
     module Helpers
-      def is_fully_authenticated?
-        !session["warden.user.user.session"].try(:[], TwoFactorAuthentication::NEED_AUTHENTICATION)
+      def is_fully_authenticated?(resource_name)
+        !session["warden.user.user.session"].try(:[], TwoFactorAuthentication::name_for(:need_authentication,
+                                                                                        resource_name))
       end
     end
   end

lib/two_factor_authentication/hooks/two_factor_authenticatable.rb

@@ -1,17 +1,22 @@
 Warden::Manager.after_authentication do |user, auth, options|
   if auth.env["action_dispatch.cookies"]
     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]
+    actual_cookie_value = auth.env["action_dispatch.cookies"]
+                              .signed[TwoFactorAuthentication::name_for(:remember_tfa_cookie_name,
+                                                                        options[:scope])]
     bypass_by_cookie = actual_cookie_value == expected_cookie_value
   end
 
   if user.respond_to?(:need_two_factor_authentication?) && !bypass_by_cookie
-    if auth.session(options[:scope])[TwoFactorAuthentication::NEED_AUTHENTICATION] = user.need_two_factor_authentication?(auth.request)
+    if auth.session(options[:scope])[TwoFactorAuthentication::name_for(
+        :need_authentication, options[:scope]
+    )] = user.need_two_factor_authentication?(auth.request) and user.class.subdomain_in_scope?(auth.request.subdomain)
       user.send_new_otp if user.send_new_otp_after_login?
     end
   end
 end
 
 Warden::Manager.before_logout do |user, auth, _options|
-  auth.cookies.delete TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME if Devise.delete_cookie_on_logout
+  auth.cookies.delete TwoFactorAuthentication::name_for(:remember_tfa_cookie_name,
+                                                        _options[:scope]) if Devise.delete_cookie_on_logout
 end

lib/two_factor_authentication/models/two_factor_authenticatable.rb

@@ -11,6 +11,7 @@ module ClassMethods
         def has_one_time_password(options = {})
           include InstanceMethodsOnActivation
           include EncryptionInstanceMethods if options[:encrypted] == true
+          extend ClassMethodsOnActivation
         end
 
         ::Devise::Models.config(
@@ -20,6 +21,13 @@ def has_one_time_password(options = {})
         )
       end
 
+      module ClassMethodsOnActivation
+        def subdomain_in_scope?(subdomain)
+          return true unless respond_to? :two_factor_subdomains
+          !!(subdomain =~ two_factor_subdomains)
+        end
+      end
+
       module InstanceMethodsOnActivation
         def authenticate_otp(code, options = {})
           return true if direct_otp && authenticate_direct_otp(code)

spec/controllers/two_factor_authentication_controller_spec.rb

@@ -18,23 +18,23 @@ def post_code(code)
       it 'returns true' do
         controller.current_user.send_new_otp
         post_code controller.current_user.direct_otp
-        expect(subject.is_fully_authenticated?).to eq true
+        expect(subject.is_fully_authenticated?("user")).to eq true
       end
     end
 
     context 'when user has not entered any OTP yet' do
       it 'returns false' do
         get :show
 
-        expect(subject.is_fully_authenticated?).to eq false
+        expect(subject.is_fully_authenticated?("user")).to eq false
       end
     end
 
     context 'when user enters an invalid OTP' do
       it 'returns false' do
         post_code '12345'
 
-        expect(subject.is_fully_authenticated?).to eq false
+        expect(subject.is_fully_authenticated?("user")).to eq false
       end
     end
   end

spec/features/two_factor_authenticatable_spec.rb

@@ -189,7 +189,7 @@ def sms_sign_in
     end
 
     it 'sets the warden session need_two_factor_authentication key to true' do
-      session_hash = { 'need_two_factor_authentication' => true }
+      session_hash = { 'need_two_factor_authentication_user' => true }
 
       expect(page.get_rack_session_key('warden.user.user.session')).to eq session_hash
     end

spec/lib/two_factor_authentication/models/two_factor_authenticatable_spec.rb

@@ -323,4 +323,31 @@ def instance.send_two_factor_authentication_code(code)
       end
     end
   end
+
+  describe 'self.subdomain_in_scope?' do
+    let(:user) { EncryptedUser.new }
+
+    it 'allows scope if no optional methods' do
+      expect(user.class.subdomain_in_scope?('test1'))
+    end
+
+    it 'correctly parses a regex (neglect example)' do
+      allow(user.class).to receive(:two_factor_subdomains).and_return /^(?!(bad)$).*$/
+
+      expect(user.class.subdomain_in_scope?('good')).to be(true)
+      expect(user.class.subdomain_in_scope?('another-one')).to be(true)
+      expect(user.class.subdomain_in_scope?('bad')).to be(false)
+      expect(user.class.subdomain_in_scope?('goodbadugly')).to be(true)
+      expect(user.class.subdomain_in_scope?('')).to be(true)
+    end
+
+    it 'correctly parses a regex (whitelist example)' do
+      allow(user.class).to receive(:two_factor_subdomains).and_return /^(good|ugly)$/
+
+      expect(user.class.subdomain_in_scope?('good')).to be(true)
+      expect(user.class.subdomain_in_scope?('ugly')).to be(true)
+      expect(user.class.subdomain_in_scope?('bad')).to be(false)
+      expect(user.class.subdomain_in_scope?('')).to be(false)
+    end
+  end
 end

spec/support/controller_helper.rb

@@ -2,7 +2,7 @@ module ControllerHelper
   def sign_in(user = create_user('not_encrypted'))
     allow(warden).to receive(:authenticated?).with(:user).and_return(true)
     allow(controller).to receive(:current_user).and_return(user)
-    warden.session(:user)[TwoFactorAuthentication::NEED_AUTHENTICATION] = true
+    warden.session(:user)[TwoFactorAuthentication::name_for(:need_authentication, "user")] = true
   end
 end
 

spec/support/features_spec_helper.rb

@@ -20,11 +20,11 @@ def get_cookie key
   end
 
   def set_tfa_cookie value
-    set_cookie TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME, value
+    set_cookie TwoFactorAuthentication::name_for(:remember_tfa_cookie_name, "user"), value
   end
 
   def get_tfa_cookie
-    get_cookie TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME
+    get_cookie TwoFactorAuthentication::name_for(:remember_tfa_cookie_name, "user")
   end
 end