specs for forced two-factor authentication and grace period

simplified code and fixed stuffs

app/controllers/application_controller.rb

@@ -226,12 +226,7 @@ class ApplicationController < ActionController::Base
   def check_tfa_requirement
     if two_factor_authentication_required? && current_user && !current_user.two_factor_enabled && !skip_two_factor?
-      grace_period_started = current_user.otp_grace_period_started_at
-      grace_period_deadline = grace_period_started + two_factor_grace_period.hours
-
-      deadline_text = "until #{l(grace_period_deadline)}" unless two_factor_grace_period_expired?(grace_period_started)
-      redirect_to new_profile_two_factor_auth_path,
-                  alert: "You must configure Two-Factor Authentication in your account #{deadline_text}"
+      redirect_to new_profile_two_factor_auth_path
     end
   end
@@ -377,7 +372,8 @@ class ApplicationController < ActionController::Base
     current_application_settings.two_factor_grace_period
   end
-  def two_factor_grace_period_expired?(date)
+  def two_factor_grace_period_expired?
+    date = current_user.otp_grace_period_started_at
     date && (date + two_factor_grace_period.hours) < Time.current
   end

app/controllers/profiles/two_factor_auths_controller.rb

@@ -10,6 +10,13 @@ class Profiles::TwoFactorAuthsController < Profiles::ApplicationController
     end
     current_user.save! if current_user.changed?
+    if two_factor_grace_period_expired?
+      flash.now[:alert] = 'You must configure Two-Factor Authentication in your account.'
+    else
+      grace_period_deadline = current_user.otp_grace_period_started_at + two_factor_grace_period.hours
+      flash.now[:alert] = "You must configure Two-Factor Authentication in your account until #{l(grace_period_deadline)}."
+    end
+
     @qr_code = build_qr_code
   end
@@ -40,7 +47,7 @@ class Profiles::TwoFactorAuthsController < Profiles::ApplicationController
   end
   def skip
-    if two_factor_grace_period_expired?(current_user.otp_grace_period_started_at)
+    if two_factor_grace_period_expired?
       redirect_to new_profile_two_factor_auth_path, alert: 'Cannot skip two factor authentication setup'
     else
       session[:skip_tfa] = current_user.otp_grace_period_started_at + two_factor_grace_period.hours

spec/features/login_spec.rb

@@ -98,4 +98,56 @@ feature 'Login', feature: true do
       expect(page).to have_content('Invalid login or password.')
     end
   end
+
+  describe 'with required two-factor authentication enabled' do
+    let(:user) { create(:user) }
+    before(:each) { stub_application_setting(require_two_factor_authentication: true) }
+
+    context 'with grace period defined' do
+      before(:each) do
+        stub_application_setting(two_factor_grace_period: 48)
+        login_with(user)
+      end
+
+      context 'within the grace period' do
+        it 'redirects to two-factor configuration page' do
+          expect(current_path).to eq new_profile_two_factor_auth_path
+          expect(page).to have_content('You must configure Two-Factor Authentication in your account until')
+        end
+
+        it 'two-factor configuration is skippable' do
+          expect(current_path).to eq new_profile_two_factor_auth_path
+
+          click_link 'Configure it later'
+          expect(current_path).to eq root_path
+        end
+      end
+
+      context 'after the grace period' do
+        let(:user) { create(:user, otp_grace_period_started_at: 9999.hours.ago) }
+
+        it 'redirects to two-factor configuration page' do
+          expect(current_path).to eq new_profile_two_factor_auth_path
+          expect(page).to have_content('You must configure Two-Factor Authentication in your account.')
+        end
+
+        it 'two-factor configuration is not skippable' do
+          expect(current_path).to eq new_profile_two_factor_auth_path
+          expect(page).not_to have_link('Configure it later')
+        end
+      end
+    end
+
+    context 'without grace pariod defined' do
+      before(:each) do
+        stub_application_setting(two_factor_grace_period: 0)
+        login_with(user)
+      end
+
+      it 'redirects to two-factor configuration page' do
+        expect(current_path).to eq new_profile_two_factor_auth_path
+        expect(page).to have_content('You must configure Two-Factor Authentication in your account.')
+      end
+    end
+  end
 end